[Juntian777]

Overview

Runs DSpark's sequential Markov sampling for DeepSeek-V4.1 on NVIDIA as the existing bias GEMM plus two Triton kernels per draft step, through the draft model's sample_sequential hook (as the ROCm model does). Sampled draft tokens are drawn by inverse CDF of the cached draft logits instead of Gumbel-max. Production code changes only under vllm/models/deepseek_v41/nvidia/ (+223). It is opt-in: --additional-config '{"dspark_fused_markov": true}' enables it, and the speculator's loop stays the default.

Claims

• DSpark Markov sampling (5 draft steps, GB200, CUDA graph): 1.6x at 1 request per rank, 2.7x to 3.9x at 16 to 253 requests per rank.
• DEP4 decode serving with AgentX KV: step time -1.7% at 64 requests x 131K context per rank, -3.2% at 128 x 32K.
• GSM8K and DSpark acceptance length unchanged at DEP4 and TEP4 (TP4 + EP).

Validation

Kernel time, DSparkSpeculator._sample_sequential (main's loop) vs this PR, CUPTI under a CUDA graph, one GB200, V = 129280, rank 256, 5 draft steps, 3 runs each (spread under 1% except 1 request):

Requests per rank	main loop (us), T=1	This PR (us), T=1	Speedup	main loop (us), T=0/0.6/1 mixed	This PR (us), mixed	Speedup
1	111.9	71.6	1.56x	107.7	64.5	1.67x
16	237.6	89.3	2.66x	233.6	88.7	2.63x
32	269.9	96.2	2.80x	256.7	95.9	2.68x
64	411.4	144.3	2.85x	385.2	144.7	2.66x
128	658.3	194.9	3.38x	601.2	196.5	3.06x
253	1158.5	298.9	3.88x	1039.9	298.4	3.48x

Re-checked on main 51c6dfd8c8 with this revision, same benchmark, one GPU, median of 3 alternating rounds (worst spread 9.0% at 1 request, 3.8% or less from 16 to 253 requests at T=1, 7.1% or less for mixed):

Requests per rank	main loop (us), T=1	This PR (us), T=1	Speedup	main loop (us), mixed	This PR (us), mixed	Speedup
1	112.9	71.0	1.59x	107.1	65.7	1.63x
16	238.0	89.3	2.66x	231.3	87.7	2.64x
64	411.6	144.4	2.85x	384.5	144.2	2.67x
128	658.7	196.2	3.36x	600.7	197.3	3.05x
253	1158.7	301.3	3.85x	1042.5	299.9	3.48x

DEP4 decode serving, DeepSeek-V4.1-Flash on one GB200 node, DSpark K=5 probabilistic drafting, synthetic acceptance length 3.51, FULL_DECODE_ONLY graphs, MegaAttn + MegaMoE, --max-num-seqs 128. The client fills the KV with AgentX token streams, then runs 6 rounds of 4096 forced output tokens with exactly the given requests on every rank (no logits processing, as AgentX); step time is the median of the last 3 rounds, measured only over windows of pure decode:

Requests x context per rank	main (ms/step)	This PR (ms/step)	Change
64 x 131K	25.26	24.84	-1.7%
128 x 32K	34.17	33.08	-3.2%

Accuracy (vllm serve, DSpark K=5, probabilistic drafting, standard rejection sampling; GSM8K is lm_eval 5-shot, greedy, max 1024 tokens, concurrency 32; acceptance length from the server's spec-decode metrics over the GSM8K run):

	GSM8K strict-match	Acceptance length
DEP4, reference (#62359's run, same flags)	0.9265	3.451
DEP4, this PR	0.9280	3.448
TEP4 (TP4 + EP), this PR	0.9257	3.463

These evals and the DEP4 decode serving numbers ran on an earlier base (66875d3127, 259 commits behind 51c6dfd8c8) and were not rerun. Since then this PR's kernels have not changed and the only production additions are the opt-in dspark_fused_markov check and one inlined local. All "This PR" numbers are the fused path, which now needs that flag. The speculator changed on main; I checked its sample_sequential contract by reading the code, and on the merged tree ran the test below and the kernel benchmark above.

The DEP4 eval ran with --jit-monitor-verbose: the new kernels compile only during warmup. The TEP4 eval ran with VLLM_GPU_SYNC_CHECK=error.

.venv/bin/python -m pytest tests/models/deepseek_v41/nvidia/test_fused_markov.py -v  # 2 passed on main 51c6dfd8c8

The test runs the speculator's loop and the fused path on the same inputs: temperature-0 tokens and every cached logit of temperature-0 requests (and step 0 of all requests) match bit for bit, and sampled tokens follow softmax(cached logits / temperature).

Details

• Per draft step, the loop runs the Markov embedding, the bias GEMM, the logits add, the Gumbel sample (its kernel, a cross-block argmax and a gather) and a copy into draft_tokens. Here the bias GEMM (cuBLAS) is kept, _markov_sample_kernel adds the base logits as the loop does (fp32 sum rounded once to the cache dtype), writes the draft-logits cache and keeps each vocab block's max and exp-sum, and _markov_select_kernel picks the token and writes its markov_w1 row into the next step's GEMM input.
• Inverse CDF instead of Gumbel-max: one uniform per request, keyed by seed, draft step and position like the loop's noise, picks a block by the prefix sum of the block masses and a token inside it from the same fp32 weights. Tokens follow softmax(cached logits / T), the distribution the rejection sampler verifies; it costs one exp per logit instead of a hash and two logs, which made the Gumbel kernel ALU-bound. The sampled tokens are therefore not those of the loop's Gumbel draw; temperature-0 tokens are.
• The hook is set at model init only when dspark_fused_markov is true and the loop would sample from a full draft-logits cache with plain fp32 Gumbel noise: probabilistic drafting, no adaptive verification, no fp64 Gumbel, no watermarking and no batch-sharded sampling (whose sharded cache keeps the loop). Draft top-k is handled by the speculator before the hook.
• Not a duplicate of: [Perf][DSv4.1] DSpark draft on the specialized decode #61869 (fuses the loop only with decode_backend=latency_specialized on GB300 TP), [Perf][Spec Decode] DSpark: run the top-k Markov walk in one kernel #64033 (the top-k Markov walk), [Perf][Spec Decode] DSpark: run the Markov walk in one kernel with int8 W2 on chip (Hopper) #62958 (int8 markov_w2 on Hopper), [ROCm] Fused DSpark draft sampling: markov GEMV + sample epilogue #63181 / [ROCm][Perf][DSv4.1] Run DSpark's fused Markov sampling on gfx950 with --monokernel #63891 and the merged [ROCm][Perf][DSv4.1] Mono decode on MI300X and MI325X (gfx942) #60731 (ROCm).

---

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.

[Aminous-001]

+++

[mergify]

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

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

[Juntian777]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99415 for commit b195a3bc7abd.

[TheEpicDolphin]

LGTM