Autoregressive generation with post-hoc early exit. Each decode step runs all layers (preserving KV cache correctness), then evaluates routers to select which layer's output to use for the next token. Compatible with all transformers versions.
(
self,
input_ids: torch.Tensor,
max_new_tokens: int = 512,
temperature: float = 1.0,
top_k: int = 50,
top_p: float = 0.9,
**kwargs,
)
| 353 | |
| 354 | @torch.no_grad() |
| 355 | def generate( |
| 356 | self, |
| 357 | input_ids: torch.Tensor, |
| 358 | max_new_tokens: int = 512, |
| 359 | temperature: float = 1.0, |
| 360 | top_k: int = 50, |
| 361 | top_p: float = 0.9, |
| 362 | **kwargs, |
| 363 | ) -> torch.Tensor: |
| 364 | """Autoregressive generation with post-hoc early exit. |
| 365 | |
| 366 | Each decode step runs all layers (preserving KV cache correctness), |
| 367 | then evaluates routers to select which layer's output to use for |
| 368 | the next token. Compatible with all transformers versions. |
| 369 | """ |
| 370 | return self._generate_with_skipping( |
| 371 | input_ids, max_new_tokens, temperature, top_k, top_p, |
| 372 | ) |
| 373 | |
| 374 | @staticmethod |
| 375 | def calibrate( |