[TheEpicDolphin]

Context

During batch-sharded sampling, the target logits are sharded along the token dimension. The requests are deterministically assigned to ranks using the formula: req_state_idx % tp_size. So during TP=4, rank 0 only samples (and in the case of spec decoding verifies) requests with state indices: 0, 4, 8, 12, ....

Currently, draft model-derived speculators allocate persistent draft logits tensors to hold all requests up to max-num-seqs, which is wasteful during batch-sharded sampling because only max-num-seqs // TP of those request logits are actually used by each rank. Sharding the draft logits can enable us to save GPU memory.

This PR

Under batch-sharded sampling, each rank now allocates draft logits only for the requests it verifies. The draft logits cache shrinks from max_num_seqs rows to cdiv(max_num_seqs, tp_size) rows per rank, and the freed memory goes to the KV cache.

Layout. A request's draft logits live only on its owner rank, at row req_state_idx // tp_size. Two Triton helpers in batch_shard.py define the pair: shard_owner (req_state_idx % num_shards, the existing owner rule) and shard_row (req_state_idx // num_shards). Every write and read of the cache goes through them. With one shard both reduce to the current behavior and compile away.

Writes. Drafting stays replicated, so every rank still samples every draft token and no new collectives are added. Each rank writes only the rows it owns:

• gumbel_sample takes logits_cache_num_shards / logits_cache_shard_idx and masks its cache store to the owner rank. DraftModelSpeculator.sample_draft and DSpark's _sample_logits pass them through.
• DFlash2's _cache_draft_logits_kernel, which writes only the top-k candidate columns, applies the same mask.

Reads. RejectionSampler takes draft_logits_num_shards, and the rejection kernels index draft logits with shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS).

Allocation. The cache moves from DraftModelSpeculator.__init__ to load_model, so it is allocated once at its final size after the target and draft models are loaded. It is sharded when all of these hold:

• batch-sharded sampling is active for the target (batch_sharded_sampling_enabled, a new helper that the model runner also uses to create the BatchSharder),
• the speculator opts in via supports_batch_sharded_draft_logits(),
• no draft watermarker is configured.

The model runner passes speculator.draft_logits_num_shards to the RejectionSampler. It also asserts that a custom rejection sampler never receives a sharded cache.

Speculators that opt in: every one whose cache writes honor the owner mask.

• TargetDependentARSpeculator: Eagle, MTP, Gemma4.
• StandaloneARSpeculator and MultiModuleMTPSpeculator.
• DFlashSpeculator, and through it DSpark, XPress, DFlash2 and LiLiCorr.
• Exception: DSpark keeps the full cache when the draft model provides its own fused sample_sequential (the AMD DeepSeek-V4.1 path), because those kernels write the full-size cache.

Memory freed per rank (TP4, max_num_seqs=256, probabilistic drafts), from server startup logs:

Target + draft	Draft logits cache, base → shard	KV cache memory, base → shard	KV cache tokens
DeepSeek-V4.1-Flash + DSpark (K=5)	0.31 → 0.08 GiB	74.44 → 74.67 GiB	+0.31%
Qwen3.6-27B + MTP (K=3)	0.36 → 0.09 GiB	150.69 → 150.96 GiB	+0.18%
Qwen3.6-27B + DFlash (K=7)	0.83 → 0.21 GiB	145.77 → 146.39 GiB	+0.43%
Qwen3.8-27B + DFlash2 (K=7, fp32 cache)	1.66 → 0.41 GiB	147.07 → 148.31 GiB	+0.85%

Benchmarks

No regression in accuracy or acceptance rate for any speculator that now writes batch-sharded draft logits.

Setup: 4x GB200, TP4, draft_sample_method=probabilistic, prefix caching off, natural EOS. Each pair ran 4 server boots in the order base, shard, shard, base. "base" is --enable-batch-sharded-sampling off and "shard" is on. Startup logs confirmed every shard boot sharded the draft logits cache on all 4 ranks.

• GSM8K: tests/evals/gsm8k/gsm8k_eval.py, 1319 questions. t0 is temperature 0, one run per boot. t1 is temperature 1 with 3 seeds per boot, 6 runs per arm. The Qwen models use chat completions with thinking disabled.
• SPEED-Bench: low_entropy subset, throughput_2k config, concurrency 64, 256 prompts, temperature 1, top_p 1.
• Acceptance length = 1 + accepted / drafts, from the server's Prometheus metrics.
• Each cell shows the mean over the arm's runs.

GSM8K accuracy

Target + draft	K	t0 base	t0 shard	t1 base	t1 shard
DeepSeek-V4.1-Flash + DSpark	5	0.907	0.906	0.816	0.807
Qwen3.6-27B + MTP	3	0.969	0.970	0.968	0.967
Qwen3.6-27B + DFlash	7	0.970	0.970	0.966	0.966
Qwen3.8-27B + DFlash2	7	0.964	0.961	0.960	0.960

Acceptance length

Target + draft	GSM8K t0 base	GSM8K t0 shard	GSM8K t1 base	GSM8K t1 shard	SPEED-Bench base	SPEED-Bench shard
DeepSeek-V4.1-Flash + DSpark	3.554	3.540	3.275	3.283	3.546	3.556
Qwen3.6-27B + MTP	3.617	3.619	3.519	3.515	3.258	3.271
Qwen3.6-27B + DFlash	5.499	5.490	4.980	4.984	3.905	3.927
Qwen3.8-27B + DFlash2	6.164	6.152	5.876	5.890	3.636	3.647

SPEED-Bench noise check (Qwen3.6-27B + MTP)

Every pair read 0.3–0.6% higher acceptance on SPEED-Bench with sharding. To check whether that was real, I ran a third arm that shards verification but keeps the draft logits cache replicated. I ran each arm with top-k on and off, with 2 boots per arm.

Arm	top_k=20 (model default)	top_k disabled
base	3.266	3.268
sharded verification, replicated draft logits	3.263	3.262
sharded verification, sharded draft logits	3.263	3.261

The three arms match. Base alone ranged from 3.254 to 3.279 across boots, so the shift is boot-to-boot noise.

Notes

• DFlash2 GSM8K t0, 0.964 → 0.961: a 4-question difference. At temperature 0, verification compares against the target argmax and never reads cached draft logits, so this can't come from this change.
• DSV4.1 GSM8K t1, 0.816 → 0.807: not significant (Welch p≈0.3). This model's t1 accuracy varies a lot from seed to seed (0.78–0.83).

[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 #97407 for commit 62e2b087e791.

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

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97436 for commit 610c19a89066.

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