[jiacao-amd]

Purpose

With a vocab-parallel LM head, greedy speculative decoding all-gathers the full-vocabulary logits across TP ranks: on the target side before rejection sampling, and (without use_local_argmax_reduction) on the draft side before the argmax. For a 200k vocabulary at TP4 that is a large gather on every decode step, although greedy verification only needs the argmax of each row.

This PR keeps the logits on each rank:

1. Sharded greedy verification (v1/worker/gpu/spec_decode/sharded_greedy_verifier.py, opt-in via VLLM_SHARDED_GREEDY_VERIFY=1). Each rank reduces its vocabulary shard to compact per-block (max, token id) statistics, only those are all-gathered, and one kernel finishes the global argmax, acceptance (including synthetic acceptance), bonus token and sampled/rejected counts. The blocks keep the native kernels' global boundaries (8192 for draft positions, 1024 for the bonus), also across shard edges, so ties resolve exactly as on the full logits.
   • Models opt in through compute_logits_local, the hook batch-sharded sampling already uses (added here to the ROCm MiniMax-M3 model, matching the NVIDIA one). Any TP size > 1 (including non-power-of-two), vocabulary < 2^24, hidden size and number of speculative tokens; CUDA and ROCm.
   • Per step, the batch falls back to the native path, before any logits are computed, unless every request is greedy with no logits processing, logprobs, structured output, draft logits, adaptive/block verification or watermarking.
2. Fused LogitsProcessor.get_top_tokens (model_executor/layers/vocab_parallel_argmax.py), i.e. the existing speculative_config.use_local_argmax_reduction draft path. The in-place padding mask, the torch reduction, the stack and the last-dim all-gather become two Triton kernels with a dim-0 all-gather. Same result (lowest id on ties, NaN wins, as torch.argmax); it applies to every model using LocalArgmaxMixin.

Results (MI355X, MiniMax-M3-MXFP4 TP4, EAGLE3-GQA, K=3)

Base image vllm/vllm-openai-rocm:nightly-rocm72-1b2a666391f2a24228a1436b64d547c34cc6b6ce plus the AITER gfx950 unified-attention routing fix (ROCm/aiter#6336) so an unrelated draft-attention regression does not dilute the numbers. This branch is based on the same vLLM commit; the measured image contains exactly these files. Server: VLLM_SHARDED_GREEDY_VERIFY=1, --speculative-config '{..., "use_local_argmax_reduction": true}'.

8k/1k C24 (240 requests, zero HTTP warmup, warmed compile caches, interleaved rounds vs same-round baseline):

	output throughput	rounds
this PR (verification + fused get_top_tokens)	+1.75%	+1.47 / +2.02
first, MiniMax-specific revision of this PR (same batch)	+1.78%	+1.68 / +1.88
verification + unfused use_local_argmax_reduction	+1.25%	+1.46 / +1.04
verification only	+0.65%	+0.65 / +0.64

Mean TPOT −3.0%, p99 E2EL −1.3%. Per decode step (profiler, 24 req × 4 tokens), NCCL in the cycle went from 385 µs to 83 µs.

AgenticX C24 (InferenceX recipe minimaxm3/vllm/mi355x-fp4-mtp/agentic.yaml, 3600 s, one run of this revision vs the same baseline): P50 interactivity +1.9%, P90 interactivity +0.7%, mean TPOT −0.5%; throughput is trace-saturated (+0.3%). TTFT regresses: mean +6.8%, P90 +8.1%. Two earlier runs that combined this change with a draft-metadata change showed the same TTFT shift (+6.3% / +5.3% mean), while other changes measured against the same baseline did not, so the TTFT cost is attributed to this PR and is being investigated before this leaves draft.

Correctness: real-acceptance greedy EAGLE3 smoke (24 concurrent ≥8k-token prompts + 4 short requests): generated text and token counts identical to the baseline (24/24, 4/4).

Test Plan

• tests/v1/worker/test_gpu_sharded_greedy_verify.py: TP2/3/4/8 emulated in one process (per-shard statistics concatenated in rank order, as all_gather(dim=0) produces), compared with rejection_sample + get_num_sampled_and_rejected on the full logits. Covers vocab 1000/50257/200064 (including shards narrower than a block), K=1/3/5, synthetic acceptance, prefilling requests and placeholder drafts, bf16 ties and planted cross-shard ties, unaligned vocabulary padding, and NaN rows staying in-vocabulary.
• tests/model_executor/layers/test_vocab_parallel_argmax.py: fused argmax vs torch.argmax for TP1/2/3/4/8, fp32/bf16, ties, NaN, padding, empty batch and all -inf rows.
• 157 passed on MI355X.

Draft: E2E measured on MiniMax-M3 / ROCm only so far.

🤖 Generated with Claude Code

[jiacao-amd]

Closing for now. The target workload for this work is the InferenceX AgenticX benchmark, which sends no sampling parameters, so requests use the model's generation_config defaults (temperature 1.0, top_p 0.95). The sharded greedy verification path therefore never triggers there, and the remaining draft-side change measured P90 interactivity +0.7% with mean TTFT +6.8% in a standalone AgenticX run. The 8k/1k greedy gain (+1.75%) stands, and the branch is kept in case greedy serving becomes a priority.