[Perf][Spec Decode] Fused vocab-parallel argmax for LogitsProcessor.get_top_tokens()

[pst2154]

Overview

On CUDA at TP > 1, compute LogitsProcessor.get_top_tokens() (the use_local_argmax_reduction path) with one Triton pass over the logits shard, a dim-0 all-gather of a [B, 16] fp32 candidate tensor and one small pick kernel, instead of the current sequence of small torch ops. Same result as torch.argmax over the full vocabulary. Split out of #59727 after review.

Claims

• Same algorithm as today (reduce each shard locally, gather per-row candidates, pick), with fewer kernels: 2 kernels plus the collective instead of 9 plus the collective (see Details).
• Same selected token as torch.argmax over the concatenated logits for every input, including ties within and across shards, inf and NaN (first NaN wins, ties go to the lowest id).
• Microbenchmark, one TP4 shard of a 131072-entry vocabulary, bf16, CUDA graph, communication excluded: the reduction takes about 75 % less time at 1-8 rows and about 60 % less at 64 rows.
• No end-to-end claim: this change was not measured end to end on its own.

Validation

On ac9126e58 (nightly image with this branch's vllm/ files overlaid), 1x / 4x B200:

• pytest -v tests/model_executor/layers/test_vocab_parallel_argmax.py: 27 passed. Simulated TP 2 / 3 / 4 / 8 shards, bf16 / fp16 / fp32, 131072 and 32001 vocab, 1-64 rows, against torch.argmax: random values, integer ties, ties across shards, all -inf, +inf ties, NaN, -0.0 / +0.0 ties; column-major and strided-column shards; 50 CUDA graph replays.
• pytest -v tests/distributed/test_vocab_parallel_argmax.py (TP2 / TP4, Ray): 2 passed. Real ParallelLMHead + LogitsProcessor.get_top_tokens() across ranks against the full-vocabulary greedy path (argmax of the gathered logits), with tie-heavy integer weights, 131072 and 32001 (padded) vocab, 1-8 rows, eagerly and from a captured CUDA graph.
• Microbenchmark (same shard as above, 20 calls captured per CUDA graph, communication replaced by a same-size copy in both paths):

Time per call relative to the existing reduction:

rows (B)	existing torch reduction	this PR
1	1.00	0.25
2	1.00	0.23
4	1.00	0.24
8	1.00	0.25
64	1.00	0.41

The existing path in this benchmark also gets 54 / 54 simulated-shard cases (NaN, ties, signed zeros) equal to torch.argmax, so the change does not alter the selected token.

Details

What the existing path does per call (after the head GEMM, soft-cap / scale and padding mask, which are unchanged):

• on each rank: logits.max(dim=-1) (1 reduction kernel), + vocab_start, .float() on the values and on the indices, torch.stack (4 elementwise kernels): 8 bytes per row;
• all_gather(dim=-1): NCCL gather along dim 0 plus the copy that moves the rank dimension next to the last dim;
• argmax over ranks, gather, .to(int64) (3 kernels).

That is 9 small kernels around the collective, each a separate launch.

This PR: local_argmax_candidates scans the shard once with 8 programs per row; each writes its best (value, global id) as fp32 (exact for bf16 / fp16 / fp32 values and ids below 2^24, the same assumption as today). The candidates (64 bytes per row and rank) are gathered along dim 0, so no rearranging copy is needed, and pick_from_candidates picks per row. The kernel takes row and column strides, so non-contiguous logits need no copy. TP = 1 and non-CUDA platforms keep the torch path.

Exactness: every comparison uses torch.argmax's total order with an index tie-break (NaN first, then larger value, then lower id, -0.0 == +0.0). The maximum under such an order over a partition of the vocabulary equals the maximum of the per-part maxima, so the result does not depend on the shard or chunk boundaries.

torch.compile: this path is not compiled. For MTP drafters such as Nemotron-H, @support_torch_compile is applied to the inner model (NemotronHMultiTokenPredictor), so only its forward is traced. The speculator's sample_draft() calls get_top_tokens() after that forward, outside the compiled region, so Inductor never fuses these ops. With FULL CUDA graphs the eager kernels are captured into the drafter's graph and replayed one by one, once per draft position; with PIECEWISE they run eagerly. The kernel count is therefore what runs at decode.

Race-free: two kernels plus vLLM's normal TP all-gather, all on the caller's stream; each kernel writes a fresh output; no atomics, no persistent state, no cross-program communication.

Related: #59727 (Nemotron-H MTP get_top_tokens(), needed to use this path with that drafter), #56995 (generic MTP drafter hook), #57770 (split-vocabulary argmax on already-gathered logits for DSpark; different call site).

Duplicate check: gh pr list --repo vllm-project/vllm --state open --search "get_top_tokens" -> #59727, #56995, #52072, #57770, #47584; --search "vocab parallel argmax" and --search "local argmax reduction" -> the same plus unrelated sampler / diffusion items. None changes the reduction inside LogitsProcessor.get_top_tokens().

AI assistance: Claude Code was used to write this change and its tests, run the validation and draft this description. The human submitter reviews every changed line before this leaves draft.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code