[TheEpicDolphin]

Context

Under batch-sharded sampling, #64356 shards verification and the draft logits cache by batch position: request i is owned by rank i % tp_size. The DSpark sequential (Markov) block still runs replicated, though. Every rank computes full-vocab draft logits and samples every request at every draft step, even though each rank only keeps the draft logits of the requests it owns.

This PR

This PR shards the DSpark sequential block along the same ownership rule, so each rank runs the Markov loop only for its own requests.

DSparkSpeculator._sample_sequential_sharded. Each rank owns at most capacity = cdiv(num_reqs, tp_size) requests, so every shape depends only on num_reqs and the step is captured in the existing draft CUDA graphs. No second graph set is needed. Each step:

1. Each rank computes its vocab shard of the base draft logits (compute_draft_logits_local) for every request, in owner-major order. Padding slots read request 0's rows and are discarded.
2. An all-to-all re-shards the logits from the vocab dimension to the batch dimension, so each rank holds full-vocab logits for its own requests.
3. The Markov loop (_sample_markov_chain, now shared with the replicated path) samples those rows and writes the owner's draft logits rows. Padding lanes pass -1 and write nothing.
4. An all-gather of the drafts (plus the confidences, under adaptive verification) restores batch order on every rank.

Single NCCL communicator.

• PyNcclCommunicator.all_to_all_single (new) does a variable-split all-to-all, and is graph-capturable when the split sizes are fixed.
• Verification's logits exchange now goes through the same TP PyNCCL communicator, via all_to_all_vocab_shards (which replaces all_to_all_logits) and get_tp_pynccl_comm().
• A separate c10d all-to-all would make NCCL allocate a second set of point-to-point buffers, about 1.1 GiB per GPU on GB200.
• The c10d path remains as a fallback when PyNCCL is unavailable.

When the loop stays replicated. It logs the reason and runs replicated when:

• the draft logits cache is replicated (PCP, a draft watermarker, or a draft model whose fused sample_sequential writes the full cache);
• the online acceptance estimator is used;
• top-k Markov sampling is configured;
• the draft model lacks a vocab-sharded compute_draft_logits_local;
• PyNCCL is unavailable.

Models. compute_draft_logits_local is added for DeepSeek-V4 and V4.1 DSpark (NVIDIA) and Qwen3 DSpark.

Benchmarks

Setup:

• DeepSeek-V4.1-Flash + DSpark, K=5, probabilistic drafting.
• 4× GB200, TP4, max_num_seqs=256, FP8 KV cache, prefix caching off, natural EOS.
• Arms: main (55b677b, loop replicated) vs this PR (loop sharded). Both use --enable-batch-sharded-sampling, so the comparison isolates the sharded loop.
• 4 boots, ordered PR, main, main, PR.
• SPEED-Bench low_entropy, throughput_2k, temperature 1. GSM8K: 1319 questions, temperatures 0 and 1.
• Latencies are medians.

Speed (this PR vs main)

	c1	c8	c16	c32	c64	c128	c256
ITL, AV off	+1.0%	−0.1%	−0.4%	−0.6%	−0.7%	−1.5%	−2.6%
tok/s, AV off	+1.6%	−0.4%	+1.5%	+2.4%	−0.4%	+4.3%	+5.4%
ITL, AV on	+1.9%	+0.7%	−0.6%	−0.4%	−0.9%	−2.2%	−1.4%
tok/s, AV on	+0.4%	+1.3%	+2.7%	+2.4%	+3.2%	+2.3%	+1.7%

With AV off, ITL at c256 goes from 48.43 to 47.16 ms; the two boots of each arm agree within about 0.05 ms. Throughput and the AV-on runs are noisier. At c1 the extra collectives cost about 1%. A minimum batch size for the sharded loop (no extra CUDA graphs needed) could remove that cost as a follow-up.

Accuracy and acceptance (main → this PR)

	GSM8K t0	GSM8K t1	Acceptance length c64	Acceptance length c256
AV off	0.898 → 0.901	0.800 → 0.797	3.54 → 3.54	3.54 → 3.54
AV on	0.898 → 0.910	0.809 → 0.810	3.28 → 3.27	3.35 → 3.36

Tests

.venv/bin/python -m pytest tests/v1/worker/test_gpu_batch_shard.py tests/v1/spec_decode/test_lilicorr.py  # 59 passed
.venv/bin/python -m pytest tests/distributed/test_pynccl.py  # 17 passed, incl. the new test_pynccl_all_to_all_single
.venv/bin/python -m pytest tests/v1/spec_decode/test_dspark_topk.py tests/models/test_dsv41_mono_decode.py \
  tests/v1/worker/test_gpu_model_runner_v2_eplb.py  # 92 passed
.venv/bin/python -m pytest tests/v1/e2e/spec_decode/test_sharded_sampling.py  # 2 passed
.venv/bin/python -m pytest tests/v1/e2e/general/test_sharded_sampling.py  # 12 passed

test_pynccl_all_to_all_single covers uneven and zero-row splits, plus CUDA graph replay. test_dspark_sharded_sampling_equivalence (new) runs Qwen3-4B-FP8 + DSpark at TP2. It checks that the draft logits cache and the loop are sharded only when the flag is on, and that greedy and seeded outputs match replicated sampling up to near-tie noise. The general e2e test covers verification's all-to-all, which now runs on PyNCCL.

[mergify]

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

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

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98523 for commit f558c8597c23.

ℹ️ This PR is 26 commits behind upstream main. Merged upstream main into this branch as f558c8597c23 before starting CI. Pull the branch before pushing more changes.