Main TIDE runtime: wraps a HuggingFace model with early-exit inference. Uses a hook-based approach: registers forward hooks on each checkpoint layer to evaluate routers and freeze converged tokens. Frozen tokens keep their hidden state from the exit point while remaining tokens continue
| 54 | |
| 55 | |
| 56 | class TIDERuntime(nn.Module): |
| 57 | """Main TIDE runtime: wraps a HuggingFace model with early-exit inference. |
| 58 | |
| 59 | Uses a hook-based approach: registers forward hooks on each checkpoint layer |
| 60 | to evaluate routers and freeze converged tokens. Frozen tokens keep their |
| 61 | hidden state from the exit point while remaining tokens continue through |
| 62 | deeper layers. This preserves attention correctness. |
| 63 | """ |
| 64 | |
| 65 | def __init__( |
| 66 | self, |
| 67 | model: nn.Module, |
| 68 | router_path: str | Path, |
| 69 | config: Optional[TIDEConfig] = None, |
| 70 | use_cuda_kernels: bool = True, |
| 71 | ): |
| 72 | super().__init__() |
| 73 | self.model = model |
| 74 | self.config = config or TIDEConfig() |
| 75 | self.adapter = get_adapter(model) |
| 76 | model_on_cuda = self._device.type == "cuda" |
| 77 | self.use_cuda = use_cuda_kernels and _cuda_available and model_on_cuda |
| 78 | |
| 79 | # Load routers |
| 80 | checkpoint = RouterCheckpoint.load(router_path, device=self._device_str) |
| 81 | self.routers: Dict[int, TokenRouter] = {} |
| 82 | for layer_idx, router in checkpoint.routers.items(): |
| 83 | router = router.to(self._device) |
| 84 | router.requires_grad_(False) |
| 85 | if self.use_cuda: |
| 86 | router.down.weight.data = router.down.weight.data.t().contiguous() |
| 87 | router.up.weight.data = router.up.weight.data.t().contiguous() |
| 88 | self.routers[layer_idx] = router |
| 89 | |
| 90 | self._layers = self.adapter.get_layers(model) |
| 91 | self._final_norm = self.adapter.get_final_norm(model) |
| 92 | self._lm_head = self.adapter.get_lm_head(model) |
| 93 | |
| 94 | self.last_stats: Optional[ExitStats] = None |
| 95 | logger.info( |
| 96 | f"TIDERuntime initialized: {len(self._layers)} layers, " |
| 97 | f"{len(self.routers)} routers, CUDA={'on' if self.use_cuda else 'off'}" |
| 98 | ) |
| 99 | |
| 100 | @property |
| 101 | def _device(self) -> torch.device: |
| 102 | return next(self.model.parameters()).device |
| 103 | |
| 104 | @property |
| 105 | def _device_str(self) -> str: |
| 106 | return str(self._device) |
| 107 | |
| 108 | def _score_router(self, hidden: torch.Tensor, layer_idx: int) -> torch.Tensor: |
| 109 | """Score tokens for exit at a checkpoint layer. |
| 110 | |
| 111 | Args: |
| 112 | hidden: [B, S, D] or [N, D] hidden states |
| 113 | Returns: |
no outgoing calls