[Model Runner V2][MTP] Share topk index buffer between draft steps

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

@@ -54,6 +54,18 @@ def load_model(self, target_model: nn.Module) -> None:
             device=self.device,
         )
 
+    # Lifecycle hooks for model-specific optimizations. Subclasses override
+    # the ones they need. These fire in both `capture` and `propose` so that
+    # any state they toggle (e.g. attention flags baked into a CUDA graph) is
+    # identical at capture time and replay time.
+    def on_prefill_begin(self, num_reqs: int) -> None: ...
+
+    def on_prefill_end(self, num_reqs: int) -> None: ...
+
+    def on_multi_step_decode_begin(self, num_reqs: int) -> None: ...
+
+    def on_multi_step_decode_end(self, num_reqs: int) -> None: ...
+
     @property
     def advance_draft_positions(self) -> bool:
         """
@@ -104,6 +116,8 @@ def capture(self) -> None:
         assert self.prefill_cudagraph_manager is not None
         if self.prefill_cudagraph_manager.use_breakable_cg:
             self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
+
+        self.on_prefill_begin(self.max_num_reqs)
         self.prefill_cudagraph_manager.capture(
             self._prefill,
             self.model_state,
@@ -113,10 +127,12 @@ def capture(self) -> None:
             self.kv_cache_config,
             progress_bar_desc="Capturing prefill CUDA graphs",
         )
+        self.on_prefill_end(self.max_num_reqs)
 
         if self.num_speculative_steps == 1:
             return
 
+        self.on_multi_step_decode_begin(self.max_num_reqs)
         # Capture the decode draft generation routine (model forward +
         # sample + update_draft_inputs) for a single
         # step.
@@ -130,6 +146,7 @@ def capture(self) -> None:
             self.kv_cache_config,
             progress_bar_desc="Capturing decode CUDA graphs",
         )
+        self.on_multi_step_decode_end(self.max_num_reqs)
 
     @torch.inference_mode()
     def propose(
@@ -159,7 +176,8 @@ def propose(
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
         is_profile: bool = False,
     ) -> torch.Tensor:
-        num_tokens = input_batch.num_tokens_after_padding
+        num_tokens = input_batch.num_tokens
+        num_tokens_padded = input_batch.num_tokens_after_padding
         num_reqs = input_batch.num_reqs
         max_query_len = input_batch.num_scheduled_tokens.max()
         max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
@@ -180,7 +198,7 @@ def propose(
             )
         else:
             hidden_states = last_hidden_states
-        self.hidden_states[:num_tokens].copy_(hidden_states)
+        self.hidden_states[:num_tokens_padded].copy_(hidden_states)
 
         self._copy_request_inputs(
             num_reqs,
@@ -208,21 +226,22 @@ def propose(
             num_reqs,
             # Use the actual number of tokens without padding added by
             # the target model during FULL cudagraph.
-            input_batch.num_tokens,
+            num_tokens,
             max_query_len,
         )
         prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
             self.prefill_cudagraph_manager,
             num_reqs,
-            num_tokens,
+            num_tokens_padded,
             uniform_token_count,
             dp_size=self.dp_size,
             dp_rank=self.dp_rank,
             need_eager=is_profile,
         )
 
-        self._prepare_eplb_forward(input_batch.num_tokens)
+        self._prepare_eplb_forward(num_tokens)
 
+        self.on_prefill_begin(num_reqs)
         if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
             # Replay the full graph for draft prefill.
             assert self.prefill_cudagraph_manager is not None
@@ -240,6 +259,7 @@ def propose(
                 cudagraph_runtime_mode=prefill_batch_desc.cg_mode,
                 mm_inputs=mm_inputs,
             )
+        self.on_prefill_end(num_reqs)
 
         if self.num_speculative_steps == 1:
             # Early exit.
@@ -268,6 +288,7 @@ def propose(
             need_eager=is_profile,
         )
 
+        self.on_multi_step_decode_begin(num_reqs)
         # Generate the remaining num_speculative_steps - 1 draft tokens.
         self._multi_step_decode(
             num_reqs,
@@ -276,6 +297,7 @@ def propose(
             num_tokens_across_dp,
             input_batch.seq_lens_cpu_upper_bound,
         )
+        self.on_multi_step_decode_end(num_reqs)
 
         return self.draft_tokens[:num_reqs]
 

vllm/v1/worker/gpu/spec_decode/mtp/speculator.py

@@ -10,9 +10,49 @@
 
 
 class MTPSpeculator(AutoRegressiveSpeculator):
+    share_mtp_topk_indices: bool = False
+
     def load_draft_model(
         self,
         target_model: nn.Module,
         target_attn_layer_names: set[str],
     ) -> nn.Module:
-        return load_eagle_model(target_model, self.vllm_config)
+        draft_model = load_eagle_model(target_model, self.vllm_config)
+        spec_config = self.vllm_config.speculative_config
+        draft_hf_config = (
+            spec_config.draft_model_config.hf_config
+            if spec_config is not None
+            else None
+        )
+        # Detect index_share_for_mtp_iteration. When True, the proposer
+        # toggles skip_topk so step 0 computes MTP's own indices and
+        # steps 1+ reuse them.
+        self.share_mtp_topk_indices = (
+            getattr(draft_hf_config, "index_share_for_mtp_iteration", False)
+            and hasattr(draft_model.model, "set_skip_topk")
+            and hasattr(draft_model.model, "compact_topk_indices")
+        )
+        return draft_model
+
+    def on_prefill_begin(self, num_reqs: int) -> None:
+        # Step 0 computes its own top-k. Unconditional, so a step that died
+        # midway cannot leave reuse mode on.
+        if self.share_mtp_topk_indices:
+            self.model.model.set_skip_topk(False)
+
+    def on_prefill_end(self, num_reqs: int) -> None:
+        # Step 0 (prefill) wrote topk indices for every query token in the
+        # multi-token batch. Compact them down to each request's last token so
+        # steps 1+ can reuse them from the shared buffer.
+        if self.share_mtp_topk_indices and self.num_speculative_steps > 1:
+            self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs])
+
+    def on_multi_step_decode_begin(self, num_reqs: int) -> None:
+        # Switch to reuse mode so draft steps 1+ skip the indexer op and read
+        # the indices that step 0 wrote into the shared buffer.
+        if self.share_mtp_topk_indices:
+            self.model.model.set_skip_topk(True)
+
+    def on_multi_step_decode_end(self, num_reqs: int) -> None:
+        if self.share_mtp_topk_indices:
+            self.model.model.set_skip_topk(False)