[Sampling] Stream sampling masks as per-request arrays

[nanjiangwill]

Motivation

return_sampling_mask returns, for every output token, the support the sampler drew from and its behavior logprob(s). That is L token IDs per token, plus L logprobs in support mode. L is the realized support size: tens under a typical top_p, up to top_k otherwise.

On main the support crosses the whole output path as Python objects:

• Scheduler: every step does a torch slice plus .tolist() per request row, and each Req keeps its full support history until it finishes.
• IPC: the scheduler pickles nested lists; the detokenizer unpickles and re-pickles them; the tokenizer manager unpickles them again.

At 256 concurrent requests with ~1000-token supports in support mode, that costs main's scheduler ~30 ms of row bookkeeping per decode step. It also stalls ~0.47 s on every 50-token emission, or 3.5 s at top_k=8192. A decode step is ~14 ms, so overlap scheduling cannot hide this, and under DP attention the slowest rank gates every rank.

Contract

• Response: byte-identical to main (see Correctness). output_token_sampling_mask, output_token_sampling_logprobs, output_token_sampling_mask_length and streaming slices come from the same float32 → Python float conversions.
• Sampling: the sampler, including mask capture and packing, is untouched.
• No opt-in, no cost: when no request asks for masks, no work, allocation or IPC payload is added.

Design

Treat the support as bulk numeric data from the sampler to the tokenizer manager, and create Python objects only where the response format requires them, once.

1. Scheduler (commit 1).
   • Each committed token copies its row, taken from numpy views of the step's host buffers, into the request's SamplingMaskRows.
   • The queue buffers double in capacity, so each byte is copied O(1) times regardless of the allocator.
   • On emission take() hands the rows off as a SamplingMaskChunk of int32/float32 arrays without copying. The handed-off buffer becomes a spare that is reused only once no view of it remains (checked by refcount), so steady-state appends write to warm pages.
2. IPC (commit 1). output_token_sampling_mask becomes Optional[List[Optional[SamplingMaskChunk]]]: one chunk per request. This follows the existing per-request buffer fields (decode_ids: List[array], input_top_logprobs_*_flat: List[np.ndarray]), so the detokenizer passes it through and multi-tokenizer routing uses the generic _extract_field_by_index unchanged.
3. Tokenizer manager (commit 1). SamplingMaskChunk.to_lists() builds the response lists.
4. GC-neutral chunk pickling (commits 3 and 4).
   • numpy's own array reducer, and wrapping a bytearray with np.frombuffer (whose base becomes a GC-tracked memoryview), each leave GC-tracked objects per array. When every message carries one short chunk per request, as with per-token streaming, those objects doubled the tokenizer manager's young-generation collections. That promoted its accumulated lists into 3.5× more full collections than on main.
   • A chunk therefore pickles its three buffers as read-only bytes and rebuilds them with np.frombuffer. The bytes are unchanged, unpickling leaves no tracked objects, and protocols below 5 keep the default reduction.

Checks removed, and why they are safe

Each removed check guarded a state that main already rules out before the check runs. Links point to main.

1. Length re-validation in materialize_sampling_mask_output (if status == OK and not (0 <= length <= packed_width)).

   • The sampler emits lengths = realized_lengths.clamp(max=packed_size).
   • Its packed IDs have exactly packed_size columns: top-k, PyTorch producer, and greedy, width 1 with length 1.
   • So 0 <= length <= packed_width always holds.

2. None rows in add_sampling_mask_return_values (None if mask is None else mask[i]). Every caller appends only after get_sampling_mask_finish_reason returned no abort, which means the status was OK:

   • decode;
   • non-PD prefill;
   • PD prefill.

   A missing capture has status None and is aborted as well. An OK status always materializes its row.

3. PD prefill set_buf: the -1 sentinel, the empty-queue and None-row branches, and the capacity RuntimeError (utils.py).

   • The first row is queued: it is appended right before the last-chunk send, or before the send after a pending bootstrap. set_buf runs only on the last chunk. Prefill streams a request only on abort (L438, L983, L1198), after which nothing is sent, or after its KV transfer.
   • The row is not None: same status guarantee as 2.
   • The row fits the buffer: the prefill buffer's width is --sampling-mask-max-tokens, the same cap the sampler uses. A longer support gets OVERFLOW status and is aborted by 2.

4. PD decode if sampling_mask_len < 0 (decode.py). It only matched the -1 sentinel, which prefill can no longer write; per 3, main never reached it either. The length the decode side reads is the one the prefill wrote for the same request, which is at least 1.

Nothing enforces that prefill and decode use the same --sampling-mask-max-tokens; the flag's help text asks for it. That holds on main too, and none of the removed checks guarded it.

The per-request queue relies on one more property: a chunk is serialized before its request's next append. sock_send encodes synchronously, and the stream accumulator is local to stream_output. The refcount check keeps buffer recycling safe even if a chunk were retained.

Alternatives measured and rejected

• One CSR payload per batch concatenates every request's rows before pickling. Pickle already writes each request's buffer into its stream, so the concat was a pure extra copy. It measured 4× slower to encode than per-request arrays at 210 MB.
• array.array queues (the Req.output_ids primitive) over-allocate by 1/16, so appending 32 KB rows reallocates nearly every step. Once glibc's dynamic mmap threshold has risen, that realloc copies the whole buffer. It was 48% of live scheduler time at top_k=8192. A microbenchmark that appended into fresh buffers missed it.
• FlashInfer radix top_k for packing matches torch.topk's (id, logprob) pairs, lengths and statuses on identical captures. However, with bf16 logits it orders tied token IDs differently in nearly every row, so responses would no longer be byte-identical to main. It saved ~0.3 ms of GPU time per step at batch 256.
• An R3-style preallocated pool indexed by KV slot is how routed experts are gathered. It fits fixed-width rows, but sampling-mask rows are up to --sampling-mask-max-tokens wide: for example, 500k slots × 4096 × 8 B = 16 GB.

Internal signature changes

Where	Before	After
BatchTokenIDOutput / BatchStrOutput	two nested-list fields	output_token_sampling_mask: Optional[List[Optional[SamplingMaskChunk]]] (output_token_sampling_logprobs removed)
Req	output_token_sampling_mask, output_token_sampling_logprobs, send_output_sampling_mask_offset	sampling_mask_rows: Optional[SamplingMaskRows]
LogitsProcessorOutput.next_token_sampling_mask_idx / _logprobs	Python lists	numpy row views

Performance

All end-to-end numbers are same-host A/B: one H100 machine runs the server from main and then from this branch.

• Setup: Qwen3-8B, temperature=1, ignore_eos, default overlap scheduler and pickle IPC; non-streaming unless noted.
• Metric: steady-state decode throughput (tok/s), from the scheduler's decode log over windows where every request is running.

TP1, 512 new tokens

batch	top_k / top_p (mean support)	mask off	selected: main → PR	support: main → PR
256	1000 / 1.0 (~1011)	18960 → 18953	11859 → 16477 (+39%)	6645 → 16032 (+141%)
64	1000 / 1.0 (~1012)	7785 → 7783	6337 → 7280 (+15%)	5298 → 7156 (+35%)
256	1024 / 0.95 (~7)	19013 → 19024	16428 → 16408 (0%)	16036 → 16350 (+2%)
256	20 / 1.0 (20)	19043 → 18992	17117 → 17061 (0%)	16488 → 17042 (+3%)
1	any	156 → 156	150 → 150	150 → 149

DP attention (--tp 2 --dp 2 --enable-dp-attention, 256 requests per rank, per-rank throughput)

top_k / top_p	mask off	selected: main → PR	support: main → PR
1000 / 1.0	16848 → 16804	9320 → 14978 (+61%)	5414 → 14242 (+163%)
1024 / 0.95	16697 → 16684	14721 → 14671 (0%)	14299 → 14629 (+2%)

In support mode the sampling return now costs 2.7 ms per decode step against a 15.2 ms step.

top_k=8192, batch 256 (256 new tokens, --sampling-mask-max-tokens 9216; mask off is 20636 → 20600)

mode	decode tok/s: main → PR	wall s: main → PR	wall with --tokenizer-worker-num 8: main → PR	scheduler RSS GB: main → PR
selected	1832 → 11872 (6.5×)	66.5 → 41.8	60.1 → 33.4	30.0 → 8.8
support	836 → 8744 (10.5×)	159.4 → 98.1	140.1 → 79.7	52.1 → 9.5

Even here, support mode adds 16.9 ms per decode step against a 12.4 ms step. Most of the remainder is copying each 50-token emission into the pickle stream and then into a zmq message (see Limitations).

RL rollouts with staggered finishes. Batch 256, ignore_eos. Each request's max_new_tokens is spread evenly over the listed range, so requests finish at different steps. The metric is wall time for the whole batch. Mask off is the same requests without return_sampling_mask, which has no mode.

top_k	output tokens	mask off	selected: main → PR	support: main → PR
128	64–1024	11.3 s	12.6 → 12.2 s	14.0 → 12.2 s
256	64–1024	11.3 s	13.4 → 12.2 s	15.7 → 12.3 s
8192	64–1024	11.5 s	97.2 → 73.0 s	not run
8192	32–512	5.3 s	not run	131.5 → 83.1 s

At top_k 128–256, the sampling return adds about 1 s to an 11.3 s rollout. At top_k=8192, the tokenizer manager's list building bounds wall time (see Limitations). The support run uses half the output tokens to bound its response size. Scheduler RSS drops from 24.9 to 8.8 GB (selected) and from 23.0 to 9.5 GB (support).

Streaming (stream_interval=1, --incremental-streaming-output, batch 256). Four runs, two with each server launched first on the machine. Steady decode tok/s, median of the four runs:

top_k	mask off	selected: main → PR	support: main → PR
1000	20910 → 21010	8330 → 18720	4520 → 18440
20	20810 → 20800	18860 → 18740	17900 → 18520

Wall time per run, main → PR:

workload	run 1	run 2	run 3 (PR first)	run 4 (PR first)
mask off, top_k=20	4.2 → 3.8 s	3.4 → 3.4 s	3.7 → 5.3 s	3.4 → 3.4 s
selected, top_k=20	5.1 → 4.7 s	3.9 → 3.8 s	4.4 → 6.8 s	3.9 → 4.0 s
support, top_k=20	5.1 → 5.2 s	4.1 → 4.2 s	4.5 → 7.0 s	4.3 → 4.5 s
selected, top_k=1000	12.1 → 10.2 s	10.0 → 8.7 s	10.9 → 10.3 s	10.9 → 9.5 s
support, top_k=1000	20.2 → 20.4 s	18.2 → 14.9 s	18.3 → 17.5 s	20.2 → 16.9 s

Run 3's slowdown includes mask off, which this PR does not touch. Two runs of main on one machine differ by up to 0.4 s on the top_k=20 rows. At top_k=20 in support mode, the PR is 0.1–0.2 s above main in the other three runs. Its tokenizer manager rebuilds each request's one-row chunk from three buffers and converts it with to_lists (1.8 µs for a 20-entry support row), which costs more than unpickling main's two short lists.

Per stage, batch 256, support mode. This is a microbenchmark of both branches' code on one H100 host. An "emission" is one 50-token chunk.

stage	top_k=1000: main → PR	top_k=8192: main → PR
scheduler row bookkeeping per step	29.6 → 0.82 ms	210 → 2.5 ms
scheduler per emission (pickle + send)	466 → 57 ms	3.47 → 0.59 s
detokenizer per emission (unpickle + re-pickle)	1.48 s → 92 ms	13.7 → 0.81 s
tokenizer manager per emission (unpickle + build lists)	1.12 → 1.32 s	10.4 → 10.3 s
GPU capture and packing per step (unchanged)	2.00 → 1.99 ms	2.20 → 2.19 ms

Profiler. A same-host torch profiler capture of 40 steady decode steps (batch 256, top_k=1000):

mode	GPU busy: main → PR	time per step: main → PR	longest process_batch_result: main → PR
mask off	95% → 95%	13.7 → 13.6 ms	22 → 22 ms
selected	37% → 95%	35.1 → 14.7 ms	250 → 10 ms
support	20% → 87%	60.6 → 17.9 ms	578 → 70 ms

Correctness

The optimization must not change what clients receive in either mode. Three tests in this PR pin that contract:

• CPU, test_clients_receive_the_sampler_rows_unchanged: drives sampler-shaped outputs through materialization, the per-request queues, stream emissions at arbitrary steps and both IPC codecs. Every request's lists must equal the direct conversion of the sampler's rows: token_ids[row, :length] with selected_logprobs[row] or support_logprobs[row, :length].
• CPU, test_sampling_mask_row_reaches_decode_unchanged (commit 5): under PD disaggregation, the prefill worker writes the first token's row into the metadata buffers, and the decode worker's handoff reads it back into its queue. The rows the decode worker streams must equal the row the prefill worker queued, in both modes.
• GPU, test_modes_and_streaming_return_identical_sampling_masks (deterministic inference): for the same seeded tokens, selected and support mode must return the same supports, and each selected logprob must equal the sampled token's support logprob. A per-token stream must match the response emitted every 50 tokens.

The oracle and GPU tests (commit 2) fail against a deliberately broken transport (queued lengths not cleared on hand-off), while the pre-existing suite still passes on it. The PD handoff test fails if:

• the decode side reads the full logprob width in selected mode;
• the decode side drops a support ID;
• the prefill side writes only the first logprob.

Separately, a fixed workload of 54 requests is sent to main twice (a reproducibility control), then to this branch, on the same machine, and the responses are diffed exactly: floats bitwise, and each number's JSON text. The fields compared are text, output_ids, finish_reason, completion_tokens, the three sampling-mask fields and output_token_logprobs.

The workload covers:

• selected, support and mask-off requests in one batch;
• top_k from 1 to 1000, top_p and temperature variants;
• natural and ignore_eos stops, and 120 tokens so every request crosses the 50-token forced emission;
• non-streaming and streaming requests;
• return_logprob together with masks.

All runs use --enable-deterministic-inference (PyTorch sampler, per-request seeds). Two rounds on this PR's head:

configuration	main vs main (both rounds)	main vs PR, round 1	main vs PR, round 2
requests sent one at a time	54/54 identical	54/54 identical	54/54 identical
concurrent requests	54/54 identical	50/54 identical¹	54/54 identical
concurrent, msgpack IPC (SGLANG_USE_PICKLE_IPC=false)	54/54 identical	54/54 identical	54/54 identical
concurrent, --tokenizer-worker-num 2	54/54 identical	53/54 identical¹	54/54 identical
concurrent, --incremental-streaming-output	54/54 identical	54/54 identical	54/54 identical

¹ Concurrent first-token differences. These come from main, not from this PR. The evidence covers all 29 concurrent runs of this harness, including earlier revisions of this PR and #41205, which builds on it.

• Frequency: main vs PR differed in 8 of the 29 runs. Two runs of main differed from each other in 2 of them.
• Location: every difference, in both kinds of comparison, is in row 0 (the token sampled at prefill), and only in its logprobs. Across 713 differing values, they are at most 4 float32 ulps apart, at most 4.5e-7 relative. Support IDs, text and output_ids are identical in every run.
• Which requests: almost all are the top_k=1000 requests with the largest supports (requests 27–31). The one exception is request 49 (top_k=200), which also differed between two runs of main in the same run. The other main-vs-main difference hit requests 27 and 28.
• Requests sent one at a time: 54/54 identical in every run.

The mechanism is prefill batching:

• Which concurrent requests share a prefill batch depends on arrival timing, and the first token's captured logprobs depend slightly on that batch, even under deterministic inference on main.
• Two runs of main run at the same speed and usually form the same prefill batches. This PR's scheduler runs at a different speed, so its batches differ from main's more often.
• The PR does not change the sampler or the capture. Its transport copies float32 bits without arithmetic, and the CPU test above checks that bit for bit.

Each configuration covers 4320 mask rows. The default FlashInfer sampler is not reproducible run to run: main itself differs from main in float32 last bits (max 1.9e-6) of a few rows, with identical support sets. Under it, this branch gives identical text/output_ids and differs only within that same noise.

Limitations

• The response format itself costs Python objects. The tokenizer manager builds one int per support entry, plus one float in support mode. At top_k=8192 × 256 this is ~10 s of tokenizer CPU per 50-token cycle, and it bounds wall time. In the staggered top_k=8192 rollouts above, it accounts for most of the gap to mask-off: 73.0 s against 11.5 s in selected mode.
  • It runs off the scheduler and scales with --tokenizer-worker-num.
  • About a quarter of it is cyclic GC walking the freshly built lists. The existing --gc-threshold 100000 cut tokenizer GC time from 2.4 to 0.5 s and wall time by 9% on a 256 × 256-token top_k=1000 support run.
  • [Sampling] Write /generate sampling masks from numpy rows #41205, stacked on this PR, removes it for /generate. Rows stay numpy arrays, and orjson writes them as the same JSON: staggered top_k=8192 rollouts drop from 64.9 to 16.9 s (selected) and 73.6 to 25.7 s (support).
• Emissions still copy once into the pickle stream. Non-streaming requests emit every 50 tokens in the same step, and pickling copies each chunk into freshly allocated memory. At top_k=8192 × 256 in support mode, that is about 5 ms per decode step amortized. Pickle protocol 5's out-of-band buffers, sent as extra zmq frames, would remove this copy on every hop. It is an IPC-wide change (received numpy arrays become read-only views of zmq frames), so it belongs in its own PR.
  • pyzmq then copies each encoded message once more into a zmq buffer, about 3.5 ms per step amortized at the same size. [IPC] Send encoded messages without copying #41206, independent of this PR, sends without that copy. On top of this PR at top_k=8192 × 256, it raises decode throughput by 21% in selected mode and 39% in support mode.
• D2H is sized by capacity, not by the realized support. The packed [B, --sampling-mask-max-tokens] buffer crosses every step; it overlaps on the copy stream.
• Capture keeps its two full-vocabulary renorm kernels, at 2.0 ms per step against 0.36 ms for sampling alone at batch 256. A fused capture kernel would be a follow-up.

Testing

• CPU unit tests: test/registered/unit/{managers,disaggregation,sampling} plus the OpenAI protocol and serving-chat tests.

  • 1620 passed with 972 subtests; main passes 1617.
  • Both trees share the same pre-existing failure, test_mm_process_config (cross-test state; it passes alone), and two collection errors from the CPU image (parameterized missing, no accelerator).

• New streamer case: queued rows cross both the pickle and msgpack codecs once, in batch order, and expand to the response lists. It also checks that rows queued later never overwrite a chunk that is still referenced; the case fails with the refcount check removed.

• GPU: test/registered/sampling/test_sampling_mask.py on 2×H100: 35 passed with 28 subtests, covering FlashInfer and PyTorch sampling, deterministic inference, TP=2 and PP=2.

• Benchmark validation: every benchmark response checks three things:

  • the mask count equals the output length;
  • the sampled token is in its support;
  • support-mode rows are aligned.

  This covers the --tokenizer-worker-num 8 run, which exercises per-request routing.

Builds on the sampling-mask work in #36630, #36631, and #40932.

---

CI States

Latest PR Test (Base): ⏳ Run #36162293505
Latest PR Test (Extra): ❌ Run #36162293185
Latest PR Test (AMD ROCm 10): ⏳ Run #36162293592

[ByronHsu]

Clean!

[ByronHsu]

/tag-and-rerun-ci

[nanjiangwill]

/rerun-failed-ci

[nanjiangwill]

the failed tests are not related to this pr.

[ByronHsu]

/tag-and-rerun-ci