Return top-p/top-k sampling mask/nucleas by hari-harmonic · Pull Request #27408 · sgl-project/sglang

hari-harmonic · 2026-06-06T00:33:02Z

Implement keep-sampling-mask / top-p & top-k replay runtime support through request parsing, sampling, scheduling, detokenization, disaggregation metadata, and the model gateway.

Motivation

Modifications

Accuracy Tests

test/registered/sampling/test_sampling_mask.py was added to check the following:

basic validations: each sampled token has an attached sampling mask, size of each sampling mask <= specified top-k, and the sampled token always exists in the sampling mask.
verifies under deterministic inference that sampled tokens when return_sampling_mask=True are identical to those returned when return_sampling_mask=False, the default path.
validates that the sampling mask and renormalized logprobs are mathematically correct. This is done by querying the server with return_logprob = True and top_logprobs_num = TOP_K with temperature 1, which returns the unnormalized logprobs of the TOP_K most probable tokens. From here, we can reconstruct the sampling mask (modulo some possible boundary off-by-ones) and compute the renormalized logprobs offline.

Speed Tests and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ⏳ Run #29291142253
Latest PR Test (Extra): ❌ Run #29291141839

gemini-code-assist · 2026-06-06T00:33:05Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

ByronHsu

Thanks for adding this. The Python/SRT side looks broadly safe to me, and the feature itself seems useful.

A few requests before this moves forward:

Can we rename the PR/title to avoid making this sound DeepSeek V3.2-specific? From the implementation, return_sampling_mask is generic sampling metadata. It is not inherently tied to one model.
Please add an end-to-end test in the usual server-launch + sample-request style. The smoke script is useful for manual validation, but I’d like a checked-in test that launches a small server/model, sends a /generate request with return_sampling_mask=true, and verifies.
Please revert the sgl-model-gateway raw-JSON/extra-body plumbing in this PR. Switching /generate from Json<GenerateRequest> to Json<Value> and threading route_generate_with_extra through the router stack feels too hacky for a known first-class field.

For Miles, let’s make this change directly in radixark/sgl-router-for-miles instead. That is the router we will use going forward, and it has the protocol structs plus router code in one repo, so we can add return_sampling_mask directly to GenerateRequest and keep the handler typed.
Please clean up the new scripts/ additions.
After this lands on main, please merge/cherry-pick it into sglang-miles.

hari-harmonic · 2026-06-12T20:01:52Z

Thanks for adding this. The Python/SRT side looks broadly safe to me, and the feature itself seems useful.

A few requests before this moves forward:

Can we rename the PR/title to avoid making this sound DeepSeek V3.2-specific? From the implementation, return_sampling_mask is generic sampling metadata. It is not inherently tied to one model.

Please add an end-to-end test in the usual server-launch + sample-request style. The smoke script is useful for manual validation, but I’d like a checked-in test that launches a small server/model, sends a /generate request with return_sampling_mask=true, and verifies.

Please revert the sgl-model-gateway raw-JSON/extra-body plumbing in this PR. Switching /generate from Json<GenerateRequest> to Json<Value> and threading route_generate_with_extra through the router stack feels too hacky for a known first-class field.
For Miles, let’s make this change directly in radixark/sgl-router-for-miles instead. That is the router we will use going forward, and it has the protocol structs plus router code in one repo, so we can add return_sampling_mask directly to GenerateRequest and keep the handler typed.

Please clean up the new scripts/ additions.

After this lands on main, please merge/cherry-pick it into sglang-miles.

Thank you for your review! Your suggestions make sense to me, this PR was more so intended to be a proof-of-concept prior to further cleanup. I will work on doing the suggestions over the coming days.

hari-harmonic · 2026-06-17T19:19:37Z

Hi @ByronHsu, I believe I have addressed the main points of your feedback and have opened a much cleaner routing PR here radixark/sgl-router-for-miles#7 to use serde. I added an e2e server test script as well in place of the old smoke test script. Please let me know if you have additional feedback that I can address -- thanks!

ByronHsu · 2026-06-23T04:05:50Z

lgtm. please turn the pr to "ready for review" and make the ci pass. current lint is failing too

gemini-code-assist · 2026-06-23T17:02:01Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

lawrence-harmonic · 2026-06-23T17:14:39Z

/tag-and-rerun-ci

Read req.return_sampling_mask directly in session_controller instead of getattr, give the streaming accumulator field a default, and set the flag in the session token-share unit test mock.

hnyls2002 · 2026-07-13T21:48:34Z

CI States

Latest PR Test (Base): ✅ Run #29158553156
Latest PR Test (Extra): ❌ Run #29158552915

# Conflicts: # python/sglang/srt/disaggregation/prefill.py

Main's sgl-project#27408 added an unconditional req.return_sampling_mask read in MetadataBuffers.set_buf; the SimpleNamespace req stub predates it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… sort The top-p nucleus capture ran a full-vocab `probs.sort` (~150k) AND a per-row `.item()` GPU->CPU sync on every decode step (both in `get_top_p_token_ids_from_probs` and `renorm_logprob_over_top_p`, which the sampler calls per token). On the Qwen3.6 codeforces S2 top-p run this slowed decode enough to push trajectories into the 1800s trajectory-timeout -> tokens_per_gpu_per_sec down, aborted-trajectory length shorter, empty_trajectory_ratio ~15% (baseline ~0%). See notes/training_logs/ioi/2026-07-16_qwen3_6_stage2_topp_mask_replay_observations.md. Port the structure of upstream sglang PR sgl-project#27408 (Return top-p/top-k sampling mask/nucleus) against our v0.5.10 tree: - `_top_p_keep_mask_bounded`: gather only the top-K (`torch.topk`, MASK_TOPK_DEFAULT=128) instead of a full-vocab sort, since the kept nucleus is a tiny prefix (measured width p99 ~11 / max ~35). K is DECOUPLED from the sampling top_k on purpose: no recipe / `--rollout-top-k` change, no distribution change -- it only bounds capture width. Silent-truncation guard: any row whose nucleus reaches the K-th column is recomputed with the full sort and spliced in (warns), so correctness never depends on K. - `get_top_p_token_ids_from_probs`: also drop the per-row `.item()` sync -- vectorized `nonzero` + a single batched `.cpu()`, split back into per-row lists on host. - `renorm_logprob_over_top_p`: same sort removed (same per-decode-step hot path); build the vocab-order keep-mask by setting only kept positions (nonzero gather) since the bounded path pads unused columns with index 0. No slime-side / launcher / wire-format change (still top_p_token_ids + offsets). Tests: test/srt/test_top_p_replay_logprob.py +2 (topk-vs-full-sort equivalence, wide-nucleus fallback); all 9 pass. New extractor verified byte-exact (order+values) vs the old one across seeds. CPU-validated only; on-GPU perf recovery pending.

github-actions Bot added the model-gateway label Jun 6, 2026

Qiaolin-Yu assigned guapisolo and ByronHsu Jun 6, 2026

This was referenced Jun 10, 2026

[CI Monitor] Daily Report - 2026-06-10 bingxche/sglang-ci-bot#99

Open

[CI Monitor] Daily Report - 2026-06-12 bingxche/sglang-ci-bot#102

Open

ByronHsu reviewed Jun 12, 2026

View reviewed changes

hari-harmonic force-pushed the hari/keep-sampling-mask branch from 09b21c1 to 1102107 Compare June 17, 2026 00:06

hari-harmonic changed the title ~~Add Deepseek V3.2's Keep Sampling mask response support~~ Return top-p/top-k sampling nucleas Jun 17, 2026

hari-harmonic changed the title ~~Return top-p/top-k sampling nucleas~~ Return top-p/top-k sampling mask/nucleas Jun 17, 2026

ByronHsu approved these changes Jun 23, 2026

View reviewed changes

hari-harmonic marked this pull request as ready for review June 23, 2026 17:01

hari-harmonic requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, ShangmingCai, Ying1123, ch-wan, hnyls2002, ispobock, merrymercy, sogalin and xiezhq-hermann as code owners June 23, 2026 17:01

hari-harmonic force-pushed the hari/keep-sampling-mask branch from 627102b to 34ff1b3 Compare June 23, 2026 17:12

hari-harmonic force-pushed the hari/keep-sampling-mask branch from b4fc262 to e175d71 Compare July 10, 2026 03:29

hari-harmonic added 11 commits July 11, 2026 15:46

feat: return keep sampling mask metadata

fb36744

test: add sampling mask server e2e test

d817e89

test: extend sampling mask server coverage

57aff36

fix: tolerate missing sampling mask flag in sessions

a9bf72c

fix tests

8a6704a

test: use real model for sampling mask

5e4312e

PR comments

456921a

fix: set return_sampling_mask flag in req mocks and session path

d37209e

Read req.return_sampling_mask directly in session_controller instead of getattr, give the streaming accumulator field a default, and set the flag in the session token-share unit test mock.

fix unconditional read

8f2af02

refactor: consolidate sampling mask branch

c89917a

fix sampling mask ipc defaults

cf117ee

hari-harmonic force-pushed the hari/keep-sampling-mask branch from e175d71 to cf117ee Compare July 11, 2026 15:47

hnyls2002 approved these changes Jul 13, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into hari/keep-sampling-mask

effb42f

# Conflicts: # python/sglang/srt/disaggregation/prefill.py

hnyls2002 merged commit a909077 into sgl-project:main Jul 13, 2026
73 of 92 checks passed

hari-harmonic mentioned this pull request Jul 22, 2026

Cherry-pick sampling mask support to sglang-miles #32108

Merged

merrymercy mentioned this pull request Jul 25, 2026

Fix flaky test_sampling_mask: mask length can legitimately be top_k + 1 #32410

Merged

1 task done

This was referenced Aug 4, 2026

[RL] Expose top-p-only sampling masks #33593

Merged

[Sampling] Add explicit bounded sampling-mask mode #34037

Closed

jvmncs mentioned this pull request Aug 10, 2026

[RL, Spec] Introduce top-p mask capture for spec and add DFlash/DSpark impl #34201

Merged

5 tasks

This was referenced Aug 17, 2026

[Sampling] Restore finite top-k requirement for sampling masks #35205

Merged

Support sampling masks without finite top-k #35765

Closed

[DO NOT MERGE] [Sampling Replay] Bounded sampling-support replay radixark/miles#2596

Closed

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

Return top-p/top-k sampling mask/nucleas (sgl-project#27408)

8b28803

Navigation Menu

Repository navigation

Return top-p/top-k sampling mask/nucleas - #27408

Return top-p/top-k sampling mask/nucleas#27408
hnyls2002 merged 12 commits into
sgl-project:mainfrom
hari-harmonic:hari/keep-sampling-mask

hari-harmonic commented Jun 6, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jun 6, 2026

Uh oh!

ByronHsu left a comment

Uh oh!

hari-harmonic commented Jun 12, 2026 •

edited

Loading

Uh oh!

hari-harmonic commented Jun 17, 2026

Uh oh!

ByronHsu commented Jun 23, 2026

Uh oh!

gemini-code-assist Bot commented Jun 23, 2026

Uh oh!

lawrence-harmonic commented Jun 23, 2026

Uh oh!

hnyls2002 commented Jul 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

hari-harmonic commented Jun 6, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

gemini-code-assist Bot commented Jun 6, 2026

Uh oh!

ByronHsu left a comment

Choose a reason for hiding this comment

Uh oh!

hari-harmonic commented Jun 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hari-harmonic commented Jun 17, 2026

Uh oh!

ByronHsu commented Jun 23, 2026

Uh oh!

gemini-code-assist Bot commented Jun 23, 2026

Uh oh!

lawrence-harmonic commented Jun 23, 2026

Uh oh!

hnyls2002 commented Jul 13, 2026

CI States

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

hari-harmonic commented Jun 6, 2026 •

edited by github-actions Bot

Loading

hari-harmonic commented Jun 12, 2026 •

edited

Loading