[Perf][Model Runner V2] Compact sampling masks on GPU instead of unpacking the full-vocab bitmask on CPU - #54901
Merged
ywang96 merged 10 commits intoSep 4, 2026
Conversation
aoshen02
requested review from
WoosukKwon,
njhill and
yewentao256
as code owners
September 2, 2026 05:52
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
aoshen02
requested review from
ProExpertProg,
houseroad,
mgoin,
robertgshaw2-redhat,
tlrmchlsmth and
youkaichao
as code owners
September 2, 2026 06:53
Collaborator
Author
|
/ci run |
|
✅ Triggered Buildkite CI #86837 for commit |
…cking the full-vocab bitmask on CPU SamplingMaskTensors.tolists() ran np.unpackbits over a [num_reqs, vocab] bitmask on every step, an O(num_reqs * vocab) CPU pass in the worker's async output thread that grew to ~130 ms/step at 512 requests and capped throughput regardless of the GPU step time. Compact the support on the GPU into [num_reqs, max_top_k] int32 ids with a per-row prefix sum, D2H only that, and keep the bit-packed mask solely as the exact fallback for rows whose support exceeds the batch top_k bound (top-k ties). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
… array per request per step - Clamp the compact row width to MAX_COMPACT_SUPPORT (2048) so a request with top_k near the vocab size cannot make the batch allocate [num_reqs, vocab] int32; wider rows go through the exact bitmask. - Fail configuration when return_sampling_mask is combined with enable_batch_sharded_sampling: gather_sampler_output() does not forward SamplingMaskTensors, so masks would silently come back as None. - SamplingMaskLists.offsets is optional: a per-step request slice carries only its token_ids, halving the per-request ndarray encode/decode between engine core and API server; merge/to_nested_list handle both forms and to_nested_list converts with a single tolist(). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
Build each row's support with one helper that reads the compact row or, for overflow rows, the bitmask; the vectorized fast path and the overflow loop did the same work in two branches. Same cost (0.46 ms at 512 rows). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
Model Runner V2 emits one position per request per step and speculative decoding is rejected with sampling replay, so slice_request asserts num_positions == 1 and merge stacks single-position chunks; the multi-position branches were dead. Drop the impossible max(1, ...) clamp and the concatenate-based cu_num_generated_tokens. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
The frontend concatenated the per-step slices into CSR only to split them back into list[list[int]]; convert each slice with one tolist() instead and drop SamplingMaskLists.merge. Fix the kernel docstring (the compact row holds at most max_num_kept ids) and cover the no-sampled-row step in the tolists test. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
The kernel already zeroes the count of rows that sampled nothing, so tolists() can emit one CSR row per request and slice_request() indexes it by req_idx directly; cu_num_generated_tokens is no longer needed. Inline the kernel's single-use temporaries. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
Collaborator
Author
|
/ci run |
aoshen02
force-pushed
the
perf/sampling-mask-gpu-compact
branch
from
September 2, 2026 11:53
f6f7ac9 to
9c10bd5
Compare
|
✅ Triggered Buildkite CI #86861 for commit |
Replace the hand-built tensor tests with one reference test: for rows of every support size (empty, one, around the compact width, past the cap, the whole vocab) the mask must equal the finite-logit set, and unsampled rows must be empty. Add a scheduler test that three requests each get their own sampler row, and make the e2e test check that the returned logprob mass lives exactly on the support and that it fits top_k. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
Collaborator
Author
|
/ci run |
aoshen02
requested review from
ApostaC,
alexm-redhat,
heheda12345 and
ywang96
as code owners
September 2, 2026 13:15
aoshen02
requested review from
AndreasKaratzas,
DarkLight1337,
NickLucche,
aarnphm,
ivanium and
orozery
as code owners
September 2, 2026 13:15
|
✅ Triggered Buildkite CI #86866 for commit |
Collaborator
Author
|
/ci run |
|
✅ Triggered Buildkite CI #86972 for commit |
Signed-off-by: vx120 <893600387@qq.com>
Member
|
Core changes LGTM |
Collaborator
Author
Any update you think we should do? |
njhill
approved these changes
Sep 3, 2026
Member
|
/ci run |
|
✅ Triggered Buildkite CI #87164 for commit |
ItsRoy69
pushed a commit
to ItsRoy69/vllm
that referenced
this pull request
Sep 10, 2026
…cking the full-vocab bitmask on CPU (vllm-project#54901) Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: vx120 <893600387@qq.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: vx120 <893600387@qq.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>
yinli-systems
pushed a commit
to yinli-systems/single-gpu-inference-lab
that referenced
this pull request
Sep 17, 2026
…LM #54901 Upstream vllm-project/vllm#54901 (merged 2026-09-04, shipped in 0.29.1) diagnosed the same host-side np.unpackbits bottleneck and landed the same top_k-bounded compact layout before this work was done against v0.29.0. README, artifact READMEs, notes, status ledger and installer now state that; the patch is for 0.29.0 only. The residual-mask wording no longer asserts a cause for the five differing masks; it reports the bitmap-vs-bitmap control and states that the experiment does not attribute them to the compact layout.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
--return-sampling-mask(#49577) made decode throughput collapse under load. PrimeIntellect reported ~2x RL step time and higher engine imbalance after integrating it (PrimeIntellect-ai/prime-rl#3431).Root cause:
SamplingMaskTensors.tolists()rannp.unpackbits+np.nonzeroover a[num_reqs, vocab]bitmask on every step. That is an O(num_reqs × vocab) CPU pass in the worker's async output thread, independent of how many tokens the nucleus actually keeps. With a 151k vocab it costs ~17 ms at 64 requests and ~130 ms at 512, so once it exceeds the GPU step time the engine is CPU-bound and throughput stops scaling with batch size. Because the cost grows with in-flight requests, the busiest engine in a multi-engine deployment gets slower per step and holds its KV longer, which is the imbalance prime-rl observed.Fix: compact the support on the GPU. The Triton kernel now writes each row's finite-logit ids in ascending order into a
[num_reqs, max_top_k]int32 buffer via a per-block prefix sum (max_top_kis the largesttop_kin the batch; sampling-mask requests already requiretop_k > 0). Only that buffer plus per-row counts are copied to the host, andtolists()becomes a boolean gather overnum_reqs × max_top_kelements. The bit-packed mask is still produced and copied, but it is only read on the host for rows whose support exceedsmax_top_k(the top-k kernel keeps ties at the boundary), so the returned support is unchanged.Test Plan
Unit tests (
tests/v1/test_outputs.py -k sampling_mask) updated for the new tensor layout, plus new cases for the tie-overflow fallback and a multi-block vocab.Kernel-level equivalence against the old implementation (one GB200, vocab 151552, fp32 logits with a random support of up to
kids per row,torch.cuda.synchronize()aroundfrom_logits+to_cpu_nonblockingfor the GPU column and wall-clock aroundtolistsfor the CPU column). Note that the bit-packed mask is still produced and copied to the host in the new code; only the host-side unpack is gone:toliststolistsServing A/B, Qwen2.5-1.5B-Instruct on one GB200, Model Runner V2,
--logprobs-mode processed_logprobs --max-num-seqs 1024, all clients sendingtemperature=1, top_p=0.95, top_k=512, logprobs=1, 128-token random prompts (a worst case: the nucleus keeps ~220 ids per position), 256 output tokens,/v1/completions. "no mask" and "mask after" were measured side by side in the same run; "mask before" is from an earlier run on the same machine:GLM-4.5-Air TP=4 (the model from the prime-rl report),
vllm/vllm-openai:nightly-aarch64container, same client settings:4x GB200,
--max-num-seqs 1024 --max-model-len 32768, 512 output tokens, natural-language prompt (the nucleus keeps ~6-8 ids per position on average):/v1/completions/v1/completions/v1/completions/inference/v1/generateBefore the fix the mask path is flat at ~3.1k tok/s regardless of concurrency, i.e. CPU-bound; after it, throughput is within 2-3% of no-mask.
Test Result
See above. Follow-up commits (review feedback): the frontend now builds the response
list[list[int]]straight from the per-step slices instead of re-assembling CSR; the compact row width is clamped toMAX_COMPACT_SUPPORT = 2048so a request withtop_knear the vocab size cannot make the batch allocate[num_reqs, vocab]int32 (wider rows still round-trip exactly through the bitmask);return_sampling_maskcombined with--enable-batch-sharded-samplingis now rejected at startup becausegather_sampler_output()does not forward the masks (#53826 can lift this); and each per-step request slice carries a singletoken_idsarray (offsets=None), halving the per-request ndarray encode/decode between engine core and API server.SamplingMaskListsstays a 3-field NamedTuple, so the rust client's opaque 3-element wire shape is unchanged.The remaining gap on the small model is frontend work proportional to the number of finished requests (chunk merge and the
list[list[int]]response), which is a few percent at the step time of a 100B-class model.Duplicate-work check
Searched open PRs for
sampling mask,return_sampling_mask,SamplingMaskTensorson 2026-09-02. #53826 (batch-sharded sampling drops masks) and #54166 (MTP support) touch the same file but neither addresses the CPU unpack cost; this change keeps theSamplingMaskTensors/SamplingMaskListsinterfaces used by both.AI assistance (Claude Code) was used for this change; the submitter reviewed every line and ran the tests above.
🤖 Generated with Claude Code