[Bugfix][Spec Decode] Implement get_top_tokens() on the ROCm DeepSeek V4 MTP drafter (#57568)
Signed-off-by: BaoYunkai <ybao@amd.com>
BaoYunkai committed
Sep 28, 2026 at 20:16 UTC
f6877f563c5b27e858773b8416dc5ee3d76233ec
2 files changed
+200
-6
tests/v1/spec_decode/test_dsv4_mtp_local_argmax.py
new
+163
@@ -0,0 +1,163 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""``get_top_tokens`` on the DeepSeek V4 MTP drafter.
4
+
5
+The speculator opts into ``use_local_argmax_reduction`` by probing the draft
6
+model for ``get_top_tokens``, and expects it to select the same greedy draft
7
+token the full-vocab argmax would. V4 reaches its LM head through hc_head and
8
+the shared head's RMSNorm, so both methods have to share that transform.
9
+"""
10
+
11
+from types import SimpleNamespace
12
+from unittest import mock
13
+
14
+import pytest
15
+import torch
16
+from transformers import PreTrainedConfig
17
+
18
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
+from vllm.model_executor.layers.mhc import HCHeadOp
20
+from vllm.model_executor.models.deepseek_mtp import SharedHead
21
+from vllm.platforms import current_platform
22
+
23
+if not current_platform.is_rocm():
24
+ pytest.skip(
25
+ "get_top_tokens is implemented on the ROCm DeepSeek V4 MTP drafter only.",
26
+ allow_module_level=True,
27
+ )
28
+
29
+from vllm.models.deepseek_v4.amd.mtp import ( # noqa: E402
30
+ DeepSeekV4MTP,
31
+ DeepSeekV4MultiTokenPredictor,
32
+)
33
+
34
+DEVICE_TYPE = current_platform.device_type
35
+START_LAYER_IDX = 61
36
+HIDDEN_SIZE = 4096
37
+HC_MULT = 4
38
+RMS_NORM_EPS = 1e-6
39
+HC_EPS = 1e-6
40
+_EMBEDDING = "vllm.model_executor.layers.vocab_parallel_embedding"
41
+
42
+
43
+def _make_predictor(vocab_size: int) -> DeepSeekV4MultiTokenPredictor:
44
+ config = PreTrainedConfig()
45
+ config.hidden_size = HIDDEN_SIZE
46
+ config.vocab_size = vocab_size
47
+ config.rms_norm_eps = RMS_NORM_EPS
48
+
49
+ # A single vocab shard, so the reduction returns the local argmax without
50
+ # touching the process group.
51
+ with (
52
+ mock.patch(f"{_EMBEDDING}.get_tensor_model_parallel_rank", return_value=0),
53
+ mock.patch(
54
+ f"{_EMBEDDING}.get_tensor_model_parallel_world_size", return_value=1
55
+ ),
56
+ ):
57
+ shared_head = SharedHead(config=config, prefix="shared_head")
58
+
59
+ generator = torch.Generator().manual_seed(7)
60
+ with torch.no_grad():
61
+ shared_head.head.weight.normal_(0.0, 0.02, generator=generator)
62
+ # The default norm weight is all ones, and the argmax is invariant to
63
+ # the norm's per-row rescale, so a uniform gain would hide a skipped
64
+ # shared_head. Vary it so the comparison below can see one.
65
+ shared_head.norm.weight.normal_(1.0, 0.5, generator=generator)
66
+ shared_head = shared_head.to(device=DEVICE_TYPE, dtype=torch.bfloat16)
67
+
68
+ hc_head_fn = torch.randn(
69
+ HC_MULT, HC_MULT * HIDDEN_SIZE, dtype=torch.float32, generator=generator
70
+ )
71
+ # Non-trivial gate scale and bias: hc_head mixes the hc_mult streams with
72
+ # these, and defaults of 0 would make a skipped hc_head harder to see.
73
+ hc_head_scale = torch.randn(1, dtype=torch.float32, generator=generator) * 0.1
74
+ hc_head_base = torch.randn(HC_MULT, dtype=torch.float32, generator=generator) * 0.1
75
+
76
+ mtp_layer = SimpleNamespace(
77
+ config=config,
78
+ hc_mult=HC_MULT,
79
+ hc_eps=HC_EPS,
80
+ rms_norm_eps=RMS_NORM_EPS,
81
+ hc_head_op=HCHeadOp(),
82
+ hc_head_fn=(hc_head_fn * 1e-4).to(DEVICE_TYPE),
83
+ hc_head_scale=hc_head_scale.to(DEVICE_TYPE),
84
+ hc_head_base=hc_head_base.to(DEVICE_TYPE),
85
+ shared_head=shared_head,
86
+ )
87
+
88
+ # Only the attributes the two methods read; a real predictor would also
89
+ # build the MLA/MoE draft block, which is irrelevant here.
90
+ predictor = object.__new__(DeepSeekV4MultiTokenPredictor)
91
+ torch.nn.Module.__init__(predictor)
92
+ predictor.layers = {str(START_LAYER_IDX): mtp_layer}
93
+ predictor.mtp_start_layer_idx = START_LAYER_IDX
94
+ predictor.num_mtp_layers = 1
95
+ predictor.logits_processor = LogitsProcessor(vocab_size)
96
+ return predictor
97
+
98
+
99
+def _hidden_states(num_tokens: int, seed: int) -> torch.Tensor:
100
+ """The flat pre-hc_head residual the MTP layer returns."""
101
+ generator = torch.Generator().manual_seed(seed)
102
+ hidden_states = torch.randn(
103
+ num_tokens, HC_MULT * HIDDEN_SIZE, dtype=torch.float32, generator=generator
104
+ )
105
+ return hidden_states.to(device=DEVICE_TYPE, dtype=torch.bfloat16)
106
+
107
+
108
+def _assert_selects_max_logit(top: torch.Tensor, logits: torch.Tensor) -> None:
109
+ """The selected token must carry the row's maximum logit.
110
+
111
+ hc_head emits bf16, so rows can tie for the maximum and then no single
112
+ index is the right answer; compare the values the two paths select.
113
+ """
114
+ assert top.shape == (logits.shape[0],)
115
+ selected = logits.gather(1, top.unsqueeze(1)).squeeze(1)
116
+ assert torch.equal(selected, logits.max(dim=-1).values)
117
+
118
+
119
+def test_get_top_tokens_matches_full_argmax(default_vllm_config):
120
+ """Same draft token as the path taken when the reduction is off.
121
+
122
+ ``get_top_tokens`` has to run the same pre-head transform as
123
+ ``compute_logits``: hc_head over the (T, hc_mult, D) residual, then the
124
+ shared head's RMSNorm. Dropping either one returns a different draft token
125
+ instead of failing, which is what this comparison catches.
126
+ """
127
+ predictor = _make_predictor(vocab_size=512)
128
+
129
+ for num_tokens in (1, 16, 128):
130
+ hidden_states = _hidden_states(num_tokens, seed=11)
131
+
132
+ logits = predictor.compute_logits(hidden_states)
133
+ top = predictor.get_top_tokens(hidden_states)
134
+
135
+ assert top.dtype == torch.int64
136
+ _assert_selects_max_logit(top, logits)
137
+
138
+
139
+def test_get_top_tokens_honors_padded_vocab(default_vllm_config):
140
+ """Padding entries above ``org_vocab_size`` must never win the argmax."""
141
+ vocab_size = 500
142
+ predictor = _make_predictor(vocab_size)
143
+ head = predictor.layers[str(START_LAYER_IDX)].shared_head.head
144
+ assert head.shard_indices.num_org_vocab_padding > 0
145
+
146
+ hidden_states = _hidden_states(32, seed=13)
147
+ logits = predictor.compute_logits(hidden_states)
148
+ top = predictor.get_top_tokens(hidden_states)
149
+
150
+ _assert_selects_max_logit(top, logits)
151
+ assert (top < vocab_size).all()
152
+
153
+
154
+def test_wrapper_forwards_spec_step_idx():
155
+ """``DeepSeekV4MTP`` is what the speculator probes, so it needs the method."""
156
+ assert hasattr(DeepSeekV4MTP, "get_top_tokens")
157
+
158
+ hidden_states = torch.empty(0)
159
+ wrapper = object.__new__(DeepSeekV4MTP)
160
+ wrapper.model = SimpleNamespace(get_top_tokens=lambda *args: args)
161
+
162
+ assert DeepSeekV4MTP.get_top_tokens(wrapper, hidden_states, 2) == (hidden_states, 2)
163
+ assert DeepSeekV4MTP.get_top_tokens(wrapper, hidden_states) == (hidden_states, 0)
vllm/models/deepseek_v4/amd/mtp.py
+37
-6
@@ -240,13 +240,11 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
240
current_step_idx,
241
)
242
243
- def compute_logits(
243
+ def _pre_head(
244
self,
245
+ mtp_layer: DeepSeekV4MultiTokenPredictorLayer,
246
hidden_states: torch.Tensor,
246
- spec_step_idx: int = 0,
247
) -> torch.Tensor:
248
- current_step_idx = spec_step_idx % self.num_mtp_layers
249
- mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
248
# MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply
249
# hc_head here so logits are computed from the dense hidden state.
250
hidden_states = hidden_states.view(
@@ -260,14 +258,38 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
258
mtp_layer.rms_norm_eps,
259
mtp_layer.hc_eps,
260
)
263
- hidden_states = _MTP_SHARED_HEAD_RMSNORM_KERNEL(
261
+ return _MTP_SHARED_HEAD_RMSNORM_KERNEL(
262
hidden_states,
263
mtp_layer.shared_head.norm.weight.data,
264
mtp_layer.shared_head.norm.variance_epsilon,
265
)
268
- logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
266
+
267
+ def compute_logits(
268
+ self,
269
+ hidden_states: torch.Tensor,
270
+ spec_step_idx: int = 0,
271
+ ) -> torch.Tensor:
272
+ current_step_idx = spec_step_idx % self.num_mtp_layers
273
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
274
+ logits = self.logits_processor(
275
+ mtp_layer.shared_head.head, self._pre_head(mtp_layer, hidden_states)
276
+ )
277
return logits
278
279
+ def get_top_tokens(
280
+ self,
281
+ hidden_states: torch.Tensor,
282
+ spec_step_idx: int = 0,
283
+ ) -> torch.Tensor:
284
+ current_step_idx = spec_step_idx % self.num_mtp_layers
285
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
286
+ # Vocab-parallel argmax for the greedy draft: per-rank head projection
287
+ # + local argmax + a [batch, 2*tp] (value, index) reduce, instead of
288
+ # all-gathering the full vocab logits.
289
+ return self.logits_processor.get_top_tokens(
290
+ mtp_layer.shared_head.head, self._pre_head(mtp_layer, hidden_states)
291
+ )
292
+
293
294
class DeepSeekV4MTP(nn.Module):
295
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -302,6 +324,15 @@ class DeepSeekV4MTP(nn.Module):
324
) -> torch.Tensor | None:
325
return self.model.compute_logits(hidden_states, spec_step_idx)
326
327
+ def get_top_tokens(
328
+ self,
329
+ hidden_states: torch.Tensor,
330
+ spec_step_idx: int = 0,
331
+ ) -> torch.Tensor:
332
+ # Greedy-draft path used when use_local_argmax_reduction is enabled:
333
+ # vocab-parallel argmax, no full-vocab logits.
334
+ return self.model.get_top_tokens(hidden_states, spec_step_idx)
335
+
336
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
337
# Weight name remapping for checkpoint compatibility.
338
# Maps checkpoint weight paths to model parameter paths.