Method _generate_with_skipping

python/TIDE/runtime.py:223–313

Generate tokens with post-hoc early exit evaluation. Runs the full model forward with output_hidden_states=True, then evaluates routers post-hoc to decide which layer's hidden state to use for logits. All layers run every step (preserving KV cache correctness), but the outpu

(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int,
        temperature: float,
        top_k: int,
        top_p: float,
    )

                )

    def _generate_with_skipping(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int,
        temperature: float,
        top_k: int,
        top_p: float,
    ) -> torch.Tensor:
        """Generate tokens with post-hoc early exit evaluation.

        Runs the full model forward with output_hidden_states=True, then evaluates
        routers post-hoc to decide which layer's hidden state to use for logits.
        All layers run every step (preserving KV cache correctness), but the output
        comes from the earliest layer whose router fires.

        Compatible with all transformers versions (no hooks or exceptions).
        """
        device = input_ids.device
        B = input_ids.shape[0]
        n_layers = len(self._layers)
        generated = input_ids.clone()
        gen_stats = ExitStats(total_tokens=0)

        # Prefill: run full model
        prefill_out = self.model(
            generated, use_cache=True, return_dict=True,
        )
        past_key_values = prefill_out.past_key_values
        next_logits = prefill_out.logits[:, -1, :]

        for step in range(max_new_tokens):
            next_token = self._sample_next_token(next_logits, temperature, top_k, top_p)
            generated = torch.cat([generated, next_token], dim=-1)

            if (next_token == 2).all():
                break

            gen_stats.total_tokens += B

            # Run full forward with hidden states for post-hoc router evaluation
            out = self.model(
                next_token,
                past_key_values=past_key_values,
                use_cache=True,
                return_dict=True,
                output_hidden_states=True,
            )
            past_key_values = out.past_key_values
            all_hidden = out.hidden_states  # [emb, layer_0, ..., layer_N]

            # Evaluate routers post-hoc to find earliest exit
            exit_hidden = None
            exit_layer = None
            for layer_idx in sorted(self.routers.keys()):
                if layer_idx < self.config.min_layers:
                    continue
                hidden = all_hidden[layer_idx + 1]  # +1 for embedding offset
                scores = self._score_router(hidden, layer_idx)