[Model Runner V2][Spec Decode] Handle tuple hidden states from MTP draft models (#46786)
Chauncey committed
Jun 28, 2026 at 09:38 UTC
11a12305c0522c5c1ed273d7d3dc2304ac0cd495
4 files changed
+85
-22
tests/v1/worker/test_gpu_autoregressive_speculator.py
new
+82
@@ -0,0 +1,82 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from contextlib import nullcontext
5
+from types import SimpleNamespace
6
+
7
+import torch
8
+
9
+from vllm.config.compilation import CUDAGraphMode
10
+from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module
11
+from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
12
+ AutoRegressiveSpeculator,
13
+)
14
+
15
+
16
+class _TestSpeculator(AutoRegressiveSpeculator):
17
+ def load_draft_model(self, target_model, target_attn_layer_names):
18
+ raise NotImplementedError
19
+
20
+
21
+class _DraftModel(torch.nn.Module):
22
+ def __init__(self, output: torch.Tensor | tuple[torch.Tensor, torch.Tensor]):
23
+ super().__init__()
24
+ self.output = output
25
+
26
+ def forward(self, **kwargs):
27
+ return self.output
28
+
29
+
30
+def _make_speculator(
31
+ monkeypatch,
32
+ output: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
33
+) -> _TestSpeculator:
34
+ monkeypatch.setattr(
35
+ spec_module,
36
+ "set_forward_context",
37
+ lambda *args, **kwargs: nullcontext(),
38
+ )
39
+
40
+ speculator = object.__new__(_TestSpeculator)
41
+ speculator.supports_mm_inputs = False
42
+ speculator.vllm_config = None
43
+ speculator.input_buffers = SimpleNamespace(
44
+ input_ids=torch.arange(4),
45
+ positions=torch.arange(4),
46
+ )
47
+ speculator.hidden_states = torch.zeros(4, 3)
48
+ speculator.model = _DraftModel(output)
49
+ return speculator
50
+
51
+
52
+def test_run_model_unpacks_tuple_return_for_mtp(monkeypatch):
53
+ logits_hidden = torch.full((4, 3), 1.0)
54
+ feedback_hidden = torch.full((4, 3), 2.0)
55
+ speculator = _make_speculator(monkeypatch, (logits_hidden, feedback_hidden))
56
+
57
+ actual_logits_hidden, actual_feedback_hidden = speculator._run_model(
58
+ 4,
59
+ attn_metadata=None,
60
+ slot_mappings=None,
61
+ num_tokens_across_dp=None,
62
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
63
+ )
64
+
65
+ assert actual_logits_hidden is logits_hidden
66
+ assert actual_feedback_hidden is feedback_hidden
67
+
68
+
69
+def test_run_model_reuses_tensor_return_for_mtp(monkeypatch):
70
+ hidden = torch.full((4, 3), 1.0)
71
+ speculator = _make_speculator(monkeypatch, hidden)
72
+
73
+ actual_logits_hidden, actual_feedback_hidden = speculator._run_model(
74
+ 4,
75
+ attn_metadata=None,
76
+ slot_mappings=None,
77
+ num_tokens_across_dp=None,
78
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
79
+ )
80
+
81
+ assert actual_logits_hidden is hidden
82
+ assert actual_feedback_hidden is hidden
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+3
-11
@@ -60,16 +60,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
60
"""
61
return True
62
63
- @property
64
- def model_returns_tuple(self) -> bool:
65
- """
66
- Whether the draft model's forward() returns a tuple.
67
-
68
- True: returns (last_hidden_states, hidden_states) — Eagle, Gemma4 MTP.
69
- False: returns a single tensor used for both — standard MTP (DeepSeek).
70
- """
71
- return True
72
-
63
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
64
# Initialize cudagraph manager for draft prefill (draft position 0).
65
self.prefill_cudagraph_manager = PrefillSpeculatorCudaGraphManager(
@@ -328,7 +318,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
318
else:
319
# Eager (NONE): call the raw model directly.
320
ret_hidden_states = self.model(**model_inputs)
331
- if self.model_returns_tuple:
321
+ # Some MTP models declare a single-tensor contract but return
322
+ # (logits_hidden, feedback_hidden) for final-norm correctness.
323
+ if isinstance(ret_hidden_states, tuple):
324
last_hidden_states, hidden_states = ret_hidden_states
325
else:
326
last_hidden_states = ret_hidden_states
vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py
-7
@@ -30,13 +30,6 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
30
# No new KV slots are written, so positions and seq_lens stay fixed.
31
return False
32
33
- @property
34
- def model_returns_tuple(self) -> bool:
35
- # forward() returns (draft_hidden_states, backbone_hidden_states).
36
- # The proposer uses draft_hidden_states for compute_logits and
37
- # backbone_hidden_states for the hidden-state feedback buffer.
38
- return True
39
-
33
def load_draft_model(
34
self,
35
target_model: nn.Module,
vllm/v1/worker/gpu/spec_decode/mtp/speculator.py
-4
@@ -10,10 +10,6 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
10
11
12
class MTPSpeculator(AutoRegressiveSpeculator):
13
- @property
14
- def model_returns_tuple(self) -> bool:
15
- return False
16
-
13
def load_draft_model(
14
self,
15
target_model: nn.Module,