Method forward

python/TIDE/runtime.py:135–192

Forward pass with per-token early exit via frozen-token approach. Tokens that converge at a checkpoint get their normed hidden state saved. They continue through the model (for attention) but their contribution to the final output comes from the exit point, not the last laye

(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    )

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> torch.Tensor:
        """Forward pass with per-token early exit via frozen-token approach.

        Tokens that converge at a checkpoint get their normed hidden state saved.
        They continue through the model (for attention) but their contribution to
        the final output comes from the exit point, not the last layer.
        """
        B, S = input_ids.shape
        device = input_ids.device

        # Run full model forward, collecting all layer hidden states
        model_output = self.model(
            input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True,
        )

        # hidden_states: [embedding, layer_0, layer_1, ..., layer_N]
        all_hidden = model_output.hidden_states
        D = all_hidden[-1].shape[-1]
        dtype = all_hidden[-1].dtype

        output_buffer = torch.zeros(B, S, D, device=device, dtype=dtype)
        exited_mask = torch.zeros(B, S, dtype=torch.bool, device=device)
        stats = ExitStats(total_tokens=B * S)

        # Check routers at each checkpoint layer (post-hoc, using saved hidden states)
        for layer_idx in sorted(self.routers.keys()):
            if layer_idx < self.config.min_layers:
                continue

            hidden = all_hidden[layer_idx + 1]  # +1 because index 0 is embeddings
            scores = self._score_router(hidden, layer_idx)  # [B, S]
            new_exits = (scores > self.config.exit_threshold) & (~exited_mask)

            if new_exits.any():
                normed = self._final_norm(hidden.float()).to(dtype)
                output_buffer[new_exits] = normed[new_exits]
                exited_mask = exited_mask | new_exits
                stats.exits_per_layer[layer_idx] = new_exits.sum().item()

        # Remaining tokens use the final hidden state
        remaining_mask = ~exited_mask
        if remaining_mask.any():
            final_hidden = all_hidden[-1]
            normed = self._final_norm(final_hidden.float()).to(dtype)
            output_buffer[remaining_mask] = normed[remaining_mask]
        stats.remaining_tokens = remaining_mask.sum().item()

        self.last_stats = stats
        logits = self._lm_head(output_buffer)
        return logits