[Bugfix][Spec Decode] Implement get_top_tokens() on the generic MTP drafter

[ZhengGong-amd]

Purpose

speculative_config.use_local_argmax_reduction is unusable with the generic MTP
drafter: enabling it does not fall back to a slower path, it fails engine startup.

ValueError: use_local_argmax_reduction is enabled but draft model DeepSeekMTP
does not implement get_top_tokens().

The reduction itself already lives in core as LogitsProcessor.get_top_tokens() — each
TP rank argmaxes its own vocab shard and only (value, index) pairs are gathered. Both
opt-in sites probe the draft model for the method and raise without it:

• vllm/v1/worker/gpu/spec_decode/speculator.py (V2 model runner)
• vllm/v1/spec_decode/llm_base_proposer.py

vllm/model_executor/models/deepseek_mtp.py has compute_logits but no
get_top_tokens, so every checkpoint that resolves to it raises. On current main that
is the deepseek_v3 family — DeepSeek-V3 and R1 — on every platform:

checkpoint model_type	architecture	resolved draft class	platform
deepseek_v3	DeepSeekMTPModel	deepseek_mtp.DeepSeekMTP	all
deepseek_v32, glm_moe_dsa	DeepseekV32MTPModel	deepseek_mtp.DeepSeekMTP	XPU, CPU only

DeepSeekMTPModel maps to deepseek_mtp in registry.py with no platform dispatch, so a
DeepSeek-V3 / R1 MTP drafter is this class on CUDA and ROCm alike. deepseek_v32 and
glm_moe_dsa are dispatched in vllm/models/deepseek_v32/__init__.py: CUDA uses
nvidia/mtp.py, which already implements the method, and since #51915 ROCm uses
amd/mtp.py. So for the DSA family only XPU and CPU still fall back here. amd/mtp.py's
own DeepseekV32MTP has the same gap, but it is a different file with its own pre-head
path and is out of scope for this PR.

This PR adds the two forwarding methods, mirroring the compute_logits pair already in
the file. The capability has been rolled out per model since the flag landed in #34049,
and deepseek_v3 was never covered on any platform.

Which checkpoints this actually reaches

Resolving real config.jsons through hf_config_override and the registry on
main@94f4170df3 (ROCm), then asking the speculator's own question
(hasattr(cls, "get_top_tokens")). Of the current flagship checkpoints that ship MTP
layers at all:

verdict	checkpoints
reached by this PR	DeepSeek-R1-0528 (fp8 and MXFP4)
already had the method	GLM-5.3-Flash (glm5_next → Glm5NextMTP)
still blocked, different class	DeepSeek-V3.2, GLM-5.2, GLM-5.3 (ROCm: amd/mtp.py); DeepSeek-V4-Flash / -Pro, DeepSeek-V4.1-Flash, MiniMax-M3

What it does not reach matters as much. On ROCm the DSA family now has its own drafter
class, as do the DeepSeek-V4 line and MiniMax-M3; on CUDA the DSA family already works
through nvidia/mtp.py. Kimi-K3, Qwen3.5-397B, Qwen3.8-Flash-Next and Qwen3-Next-80B ship
no MTP layers in the checkpoints checked.

Scope

Resolving every entry of _SPECULATIVE_DECODING_MODELS on main@94f4170df3 and asking
hasattr, 27 of the 62 speculative architectures satisfy the guard and 34 do not, so this
flag is broadly unavailable rather than uniquely broken here. 24 of the 27 get it from
LocalArgmaxMixin (the EAGLE, DFlash and DSpark families, plus Qwen3.5 MTP); only
Gemma4MTP, InklingMTP and Glm5NextMTP write the method themselves (plus
DeepseekV32MTP in the CUDA-only nvidia/mtp.py).

This PR fixes one of the 34 rather than sweeping them, because the correct call differs
per family and getting it wrong is silent:

• deepseek_mtp.py's draft layer returns (pre-final-norm, post-final-norm) and the
  proposer consumes the first element, so compute_logits — and therefore
  get_top_tokens — must apply shared_head (the final RMSNorm).
• deepseek_v32/nvidia/mtp.py's layer returns the post-norm hidden in both tuple
  positions, so its get_top_tokens correctly applies the LM head only.

Reusing the second form in the first file does not raise; it just returns a different
draft token. The added test is built to fail in that case (see Test Result). Extending
this to the remaining families needs the same per-family check, which is better done in
separate PRs by people who can run those models.

A second commit amends the flag's own docstring. It currently says only "reduces
communication from O(vocab_size) to O(2 * tp_size)", which reads as unconditional; the
measurements below show it is a break-even-dependent trade, and low-concurrency serving
at high TP is the case where the current wording would mislead someone into enabling it.

LocalArgmaxMixin in models/interfaces.py is the right tool where it fits, and
Qwen3_5MTP shows it fits some MTP models: it owns a flat self.lm_head and its
compute_logits is just self.logits_processor(self.lm_head, hidden_states). This class
is not like that. It has no self.lm_head at all — the head is
self.layers[str(mtp_start + spec_step_idx)].shared_head.head, chosen per draft step,
and the hidden state has to pass through shared_head's norm first. The mixin reads
self.lm_head and applies no norm, so it cannot serve this class.

What this cannot affect

use_local_argmax_reduction defaults to False, and every call site of
get_top_tokens is behind it — no site probes for the method to auto-enable anything.
With the flag off, this PR adds two unreachable methods and changes nothing.

Test Plan

pytest tests/v1/spec_decode/test_mtp_local_argmax.py
pytest tests/v1/sample/test_head_dtype.py   # existing coverage of get_top_tokens

New tests/v1/spec_decode/test_mtp_local_argmax.py builds a real SharedHead
(RMSNorm + ParallelLMHead) and LogitsProcessor on a single vocab shard and asserts
get_top_tokens() picks the same token as compute_logits().argmax(-1), including with
a padded vocab. The norm's weight is given a non-uniform gain on purpose: the argmax is
invariant to the norm's per-row rescale, so a default all-ones weight would let a
skipped shared_head pass unnoticed.

End to end, on a deepseek_v3 checkpoint (DeepSeek-R1-0528, official fp8 weights) at TP=4,
with a fresh server per run:

vllm serve deepseek-ai/DeepSeek-R1-0528 -tp 4 --max-model-len 4096 \
  --speculative-config '{"method":"mtp","num_speculative_tokens":1,
                         "use_local_argmax_reduction":true}'
vllm bench serve --model deepseek-ai/DeepSeek-R1-0528 --dataset-name random \
  --random-input-len 1024 --random-output-len 512 --num-prompts 288 \
  --max-concurrency 64 --ignore-eos --seed 1234 \
  --percentile-metrics ttft,tpot,itl,e2el

The base arm is upstream main with the flag off; with it on, main raises at engine
init.

Test Result

Hardware for every number below: MI355X (gfx950), ROCm. Correctness and micro numbers
are at real model shapes in bf16 at TP=2 and TP=4; GSM8K and end to end are
DeepSeek-R1-0528 (official fp8 weights) at TP=4 against main@94f4170df3.

Correctness

tests/v1/spec_decode/test_mtp_local_argmax.py: 3 passed. Mutating the patch to drop
shared_head from get_top_tokens turns 2 of the 3 into failures, so the equivalence
assertion is load-bearing rather than decorative.

tests/v1/sample/test_head_dtype.py: identical results with and without this PR
(7 passed, 1 pre-existing failure in the core_model e2e test, which downloads
facebook/opt-125m).

Token-for-token equivalence over a sweep at real model shapes, TP=2, batch sizes 1 to
1024: agreement 1.0000 in all 20 cases. Cross-rank tie-breaking is the part a
single-shard test cannot reach, and it agrees too. The same sweep run against a variant
that skips shared_head agrees on as little as 0 % of tokens.

Before this PR, the vllm serve / vllm bench invocation above raises the ValueError
at engine init. After it, both workers log

Using local argmax reduction for draft token generation
(communication: O(2*tp_size) vs O(vocab_size)).

and generation completes.

Accuracy (GSM8K)

A vocab-parallel argmax is a reduction-order change, so "the unit tests pass" does not
cover it. On DeepSeek-R1-0528, all 1319 items, 5-shot greedy: no measurable accuracy
change on GSM8K, within the base-repeat noise floor. Patched scores 0.9568 against
0.9530 / 0.9560 for two runs of main — a base-vs-base drift of 0.0030, which the
patched-vs-second-base difference (+0.0008) sits well inside. Per-sample agreement is
inside the same control range. Full table in the comment below.

Micro benchmark

The two providers on a real vLLM process group with a real SharedHead and
LogitsProcessor, driven through the patched predictor so the only difference between
arms is the collective:

• baseline — compute_logits(h).argmax(-1), what the drafter does today
• patched — get_top_tokens(h)

Five interleaved pairs per point, arm order alternated per pair, 300 iterations after
100 warmup, no other tenant on the box. The delta column is what actually matters:
both arms pay the same per-call launch overhead (~200 µs in this harness, which times
the sampling step in isolation), so the percentage exaggerates the cost at small batch
and the absolute difference does not.

TP=2, GLM-5.2 shapes (hidden_size 6144, vocab_size 154880) / DeepSeek-V3 shapes
(7168, 129280):

batch	GLM base	GLM patch	GLM delta	DSV3 base	DSV3 patch	DSV3 delta
1	204.2 µs	228.4 µs	+24 µs	222.4 µs	219.2 µs	-3 µs
8	252.4 µs	270.7 µs	+18 µs	251.7 µs	271.1 µs	+19 µs
16	287.4 µs	269.1 µs	-18 µs	273.0 µs	273.4 µs	+0 µs
32	354.1 µs	271.7 µs	-82 µs	317.0 µs	270.4 µs	-47 µs
64	473.8 µs	287.0 µs	-187 µs	398.8 µs	271.7 µs	-127 µs
128	713.5 µs	318.1 µs	-395 µs	598.6 µs	299.1 µs	-300 µs
256	1202.4 µs	364.8 µs	-838 µs	922.6 µs	298.7 µs	-624 µs
512	2166.8 µs	535.2 µs	-1632 µs	1778.0 µs	494.3 µs	-1284 µs
1024	4124.6 µs	848.3 µs	-3276 µs	3289.3 µs	811.0 µs	-2478 µs

Within-arm spread over the five pairs is ≤ 1 % for 34 of the 40 arms and ≤ 4.7 %
everywhere, and the per-pair deltas are tight (e.g. GLM bs=128: +55.2, +55.4, +55.4,
+55.3, +55.4 %).

TP=4, same shapes, same protocol:

batch	GLM base	GLM patch	GLM delta	DSV3 base	DSV3 patch	DSV3 delta
1	218.5 µs	308.2 µs	+90 µs	224.5 µs	314.0 µs	+90 µs
32	286.3 µs	382.3 µs	+96 µs	294.7 µs	395.2 µs	+100 µs
64	317.6 µs	384.5 µs	+67 µs	295.2 µs	393.3 µs	+98 µs
128	443.9 µs	392.2 µs	-52 µs	379.6 µs	385.2 µs	+6 µs
256	710.0 µs	387.6 µs	-322 µs	573.1 µs	391.7 µs	-181 µs
512	1322.0 µs	385.9 µs	-936 µs	1003.5 µs	386.8 µs	-617 µs
1024	2583.3 µs	579.0 µs	-2004 µs	1854.6 µs	488.9 µs	-1366 µs

So the trade is a fixed per-step overhead — one extra small collective, ~20 µs at TP=2
and ~95 µs at TP=4 — against a gather that grows with batch × vocab_size. The patched
arm is nearly batch-independent (382–392 µs across bs=1..512 at TP=4); the baseline is
not. Break-even draft batch is ≈ 16–32 at TP=2 and ≈ 128–256 at TP=4, and it moves
right as tp grows because the reduction's own collective is latency-bound while the
gather it replaces is bandwidth-bound.

Being explicit about the cost: below break-even this is a regression, up to -41 % of
the sampling step at TP=4 bs=1 (which is +90 µs in absolute terms). A batch-size median
is a bad summary here since it just reflects which batch sizes one tabulates; the
crossover and the absolute delta are the honest statement. The flag is opt-in, so the
decision can be made per deployment, and the numbers above are what it should be made on.

Token agreement was asserted at every point before timing: 1.0000 in all 40 cases
across TP=2 and TP=4. Cross-rank tie-breaking is what a single-shard test cannot
reach, and it agrees at both rank counts.

Communication per draft step, derived from logits_processor.py rather than measured:
the default path materializes a [batch, vocab_size] logits tensor on every rank, the
reduction gathers [batch, 2 * tp_size] fp32 instead.

model	vocab_size	tp	batch	default	reduction
DeepSeek-V3	129280	2	512	126.2 MiB	8.0 KiB
GLM-5.2	154880	2	512	151.2 MiB	8.0 KiB
GLM-5.2	154880	8	1024	302.5 MiB	64.0 KiB

End to end

DeepSeek-R1-0528 at TP=4, the invocation in the Test Plan, median of three runs per arm
after one warm-up:

metric	main	this PR, flag on	Δ
Total tok/s	5764.9	5812.4	+0.8 %
Median TPOT (ms)	29.49	29.49	0.0 %
Median TTFT (ms)	332.0	328.7	−1.0 %

Flat, which is what the micro numbers predict for this point. Concurrency 64 puts the
draft batch at 64, below the TP=4 break-even, where the reduction costs ~+70–100 µs per
step — about 0.3 % of a 29.5 ms TPOT, under the ~5 % within-arm spread of throughput.
So at this operating point the flag costs nothing measurable. The gain needs a draft batch
above break-even, i.e. higher concurrency at TP=4, or TP=2. Per-run numbers, including
mean TPOT and TTFT, are in the comment below.

---

Found by Hyperloom; reviewed and measured by hand.

[ZhengGong-amd]

Why this cannot regress NVIDIA

The measurements in the description are MI355X, so it is worth being explicit about what
this change can and cannot do on CUDA. There is no NVIDIA card on the machine I used, so
everything below is read off the dispatch tables and the call sites rather than measured
— stating which is which.

1. The default path is untouched on every platform.
SpeculativeConfig.use_local_argmax_reduction defaults to False, and every call site
of get_top_tokens sits behind it (speculator.py, llm_base_proposer.py,
step3p5.py). Nothing probes for the method to auto-enable anything — the two
hasattr checks run only inside if self.use_local_argmax_reduction. With the flag
unset, this PR adds two methods that are never called.

2. The models that already had the method on CUDA are not touched.
deepseek_v32 and glm_moe_dsa checkpoints resolve to DeepseekV32MTPModel, and
vllm/models/deepseek_v32/__init__.py dispatches that to
vllm/models/deepseek_v32/nvidia/mtp.py under current_platform.is_cuda(). That module
implements get_top_tokens already. This PR edits
vllm/model_executor/models/deepseek_mtp.py, which is the non-CUDA target of that same
dispatch, so DeepSeek-V3.2 and GLM-5.2 on NVIDIA are unaffected either way.

3. The one NVIDIA-visible change is a configuration that used to be impossible.
To be accurate rather than reassuring: deepseek_v3 checkpoints (DeepSeek-V3, R1) are
not platform-dispatched. SpeculativeConfig.hf_config_override gives them
DeepSeekMTPModel, and registry.py maps that to ("deepseek_mtp", "DeepSeekMTP") with
no platform branch — so on CUDA too, a DeepSeek-V3 MTP drafter is this class. Today
use_local_argmax_reduction=True raises for it on NVIDIA exactly as it does on ROCm;
after this PR it starts and runs. That is a strict addition: no NVIDIA user who does not
set the flag sees any difference, and no NVIDIA user who does set it was working before.

This is also the case I exercised most heavily: DeepSeek-R1-0528 is a deepseek_v3
checkpoint, so the end-to-end runs in the description go through exactly the dispatch an
NVIDIA user would, and differ from CUDA only in the interconnect underneath the
collective. The logic being validated is platform-independent; only the constants are
not.

4. Nothing platform-specific is added.
The diff is two Python forwarding methods. No kernel, no current_platform branch, no
HIP or CUDA code. They call LogitsProcessor.get_top_tokens(), which is existing core
code with existing CUDA coverage in tests/v1/sample/test_head_dtype.py, and which four
other MTP families already call on CUDA.

5. What a CUDA reviewer may want to re-measure.
The mechanism — replacing a [batch, vocab_size] all-gather with a
[batch, 2 * tp_size] one — is interconnect-bound, so the crossover batch size will sit
somewhere else on NVLink than it does on this box. On MI355X the reduction starts paying
at draft batch ≈ 16–32 at TP=2 and ≈ 128–256 at TP=4; below that it costs a fixed
~20 µs (TP=2) to ~95 µs (TP=4) per draft step. The sign of the effect at small batch is
the part most likely to differ on NVIDIA, and it is also the part that matters least,
since the flag is opt-in.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[ZhengGong-amd]

Full GSM8K and end-to-end results for DeepSeek-R1-0528 (the deepseek_v3 → DeepSeekMTP path this PR changes), summarised in the description.

Setup. MI355X (gfx950), ROCm 7.2.3, torch 2.12.0, aiter 0.1.21.post2. deepseek-ai/DeepSeek-R1-0528, official fp8 weights, TP=4, mtp with one draft token. The base arm is upstream main@94f4170df3, meaning main's Python on the image's compiled v0.30.0 extensions. Between that tag and main no op schema changed in either torch_bindings.cpp, and every new _C/aiter symbol main calls is present or unreachable on ROCm. The patched arm is that tree plus this PR's diff, and diff -r between the two shows only deepseek_mtp.py and speculative.py. main rejects use_local_argmax_reduction for this drafter, so the base arm runs with the flag off.

GSM8K

lm_eval, 5-shot, greedy, all 1319 items, max_model_len=4096, max_gen_toks=256, prefix caching off, seed=1234. There are four arms: base2 repeats base to give the noise floor, and noop is the patched tree with the flag off, where the new method exists but is never called.

arm	tree	flag	strict-match	flexible-extract
base	main	off	0.9530	0.9568
base2	main	off	0.9560	—
noop	this PR	off	0.9538	—
patched	this PR	on	0.9568	0.9575

The base-vs-base drift is +0.0030. Patched is +0.0038 from base and +0.0008 from base2, both inside that drift. All four TP workers of the patched arm logged Using local argmax reduction.

Per-sample agreement (strict-match). Prompt hashes match on all 1319 items in every pair.

pair	identical text	same verdict
base vs base2	30.78 %	98.48 %	control: same build twice
base vs noop	31.54 %	97.95 %	control: patch present, never fires
base vs patched	29.19 %	98.26 %	measurement
base2 vs patched	31.01 %	97.95 %	measurement

Identical text is low even for the same build twice. vLLM's batching makes the reductions batch-shape dependent, so long generations diverge between runs of one build. The patched pairs sit inside the control range on both columns. GSM8K is 5-shot greedy short-answer math, so this says nothing about long context, sampling or multi-turn.

End to end

vllm bench serve, random dataset, 1024 in / 512 out, 288 prompts, --max-concurrency 64, --ignore-eos --seed 1234. Every run starts a fresh server and waits for /health, so no KV or prefix cache carries between runs; only the on-disk compile caches are kept, one per arm. Each arm is warmed once and then measured three times in a row, base then patched.

run	tok/s	mean TPOT	median TPOT	mean TTFT	median TTFT	acceptance
base 1	5568.9	29.92	29.59	1103.1	332.0	82.5 %
base 2	5868.0	29.11	29.28	689.7	329.3	82.9 %
base 3	5764.9	29.58	29.49	717.5	339.0	82.2 %
patched 1	5839.7	29.32	29.34	687.6	311.8	82.3 %
patched 2	5510.1	30.19	30.06	1094.3	328.7	81.1 %
patched 3	5812.4	29.32	29.49	702.3	341.1	82.0 %

Medians over the three runs (times in ms):

metric	base	spread	patched	spread	Δ
Total tok/s	5764.9	5.2 %	5812.4	5.7 %	+0.8 %
Mean TPOT	29.58	2.8 %	29.32	3.0 %	−0.9 %
Median TPOT	29.49	1.1 %	29.49	2.4 %	0.0 %
Mean TTFT	717.5	58 %	702.3	58 %	−2.1 %
Median TTFT	332.0	2.9 %	328.7	8.9 %	−1.0 %
Acceptance	82.5 %		82.0 %

Each arm has one run in a slower mode, with mean TTFT around 1.1 s and throughput around 5.5k (base 1, patched 2). That mode is what the 58 % mean-TTFT spread is. Median TTFT does not move with it, so it is a tail in a few requests, not a slower decode.

Reading: with the draft batch at 64, below the TP=4 break-even, the reduction's per-step cost (~70–100 µs in the micro benchmark) is about 0.3 % of TPOT, and nothing moves beyond the within-arm spread. The flag is neutral at this operating point. It does not show the gain, which needs a larger draft batch.