[Spec Decode][Perf] Fuse the MTP trailing all-reduce; local-argmax draft tokens

[zhou9402]

Two optimizations on the DeepSeek-V3.2 / GLM-5.2 MTP draft path.

• Fuse the trailing all-reduce into the final RMSNorm on the
  non-sequence-parallel path, as the main model already does at layer
  boundaries. The sequence-parallel path is unchanged.
• Greedy draft tokens via vocab-parallel local argmax (get_top_tokens),
  skipping the full-vocab all-gather in compute_logits. The proposer already
  probes for this via use_local_argmax_reduction; this adds the model side.

Not a duplicate: an earlier revision also carried the
index_share_for_mtp_iteration gates and the V2 speculator lifecycle hooks,
both of which have since landed on main; that part is dropped. Nothing open
overlaps with what remains.

gh pr list --repo vllm-project/vllm --state open --search "MTP all-reduce RMSNorm fuse"
gh pr list --repo vllm-project/vllm --state open --search "local argmax draft tokens"

Benchmarks

2×GB300, TP8 across 2 nodes, nvidia/GLM-5.2-NVFP4, MTP=5,
VLLM_USE_V2_MODEL_RUNNER=1, fp8 KV cache, no prefix caching. The baseline arm
is the same tree with the fusion reverted to an explicit
tensor_model_parallel_all_reduce + RMSNorm; local argmax is toggled with
--speculative-config.use_local_argmax_reduction.

Batch 1 (8192 in / 1024 out, c=1) — no measurable change. Three runs per
arm, output tok/s and MTP acceptance length:

arm	fusion	local argmax	output tok/s	acceptance length
baseline	✗	✗	456.7 / 451.6 / 463.3	4.78 / 4.70 / 4.86
+fusion	✓	✗	487.0 / 459.0 / 461.3	5.05 / 4.76 / 4.79
+both	✓	✓	460.1 / 472.8 / 470.2	4.80 / 4.95 / 4.90

The arms overlap completely. At c=1 the saved all-gather is one vocab row and
the fused all-reduce is one small tensor, so this is the expected result;
median ITL sat at 9.08–9.12 ms across all nine runs.

Concurrency 64 (1024 in / 1024 out, 256 prompts) — this is where it pays.
One A/B pair on the same node pair:

arm	output tok/s	acceptance length	median ITL
baseline	2923.2	4.86	31.83 ms
+both	3321.1 (+13.6%)	4.69	31.32 ms

The gain cannot be attributed to speculation luck: acceptance length is 3.5%
lower in the faster arm. Normalizing throughput by acceptance length
(steps/s) gives +17.8%. At c=64 the draft's per-step full-vocab all-gather is
64 × 151k × 2 B ≈ 19 MB, five times per accepted token — that is what
get_top_tokens removes.

Caveats, stated plainly: the c=64 result is a single A/B pair, and a repeat was
lost to cluster contention. Cross-node variability on this cluster is large
(the +fusion arm measured 4226 tok/s on a different node pair with the same
acceptance length as the baseline), so only within-pair comparisons above are
meaningful — do not compare numbers across the two tables.

Accuracy

gsm8k, 5-shot, full 1319 questions, via lm_eval --model local-completions:

arm	flexible-extract	strict-match
baseline	0.9439 ± 0.0063	0.9424 ± 0.0064
+both	0.9431 ± 0.0064	0.9409 ± 0.0065

Unchanged within stderr, which is the point: the fused kernel accumulates in
fp32 and local argmax breaks vocab-boundary ties differently, so both are
output-affecting in principle.

AI assistance (Claude) was used; every changed line has been reviewed.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhou9402.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[zhou9402]

Part of the #48597 re-split. See that PR for the tracker: merge order, measured numbers, and what was dropped.

[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.

[mergify]

Hi @zhou9402, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhou9402.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ZJY0516]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82463 for commit 54d503b96c47.

[mergify]

Hi @zhou9402, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[ZJY0516]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82503 for commit 99b66ac54de4.

[WoosukKwon]

To clarify, we are going to merge #47352 first

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhou9402.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[WoosukKwon]

Is this PR still needed?

[zhou9402]

Is this PR still needed?

Yes, but much simpler now — the index-sharing part has already merged into main, so all that's left here is the RMSNorm fusion and the local argmax.