[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>
Jiangyun Zhu committed
Oct 1, 2026 at 16:09 UTC
c4fc3c7d322bbcd1b25307ccbd9770c8eebc5c09
6 files changed
+60
-10
vllm/model_executor/models/step3p5_mtp.py
+16
-1
@@ -22,7 +22,11 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
22
from vllm.sequence import IntermediateTensors
23
24
from .step3p5 import Step3p5DecoderLayer
25
-from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix
25
+from .utils import (
26
+ get_draft_quant_config,
27
+ get_spec_layer_idx_from_weight_name,
28
+ maybe_prefix,
29
+)
30
31
logger = init_logger(__name__)
32
@@ -152,8 +156,19 @@ class Step3p5AMultiTokenPredictor(nn.Module):
156
157
158
class Step3p5MTP(nn.Module):
159
+ # Each MTP layer's shared_head carries its own LM head weights.
160
+ has_own_lm_head = True
161
+
162
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
163
super().__init__()
164
+ # The target config may be a multimodal wrapper (Step3.7); build the
165
+ # draft from its own text and quant configs.
166
+ assert vllm_config.speculative_config is not None
167
+ draft_quant_config = get_draft_quant_config(vllm_config)
168
+ vllm_config = vllm_config.with_hf_config(
169
+ vllm_config.speculative_config.draft_model_config.hf_config
170
+ )
171
+ vllm_config.quant_config = draft_quant_config
172
self.config = vllm_config.model_config.hf_config
173
self.vllm_config = vllm_config
174
self.model = Step3p5AMultiTokenPredictor(
vllm/models/inkling/amd/mtp.py
+5
-1
@@ -323,7 +323,11 @@ class InklingMTP(nn.Module, SupportsMultiModalEmbeddings):
323
logits = logits[..., : self.logits_processor.org_vocab_size]
324
return logits
325
326
- def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
326
+ def get_top_tokens(
327
+ self,
328
+ hidden_states: torch.Tensor,
329
+ spec_step_idx: int = 0,
330
+ ) -> torch.Tensor:
331
"""Greedy draft tokens via rank-local argmax + tiny (value, index)
332
reduction — no full-vocab logits all-gather. The muP divisor is a
333
positive scalar, so the argmax is invariant and the scaling is
vllm/models/inkling/nvidia/mtp.py
+5
-1
@@ -323,7 +323,11 @@ class InklingMTP(nn.Module, SupportsMultiModalEmbeddings):
323
logits = logits[..., : self.logits_processor.org_vocab_size]
324
return logits
325
326
- def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
326
+ def get_top_tokens(
327
+ self,
328
+ hidden_states: torch.Tensor,
329
+ spec_step_idx: int = 0,
330
+ ) -> torch.Tensor:
331
"""Greedy draft tokens via rank-local argmax + tiny (value, index)
332
reduction — no full-vocab logits all-gather. The muP divisor is a
333
positive scalar, so the argmax is invariant and the scaling is
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+7
-3
@@ -13,10 +13,14 @@ from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
13
14
15
def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
16
- """Share when the draft has no own copy, or its copy matches the target."""
17
- if not getattr(eagle, flag, False) or draft is None:
16
+ """Share unless the draft declares its own copy that differs from the target.
17
+
18
+ A draft that declares its own copy but has no top-level one (e.g. MTP
19
+ heads stored per layer) keeps it.
20
+ """
21
+ if not getattr(eagle, flag, False):
22
return True
19
- if target is None:
23
+ if draft is None or target is None:
24
return False
25
# torch.equal on GPU allocates a bool mask the size of the input.
26
# Use the faster GPU path when there is plenty of headroom;
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
+12
@@ -253,6 +253,17 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
253
)
254
return self.draft_tokens[:num_reqs]
255
256
+ # Each MTP module may carry its own LM head, selected by spec_step_idx.
257
+ def compute_draft_logits(
258
+ self, hidden_states: torch.Tensor, spec_step_idx: int
259
+ ) -> torch.Tensor:
260
+ return self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx)
261
+
262
+ def get_draft_top_tokens(
263
+ self, hidden_states: torch.Tensor, spec_step_idx: int
264
+ ) -> torch.Tensor:
265
+ return self.model.get_top_tokens(hidden_states, spec_step_idx=spec_step_idx)
266
+
267
@torch.inference_mode()
268
def _run_model(
269
self,
@@ -427,6 +438,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
438
self.seeds,
439
self.current_draft_step,
440
self.draft_logits,
441
+ spec_step_idx=step,
442
)
443
444
self.draft_tokens[:num_reqs, step] = draft_tokens
vllm/v1/worker/gpu/spec_decode/speculator.py
+15
-4
@@ -403,6 +403,16 @@ class DraftModelSpeculator(BaseSpeculator):
403
"(communication: O(2*tp_size) vs O(vocab_size))."
404
)
405
406
+ def compute_draft_logits(
407
+ self, hidden_states: torch.Tensor, spec_step_idx: int
408
+ ) -> torch.Tensor:
409
+ return self.model.compute_logits(hidden_states)
410
+
411
+ def get_draft_top_tokens(
412
+ self, hidden_states: torch.Tensor, spec_step_idx: int
413
+ ) -> torch.Tensor:
414
+ return self.model.get_top_tokens(hidden_states)
415
+
416
def sample_draft(
417
self,
418
hidden_states: torch.Tensor,
@@ -412,9 +422,13 @@ class DraftModelSpeculator(BaseSpeculator):
422
seeds: torch.Tensor,
423
draft_step: torch.Tensor,
424
draft_logits: torch.Tensor | None,
425
+ spec_step_idx: int = 0,
426
) -> torch.Tensor:
427
+ if draft_logits is None and self.use_local_argmax_reduction:
428
+ return self.get_draft_top_tokens(hidden_states, spec_step_idx)
429
+
430
+ logits = self.compute_draft_logits(hidden_states, spec_step_idx)
431
if draft_logits is not None:
417
- logits = self.model.compute_logits(hidden_states)
432
sampler = (
433
gumbel_sample
434
if self.draft_watermarker is None
@@ -432,10 +446,7 @@ class DraftModelSpeculator(BaseSpeculator):
446
logits_cache_col=draft_step,
447
use_fp64=self.use_fp64_gumbel,
448
)
435
- elif self.use_local_argmax_reduction:
436
- return self.model.get_top_tokens(hidden_states)
449
else:
438
- logits = self.model.compute_logits(hidden_states)
450
sampled = logits.argmax(dim=-1)
451
self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
452
return sampled