[zyongye]

Overview

This PR speeds up the DSpark/DFlash draft step under MRV2 and fixes a top-p masking bug. It contains:

• the Markov bias added inside the Gumbel kernel;
• wider DFlash input padding;
• is_padding written by the positions kernel;
• a fix for the Triton top-k/top-p kernel.

Split out of #63232 (MiniMax-M3.1-Flash). Based on main; it does not depend on the other PRs of that split. Batch-sharded DSpark draft sampling was dropped from this PR; #63886 covers it.

Claims

• Markov sampling:
  • The Gumbel kernel adds the Markov bias itself and caches the biased logits it samples from, so verification stays lossless.
  • One kernel combines the per-block maxima, replacing argmax + gather.
  • The temperature divide is skipped at 1.0.
  • Draft ids and cached draft logits are bit-identical to adding the bias first.
• DFlash: prepare_dflash_inputs pads the per-request buffers in 1024-wide blocks. Before, it used BLOCK_SIZE, which is 8 in a decode step, so the padding took thousands of serial iterations.
• MRV2: with VLLM_MOE_SKIP_PADDING (on by default), the positions kernel writes is_padding, replacing two fill_ launches per step. The result is the same, including after adaptive compaction.
• Triton top-k/top-p, more than 64 rows (bug fix): a top-p below the fp32 resolution of the kept mass (e.g. 1e-9) could mask every token of a row. The kernel now keeps the max token(s) in that case. It is also warmed up with V2's contiguous FP32 row stride.

Validation

Unit tests on GB300:

pytest tests/v1/spec_decode/test_dspark_markov_sampling.py tests/v1/worker/test_gpu_batch_shard.py \
  tests/v1/worker/test_gpu_gumbel_sample.py tests/v1/spec_decode/test_rejection_sampler_utils.py \
  tests/v1/worker/test_gpu_input_batch_v2.py tests/v1/sample/test_topk_topp_sampler.py \
  tests/v1/spec_decode/test_dflash2.py tests/v1/spec_decode/test_dflash_prepare_inputs.py \
  tests/v1/spec_decode/test_dspark_topk.py tests/v1/worker/test_gpu_model_runner_v2_eplb.py

460 passed, 1 skipped, 0 failed.

New tests:

• test_dspark_markov_sampling.py::test_markov_sampling_matches_unfused runs the Markov loop with the real LM-head, Markov embedding and Markov bias layers. For greedy and probabilistic drafting, with full and reduced draft vocabularies, it checks that the loop gives the same draft ids and cached draft-logits rows as adding the bias first.
• test_gpu_gumbel_sample.py::test_gumbel_combine_preserves_first_nan_and_maximum checks that the combine kernel follows torch.argmax: the first NaN wins, otherwise the first maximum.
• test_rejection_sampler_utils.py::test_block_verification_is_unbiased_across_steps drafts with gumbel_sample over many steps and checks that block verification keeps the target distribution.
• test_gpu_input_batch_v2.py::test_prepare_positions_preserves_padding_after_compaction checks the is_padding writes.
• test_topk_topp_sampler.py::test_tiny_p_keeps_max_token fails on main's kernel (2 failed) and passes here.

Performance:

	before	after	measured on
Markov loop, kernels per draft step	7	4	this branch, torch profiler, 1 request, 3 steps: 17 kernels for the loop vs 24 with the bias added outside the kernel; main's 7 also has argmax + gather in place of the combine kernel
DFlash input preparation per draft step	21 us	6 us	the internal branch these changes come from, GB300, MiniMax-M3.1 DSpark, 1 request, nsys

The 4 kernels per step are the Markov embedding, the bias GEMM, the Gumbel kernel and the combine kernel. The fused path removes the bias add and the per-step copy into draft_tokens.

Not covered: no end-to-end benchmark or accuracy eval on this branch.

Details

• Not a duplicate:

  • [Spec Decode][DSpark] Batch-sharded sequential block #63886 (WIP) batch-shards DSpark draft sampling across TP ranks. This PR no longer touches that; its Markov-loop change is in the replicated loop and the shared Gumbel kernel.
  • Other open PRs cut the cost of DSpark's Markov sampling in different ways: [MRV2][Spec Decode] Optimize DSpark with local argmax reduction #61824 (draft) is a greedy-only local argmax over vocab shards in the generic DSpark speculator; [ROCm][Perf][Kimi-K3] DSpark Markov sampling sharded over the vocabulary, one launch #64031 is Kimi-K3 on ROCm, sharded over the vocabulary inside the model; [Perf][DSpark][DSv4.1] Fuse the Markov sampling loop on NVIDIA by inverse CDF #64442 and [Perf][DSv4.1] DSpark draft on the specialized decode #61869 are fused loops specific to DeepSeek-V4.1.

  This PR keeps the generic Gumbel loop and its draws, so it works with probabilistic drafting and any DSpark draft model.

• When the bias is fused: probabilistic full-vocab drafting with no watermarking, acceptance estimator or output logit transform. Greedy drafting, a reduced draft vocab and those features still add the bias first, because they need the biased logits.

• Shared code touched:

  • gumbel_sample gains an optional bias.
  • Every caller now uses the combine kernel, which has torch.argmax's tie and NaN semantics.
  • prepare_dflash_inputs padding.
  • prepare_pos_seq_lens gains optional is_padding outputs.
  • The Triton top-k/top-p kernel.

• Size: about 220 non-test lines.

Credits: qizixi (Markov bias in the Gumbel kernel), plus Dao Le, Taoyu Zhu, Zhewen Li, Zijing Liu and xiaozhoupy, whose DSpark work this comes from, as co-authors in the commit. AI assistance (Claude) was used to port these changes onto main, and to write the port's tests and this description. The submitter reviewed the change and ran the tests above.

---

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

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[zixi-qi]

Current PR overlaps with #63554 and probably needs to be rebased

[TheEpicDolphin]

chatted with @zyongye offline, I'll takeover batch-sharded DSpark sampling in #63886, and this PR will keep the remaining spec decode perf optimizations

[WoosukKwon]

😠

[mergify]

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

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