Repository navigation
[cherrypick from #41235] [PD] Keep the sampling mask of a replayed rebootstrap token - #41271
Merged
yueming-yuan merged 3 commits intoSep 26, 2026
Conversation
…ect#41235) After a decode-side retract, a PD rebootstrap replays the request's already-emitted boundary token. The decode worker skipped re-committing that token's logprob but still appended the prefill worker's fresh sampling-mask row. Each such request then returned one more row than output tokens. Skip the sampling mask the same way, so the token keeps its original row.
nanjiangwill
marked this pull request as ready for review
September 25, 2026 19:36
nanjiangwill
requested review from
BBuf,
ByronHsu,
Duyi-Wang,
Edwardf0t1,
Fridge003,
HaiShaw,
ShangmingCai,
Ying1123,
ch-wan,
hnyls2002,
ispobock,
merrymercy,
sogalin and
xiezhq-hermann
as code owners
September 25, 2026 19:36
yueming-yuan
merged commit Sep 26, 2026
571212b
into
sgl-project:sglang-miles
78 of 87 checks passed
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
2 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.
Summary
Port #41235 to the
sglang-milesbranch.This is stacked on #41270 (the #40986 port), because the fix applies to the decode-side block that #40986 introduces. Only the last commit is this change.
After a decode-side retract, a PD rebootstrap replays the request's already-emitted boundary token, for example after
pause_generation(mode="retract")around weight updates. The decode worker skipped re-committing that token's logprob, but it still appended the prefill worker's fresh sampling-mask row. Each such request then returned one more row than output tokens.The fix skips the sampling mask on a replayed boundary, as the logprob path already does, so the token keeps its original row.
Port notes
Testing
test/registered/unit/{managers,disaggregation,sampling}plus the OpenAI protocol and serving-chat tests. 1435 passed, against 1434 on [cherrypick from #40986] [Sampling] Stream sampling masks as per-request arrays #41270 and 1431 onsglang-milesin the same run.test_rebootstrap_replay_keeps_one_sampling_mask_row_per_token.sglang-mileswithout either port.test/registered/sampling/test_sampling_mask.pyon 2×H100: 35 passed with 28 subtests.CI States
Latest PR Test (Base): ❌ Run #36246207751
Latest PR Test (Extra): ❌ Run #36246207579
Latest PR Test (AMD ROCm 10): ❌ Run #36246207739