[Model Runner V2] batch-sharded sample

[TheEpicDolphin]

Summary

Currently before sampling, the full target logits are materialized, allocating memory of size O(batch_size * (num_spec_tokens + 1) * vocab_size). This PR reduces that per-step memory allocation by a factor of 1/P during tensor parallelism by sharding the logits and sampler inputs along the batch dimension, and sampling in parallel on a subset of the batch.

Sharding the sampler inputs (InputBatch) was straightforward, and is handled by BatchSharder.shard_sampler_inputs. Requests are assigned to ranks by partitioning the request state index mapping (using idx_mapping % tp_size) in a deterministic way. From there, we derive the new sharded input batch CPU arrays & GPU tensors for the local rank, based on the requests that it "owns".

Sharding the logits along the batch dimension is accomplished by having the model compute logits for only its slice of the vocab size by implementing the compute_logits_local method. After that, each rank has a [num_tokens, vocab_size // P] slice of the full logits. However, instead of all-gathering to get the full [num_tokens, vocab_size], we instead all-to-all along the batch dimension, so that each rank ends up with [num_tokens // P, vocab_size].

In summary, this PR (1) reduces peak target logits memory from O(batch_size * (num_spec_tokens + 1) * vocab_size) => O(batch_size * (num_spec_tokens + 1) * vocab_size / P), and (2) parallelizes the sampling operation, which can be particularly beneficial for TPOT whenever heavy sampling operations (e.g. top-k, top-p, block verification, etc) are performed. No changes had to be made to the Sampler or RejectionSampler because they work perfectly fine with smaller batch sizes as long as we shard the InputBatch properly, which we do.

NOTE: A big motivation for this PR is to enable sharding of the draft logits. The cached draft logits from the speculator forward pass(es) are preallocated in memory for the worst case, with a size of O(max_num_reqs * num_spec_tokens * vocab_size). We can reduce that by a factor of 1/P as well, but sharded sampling is a prerequisite.

Speed-Bench 2K/2K, Temp=1.0, Top-p=0.95

DSV4 + DSpark (7 spec tokens)

Concurrency 64

Metric	Baseline	Sharded Sampling	Δ %
Request throughput (req/s)	2.62	2.66	+1.53%
Output token throughput (tok/s)	5229.81	5286.66	+1.09%
Median TPOT (ms)	12.27	12.20	-0.57%
Median TTFT (ms)	230.80	234.14	+1.45%
Median ITL (ms)	28.46	27.58	-3.09%
Acceptance length	3.62	3.64	+0.55%
Per-position acceptance (%)
Position 0	79.01	79.11	+0.13%
Position 1	59.94	60.09	+0.25%
Position 2	44.22	44.59	+0.84%
Position 3	31.80	32.28	+1.51%
Position 4	22.42	22.85	+1.92%
Position 5	15.21	15.56	+2.30%
Position 6	9.54	9.83	+3.04%

Low Concurrencies (4-16)

Batch-sharded sampling is not expected to (and indeed does not) yield improvements at low concurrencies.

Concurrency 4 (mean of 3 runs)

Metric	Baseline	Sharded Sampling	Δ %
Request throughput (req/s)	1.41	1.45	+3.32%
Output token throughput (tok/s)	719.37	743.42	+3.34%
Median TPOT (ms)	4.93	4.87	-1.35%
Median TTFT (ms)	222.53	225.18	+1.19%
Median ITL (ms)	13.82	13.83	+0.05%
Acceptance length	3.49	3.65	+4.68%
Per-position acceptance (%)
Position 0	77.40	78.61	+1.57%
Position 1	57.94	60.77	+4.90%
Position 2	42.75	46.05	+7.71%
Position 3	30.36	32.57	+7.27%
Position 4	20.15	23.04	+14.32%
Position 5	13.07	15.48	+18.47%
Position 6	7.47	8.69	+16.34%

Concurrency 8 (mean of 3 runs)

Metric	Baseline	Sharded Sampling	Δ %
Request throughput (req/s)	2.08	2.08	+0.00%
Output token throughput (tok/s)	1062.54	1063.37	+0.08%
Median TPOT (ms)	6.97	6.99	+0.29%
Median TTFT (ms)	224.52	227.87	+1.49%
Median ITL (ms)	16.59	16.52	-0.42%
Acceptance length	3.41	3.48	+2.25%
Per-position acceptance (%)
Position 0	76.54	77.26	+0.94%
Position 1	56.94	57.72	+1.37%
Position 2	41.07	42.10	+2.49%
Position 3	28.13	29.76	+5.82%
Position 4	18.88	20.21	+7.04%
Position 5	12.11	13.15	+8.56%
Position 6	7.04	7.87	+11.89%

Concurrency 16 (mean of 3 runs)

Metric	Baseline	Sharded Sampling	Δ %
Request throughput (req/s)	3.11	3.09	-0.54%
Output token throughput (tok/s)	1590.46	1583.13	-0.46%
Median TPOT (ms)	9.50	9.51	+0.14%
Median TTFT (ms)	229.14	230.37	+0.54%
Median ITL (ms)	18.32	18.16	-0.87%
Acceptance length	3.54	3.47	-1.89%
Per-position acceptance (%)
Position 0	77.66	77.03	-0.81%
Position 1	58.88	57.91	-1.64%
Position 2	43.44	41.99	-3.34%
Position 3	30.76	29.60	-3.78%
Position 4	21.07	20.15	-4.35%
Position 5	13.84	13.05	-5.69%
Position 6	8.01	7.44	-7.04%

Minimax M3 + DSpark (8 spec tokens)

Metric	Baseline	Sharded Sampling	Δ %
Request throughput (req/s)	2.60	2.74	+5.38%
Output token throughput (tok/s)	4081.74	4354.00	+6.67%
Median TPOT (ms)	15.61	14.31	-8.33%
Median TTFT (ms)	299.06	284.43	-4.89%
Median ITL (ms)	32.87	31.17	-5.17%
Acceptance length	3.91	3.96	+1.28%
Per-position acceptance (%)
Position 0	77.12	77.43	+0.40%
Position 1	57.85	58.43	+1.00%
Position 2	43.96	44.59	+1.43%
Position 3	33.96	34.71	+2.21%
Position 4	26.57	27.41	+3.16%
Position 5	21.10	21.96	+4.08%
Position 6	16.88	17.56	+4.03%
Position 7	13.43	14.00	+4.24%

[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

[TheEpicDolphin]

Moved this logic into the get_logprobs_dims helper method that is used for sharding the batch for sampling

[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 #84330 for commit 6c447e62794b.

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84340 for commit 285609a0ca58.

[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 #84430 for commit dc100e634118.

[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

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84867 for commit 5f2e762d838b.

[WoosukKwon]

thanks for the PR and sorry for the late review.

[WoosukKwon]

Can we add an extra test for bitwise correctness (without top-p and top-k)?