| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch.nn as nn |
| 5 | |
| 6 | from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( |
| 7 | AutoRegressiveSpeculator, |
| 8 | ) |
| 9 | from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model |
| 10 | |
| 11 | |
| 12 | class MTPSpeculator(AutoRegressiveSpeculator): |
| 13 | share_mtp_topk_indices: bool = False |
| 14 | |
| 15 | def load_draft_model( |
| 16 | self, |
| 17 | target_model: nn.Module, |
| 18 | target_attn_layer_names: set[str], |
| 19 | ) -> nn.Module: |
| 20 | draft_model = load_eagle_model(target_model, self.vllm_config) |
| 21 | spec_config = self.vllm_config.speculative_config |
| 22 | draft_hf_config = ( |
| 23 | spec_config.draft_model_config.hf_config |
| 24 | if spec_config is not None |
| 25 | else None |
| 26 | ) |
| 27 | # Detect index_share_for_mtp_iteration. When True, the proposer |
| 28 | # toggles skip_topk so step 0 computes MTP's own indices and |
| 29 | # steps 1+ reuse them. |
| 30 | self.share_mtp_topk_indices = ( |
| 31 | self.vllm_config.parallel_config.prefill_context_parallel_size == 1 |
| 32 | and getattr(draft_hf_config, "index_share_for_mtp_iteration", False) |
| 33 | and hasattr(draft_model.model, "set_skip_topk") |
| 34 | and hasattr(draft_model.model, "compact_topk_indices") |
| 35 | ) |
| 36 | return draft_model |
| 37 | |
| 38 | def on_prefill_begin(self, num_reqs: int) -> None: |
| 39 | # Step 0 computes its own top-k. Unconditional, so a step that died |
| 40 | # midway cannot leave reuse mode on. |
| 41 | if self.share_mtp_topk_indices: |
| 42 | self.model.model.set_skip_topk(False) |
| 43 | |
| 44 | def on_prefill_end(self, num_reqs: int) -> None: |
| 45 | # Step 0 (prefill) wrote topk indices for every query token in the |
| 46 | # multi-token batch. Compact them down to each request's last token so |
| 47 | # steps 1+ can reuse them from the shared buffer. |
| 48 | if self.share_mtp_topk_indices and self.num_speculative_steps > 1: |
| 49 | self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs]) |
| 50 | |
| 51 | def on_multi_step_decode_begin(self, num_reqs: int) -> None: |
| 52 | # Switch to reuse mode so draft steps 1+ skip the indexer op and read |
| 53 | # the indices that step 0 wrote into the shared buffer. |
| 54 | if self.share_mtp_topk_indices: |
| 55 | self.model.model.set_skip_topk(True) |
| 56 | |
| 57 | def on_multi_step_decode_end(self, num_reqs: int) -> None: |
| 58 | if self.share_mtp_topk_indices: |
| 59 | self.model.model.set_skip_topk(False) |