[Spec Decode] NemotronH MTP: implement get_top_tokens()

[pst2154]

Overview

Implement NemotronHMTP.get_top_tokens() so that --speculative_config.use_local_argmax_reduction true works with the Nemotron-H MTP drafter. Split after review: the fused vocab-parallel argmax kernels that were also in this PR are now #60099. Related: #56995 (same hook for the generic MTP drafter).

Claims

• Nemotron-H MTP can draft greedily without all-gathering the full-vocabulary draft logits, through the existing use_local_argmax_reduction switch (no new flag). Before this change the speculator rejected the option because the drafter had no get_top_tokens().
• The selected draft tokens are the same as compute_logits(...).argmax(-1): the method delegates to the existing LogitsProcessor.get_top_tokens() with the shared lm_head, as other MTP drafters do.
• No performance claim for this change alone.

Validation

• Served on ac9126e58 (nightly image with this file overlaid), Nemotron-H (NVFP4) with MTP (7 draft tokens), TP4 on 4x B200, MRV2, --speculative_config.use_local_argmax_reduction true: startup, including CUDA graph capture of the drafter, succeeded, and the speculator's Using local argmax reduction for draft token generation ... line appeared on all 4 ranks.
• Without this change, the speculator's _validate_local_argmax_reduction() rejects the option for this drafter (... does not implement get_top_tokens()).

Details

• 8 lines in nemotron_h_mtp.py, mirroring compute_logits() (same lm_head assertion).
• Whether the flag pays off depends on batch size and TP size (see the discussion in [Bugfix][Spec Decode] Implement get_top_tokens() on the generic MTP drafter #56995); this PR does not change its default.
• The kernel count of LogitsProcessor.get_top_tokens() itself is addressed separately in [Perf][Spec Decode] Fused vocab-parallel argmax for LogitsProcessor.get_top_tokens() #60099.

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

[pst2154]

Thanks. Splitting as suggested: this PR will keep only the NemotronH MTP get_top_tokens() (enabling use_local_argmax_reduction), and the vocab-parallel argmax kernels move to a separate PR.

On the kernels: you're right that the existing path already reduces per shard first. The new kernels keep that algorithm and only cut the kernel count. Today it's max, +vocab_start, two .float(), stack, an NCCL all-gather along dim -1 plus the copy that rearranges it, then argmax, gather, .to(int64): 9 small kernels plus the collective. The new version does one Triton pass over the shard, a dim-0 all-gather of 64 B/row and one pick kernel, with the same result as torch.argmax (ties and NaN included). One TP4 shard of a 131k vocab under CUDA graph, comm excluded: about 22 µs -> 5.5 µs at BS 1-8. Strides: done, the kernel takes row and column strides and .contiguous() is gone. torch.compile: this section isn't compiled. @support_torch_compile wraps only the drafter's inner model forward, and get_top_tokens() is called from the speculator's sample_draft() outside that region, so under FULL CUDA graphs these kernels are captured as-is and replayed once per draft position. Updated branches coming shortly.

[mergify]

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

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