[Bugfix][Spec Decode] Per-module LM heads for multi-layer MTP on Model Runner V2 (#58921)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: chanh <channguyen@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

vllm/model_executor/models/step3p5_mtp.py

@@ -22,7 +22,11 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.sequence import IntermediateTensors
 
 from .step3p5 import Step3p5DecoderLayer
-from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix
+from .utils import (
+    get_draft_quant_config,
+    get_spec_layer_idx_from_weight_name,
+    maybe_prefix,
+)
 
 logger = init_logger(__name__)
 
@@ -152,8 +156,19 @@ class Step3p5AMultiTokenPredictor(nn.Module):
 
 
 class Step3p5MTP(nn.Module):
+    # Each MTP layer's shared_head carries its own LM head weights.
+    has_own_lm_head = True
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
+        # The target config may be a multimodal wrapper (Step3.7); build the
+        # draft from its own text and quant configs.
+        assert vllm_config.speculative_config is not None
+        draft_quant_config = get_draft_quant_config(vllm_config)
+        vllm_config = vllm_config.with_hf_config(
+            vllm_config.speculative_config.draft_model_config.hf_config
+        )
+        vllm_config.quant_config = draft_quant_config
         self.config = vllm_config.model_config.hf_config
         self.vllm_config = vllm_config
         self.model = Step3p5AMultiTokenPredictor(

vllm/models/inkling/amd/mtp.py

@@ -323,7 +323,11 @@ class InklingMTP(nn.Module, SupportsMultiModalEmbeddings):
             logits = logits[..., : self.logits_processor.org_vocab_size]
         return logits
 
-    def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def get_top_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
         """Greedy draft tokens via rank-local argmax + tiny (value, index)
         reduction — no full-vocab logits all-gather. The muP divisor is a
         positive scalar, so the argmax is invariant and the scaling is

vllm/models/inkling/nvidia/mtp.py

@@ -323,7 +323,11 @@ class InklingMTP(nn.Module, SupportsMultiModalEmbeddings):
             logits = logits[..., : self.logits_processor.org_vocab_size]
         return logits
 
-    def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def get_top_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
         """Greedy draft tokens via rank-local argmax + tiny (value, index)
         reduction — no full-vocab logits all-gather. The muP divisor is a
         positive scalar, so the argmax is invariant and the scaling is

vllm/v1/worker/gpu/spec_decode/eagle/utils.py

@@ -13,10 +13,14 @@ from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
 
 
 def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
-    """Share when the draft has no own copy, or its copy matches the target."""
-    if not getattr(eagle, flag, False) or draft is None:
+    """Share unless the draft declares its own copy that differs from the target.
+
+    A draft that declares its own copy but has no top-level one (e.g. MTP
+    heads stored per layer) keeps it.
+    """
+    if not getattr(eagle, flag, False):
         return True
-    if target is None:
+    if draft is None or target is None:
         return False
     # torch.equal on GPU allocates a bool mask the size of the input.
     # Use the faster GPU path when there is plenty of headroom;

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

@@ -253,6 +253,17 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
             )
         return self.draft_tokens[:num_reqs]
 
+    # Each MTP module may carry its own LM head, selected by spec_step_idx.
+    def compute_draft_logits(
+        self, hidden_states: torch.Tensor, spec_step_idx: int
+    ) -> torch.Tensor:
+        return self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx)
+
+    def get_draft_top_tokens(
+        self, hidden_states: torch.Tensor, spec_step_idx: int
+    ) -> torch.Tensor:
+        return self.model.get_top_tokens(hidden_states, spec_step_idx=spec_step_idx)
+
     @torch.inference_mode()
     def _run_model(
         self,
@@ -427,6 +438,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
                 self.seeds,
                 self.current_draft_step,
                 self.draft_logits,
+                spec_step_idx=step,
             )
 
             self.draft_tokens[:num_reqs, step] = draft_tokens

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

@@ -403,6 +403,16 @@ class DraftModelSpeculator(BaseSpeculator):
             "(communication: O(2*tp_size) vs O(vocab_size))."
         )
 
+    def compute_draft_logits(
+        self, hidden_states: torch.Tensor, spec_step_idx: int
+    ) -> torch.Tensor:
+        return self.model.compute_logits(hidden_states)
+
+    def get_draft_top_tokens(
+        self, hidden_states: torch.Tensor, spec_step_idx: int
+    ) -> torch.Tensor:
+        return self.model.get_top_tokens(hidden_states)
+
     def sample_draft(
         self,
         hidden_states: torch.Tensor,
@@ -412,9 +422,13 @@ class DraftModelSpeculator(BaseSpeculator):
         seeds: torch.Tensor,
         draft_step: torch.Tensor,
         draft_logits: torch.Tensor | None,
+        spec_step_idx: int = 0,
     ) -> torch.Tensor:
+        if draft_logits is None and self.use_local_argmax_reduction:
+            return self.get_draft_top_tokens(hidden_states, spec_step_idx)
+
+        logits = self.compute_draft_logits(hidden_states, spec_step_idx)
         if draft_logits is not None:
-            logits = self.model.compute_logits(hidden_states)
             sampler = (
                 gumbel_sample
                 if self.draft_watermarker is None
@@ -432,10 +446,7 @@ class DraftModelSpeculator(BaseSpeculator):
                 logits_cache_col=draft_step,
                 use_fp64=self.use_fp64_gumbel,
             )
-        elif self.use_local_argmax_reduction:
-            return self.model.get_top_tokens(hidden_states)
         else:
-            logits = self.model.compute_logits(hidden_states)
             sampled = logits.argmax(dim=-1)
         self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
         return sampled