Fix spec decoding with grammar in disagg by ashishdatta · Pull Request #24082 · sgl-project/sglang

ashishdatta · 2026-04-29T19:04:54Z

Motivation

Fix structured output generation when speculative decoding is used in disaggregated decode mode.
Spec-v2 decode can accept multiple draft tokens in one step. With grammar-constrained decoding, the scheduler must process grammar state updates before launching the next overlapping disagg decode batch, and it must stop accepting speculative tokens once the grammar reaches a terminal state.

Modifications

Process the previous overlapped disagg decode result before launching a new spec-v2 grammar decode batch.
For spec-v2 requests with grammar constraints, accept tokens one at a time and stop once the request finishes.
Trim speculative tokens accepted after grammar completion so stale over-advanced state is not reused.

Accuracy Tests

Not a model-forward or kernel accuracy change. Adds a CPU unit test and a PD-disagg + overlap + Spec V2 (topk=1) + grammar regression; validated single-node and cross-node on H200 (see PR comment). Validated changed files with:

python3 -m py_compile \
  python/sglang/srt/disaggregation/decode.py \
  python/sglang/srt/managers/scheduler_components/batch_result_processor.py

Speed Tests and Profiling

No expected steady-state speed impact. The grammar sync path only applies to spec-v2 + grammar + disaggregated overlap decode when a prior result is pending.
No benchmarking run.

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.

CI States

Latest PR Test (Base): ✅ Run #27744167980
Latest PR Test (Extra): ✅ Run #27744167737

gemini-code-assist · 2026-04-29T19:04:58Z

Warning

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

…ort #24082) Port PR #24082 onto the post-refactor scheduler component layout. The v0.5.11 fix only covered the Spec V1 grammar-finish path; it did not cover Spec V2 or the PD-disaggregated overlap decode event loop. - decode.py: in event_loop_overlap_disagg_decode, process the previously queued batch result (advancing the grammar) before launching the next Spec V2 grammar decode batch, and skip the duplicate last-batch pop so it is not processed twice. - batch_result_processor.py: in process_batch_result_decode, for is_spec_v2 and req.grammar is not None, accept proposed tokens one at a time and stop at grammar completion, trimming output_ids, the grammar FSM, reasoning state, and logprob bookkeeping to the accepted prefix. Tests: - CPU unit test for the trimming loop (runs the real process_batch_result_decode). - PD disagg + overlap + EAGLE Spec V2 (topk=1) + grammar regression test. Co-authored-by: Ashish Datta <1856117+ashishdatta@users.noreply.github.com>

…ort sgl-project#24082) Port PR sgl-project#24082 onto the post-refactor scheduler component layout. The v0.5.11 fix only covered the Spec V1 grammar-finish path; it did not cover Spec V2 or the PD-disaggregated overlap decode event loop. - decode.py: in event_loop_overlap_disagg_decode, process the previously queued batch result (advancing the grammar) before launching the next Spec V2 grammar decode batch, and skip the duplicate last-batch pop so it is not processed twice. - batch_result_processor.py: in process_batch_result_decode, for is_spec_v2 and req.grammar is not None, accept proposed tokens one at a time and stop at grammar completion, trimming output_ids, the grammar FSM, reasoning state, and logprob bookkeeping to the accepted prefix. Tests: - CPU unit test for the trimming loop (runs the real process_batch_result_decode). - PD disagg + overlap + EAGLE Spec V2 (topk=1) + grammar regression test. Co-authored-by: jimmy.shong <jimmy.shong@radixark.ai>

…ort #24082) Port PR #24082 onto the post-refactor scheduler component layout. The v0.5.11 fix only covered the Spec V1 grammar-finish path; it did not cover Spec V2 or the PD-disaggregated overlap decode event loop. - decode.py: in event_loop_overlap_disagg_decode, process the previously queued batch result (advancing the grammar) before launching the next Spec V2 grammar decode batch, and skip the duplicate last-batch pop so it is not processed twice. - batch_result_processor.py: in process_batch_result_decode, for is_spec_v2 and req.grammar is not None, accept proposed tokens one at a time and stop at grammar completion, trimming output_ids, the grammar FSM, reasoning state, and logprob bookkeeping to the accepted prefix. Tests: - CPU unit test for the trimming loop (runs the real process_batch_result_decode). - PD disagg + overlap + EAGLE Spec V2 (topk=1) + grammar regression test. Co-authored-by: Ashish Datta <1856117+ashishdatta@users.noreply.github.com>

Jiminator · 2026-05-28T00:26:54Z

Validated on current main (8×H200, PD disagg + overlap + EAGLE Spec V2 topk=1 + xgrammar):

code	CPU unit	PD single-node	PD cross-node
this PR	✅ 2 passed	✅ 3 passed	✅
2 source changes reverted	❌ 1 failed	❌ 2 failed (`Tokens not accepted` → invalid JSON)	❌ bug repro

Reverting just the two source changes reproduces the bug, so the fix is load-bearing. Both changes are independently necessary — the disagg grammar-sync is guarded by the PD regression, the Spec V2 token-trim by the unit test.

Jiminator · 2026-05-28T00:37:24Z

/tag-and-rerun-ci

…st test ctx-len - batch_result_processor.py: extract the Spec V2 grammar token-accept loop into _accept_spec_v2_grammar_tokens(); replace the grammar_advanced flag with an is_spec_v2_grammar predicate; de-duplicate the accept_token error handling. - test_disaggregation_spec_grammar.py: use --context-length 2048 so the target and EAGLE3-draft ModelConfigs agree, instead of overriding SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN to ignore the mismatch. - decode.py: drop a restating comment on the need_grammar_sync guard. The need_grammar_sync overlap-ordering mechanism (from sgl-project#24082) is unchanged; turning it into an enforced scheduler invariant is tracked as follow-up.

_resolve_spec_overlap_tokens advances req.kv_committed_len by (accept_lens - 1) based on the full spec-accepted count, before _accept_spec_v2_grammar_tokens knows the grammar may terminate partway through. When grammar terminates mid-list, the dropped tokens are committed-but-unused and leak when the request releases — caught by SchedulerInvariantChecker._report_leak("pool") on current main (added in sgl-project#25623 after the original port was written). Decrement kv_committed_len by len(proposed) - len(accept_tokens) at the end of the helper so the resolver's over-commit is undone for the dropped suffix. No effect when grammar doesn't terminate mid-list (dropped == 0). Repro on the merge head (7e2f8df) without the fix: TestEagleConstrainedDecodingV2 hits `pool memory leak detected ... total=187911, available=187648, evictable=262` (1 slot imbalance). With the fix: 12/12 V2 cases pass on the same hardware.

ShangmingCai · 2026-05-30T08:12:13Z

+                if not need_grammar_sync:
+                    tmp_batch, tmp_result = self.result_queue.popleft()
+                    self.process_batch_result(tmp_batch, tmp_result)


Why not just set self.last_batch = None in previous if need_grammar_sync block?

ShangmingCai · 2026-05-30T08:14:24Z

+            need_grammar_sync = (
+                batch
+                and batch.is_spec_v2
+                and batch.has_grammar
+                and batch.forward_mode.is_decode()
+                and len(self.result_queue) > 0
+            )
+            if need_grammar_sync:
+                tmp_batch, tmp_result = self.result_queue.popleft()
+                self.process_batch_result(tmp_batch, tmp_result)
+


Something like this:

Suggested change

need_grammar_sync = (

batch

and batch.is_spec_v2

and batch.has_grammar

and batch.forward_mode.is_decode()

and len(self.result_queue) > 0

)

if need_grammar_sync:

tmp_batch, tmp_result = self.result_queue.popleft()

self.process_batch_result(tmp_batch, tmp_result)

need_grammar_sync = (

batch

and batch.is_spec_v2

and batch.has_grammar

and batch.forward_mode.is_decode()

and self.last_batch is not None

and len(self.result_queue) > 0

)

if need_grammar_sync:

tmp_batch, tmp_result = self.result_queue.popleft()

self.process_batch_result(tmp_batch, tmp_result)

self.last_batch = None

ShangmingCai · 2026-05-30T08:17:24Z

On second thought, without token-by-token acceptance, grammar could terminate at a token, but the bonus following token is already appended to output_ids. The trimming logic is correct in this case. But I am not fully sure about the modification of these if not is_spec_v2_grammar blocks, you should find some experts to review this part.

cc @fzyzcjy if you have time, can you take a look at this?

ShangmingCai · 2026-05-30T08:18:35Z

I think we'd better add this test to test_disaggregation_basic, instead of creating another one. This could help us save some CI resources.

Done, TestDisaggregationSpecV2Grammar lives in test_disaggregation_basic.py now. Bumped the module's est_time to 750 to cover the added class. Each class still runs different server flags.

ShangmingCai

cc: @Ubospica, do you have time to help us review this PR?

@ShangmingCai

…gg sync, merge tests - decode.py (per @ShangmingCai): tighten `need_grammar_sync` predicate with `self.last_batch is not None`, set `self.last_batch = None` inside the block, and remove the `if not need_grammar_sync` guard from the later "Process the last batch" branch. Reads cleaner and removes the only special-case fallthrough. - batch_result_processor.py (per @ShangmingCai): revert the Spec V2 grammar trimming branch + `_accept_spec_v2_grammar_tokens` helper + kv_committed_len rollback. Experiment: with the disagg sync in place, the trim never fires in the PR's primary path (Option A: 0 TRIM_PROBE hits across 3 passes). In the non-disagg Eagle V2 path it fired once across 12 tests, but plain main passes that test without the trim too (12 passed, 0 leaks, 0 grammar errors), and Option A passes equivalently (3 passed, 0 leaks). The kv leak the rollback patched was self-induced by the trim's over-commit. - test_disaggregation_basic.py (per @ShangmingCai): merge `TestDisaggregationSpecV2Grammar` here as a new class so the PD fixture amortizes across the existing Spec PD tests instead of spinning up another module. Bump `est_time` to 750s. - delete `test_disaggregation_spec_grammar.py` (merged) and `test_batch_result_processor_spec_grammar.py` (no helper to test).

@ShangmingCai

@ShangmingCai edited their earlier comment to "On second thought, without token-by-token acceptance, grammar could terminate at a token, but the bonus following token is already appended to output_ids. The trimming logic is correct in this case." Restoring the trim/helper/rollback I had removed in 7e7a487: the spec-v2 bonus-token after grammar termination needs to be dropped from output_ids, the grammar FSM, and KV bookkeeping — my earlier experiment didn't trigger that exact case so it looked dead, but the scenario is real and the trim is the correct guard. Re-adds: - batch_result_processor.py: _accept_spec_v2_grammar_tokens helper, the is_spec_v2_grammar predicate + elif branch in process_batch_result_decode, the `if not is_spec_v2_grammar` guards on reasoning/finish, the grammar terminal-flag sync (in lieu of _apply_decode_grammar), and the kv_committed_len rollback for the dropped tokens. - test_batch_result_processor_spec_grammar.py: CPU unit test exercising the helper (3 tokens proposed, grammar terminates after 2, assert trimmed prefix + finish state + grammar advance). The decode.py simplification and the test_disaggregation_basic.py merged class from 7e7a487 are kept as-is. @ShangmingCai flagged the `if not is_spec_v2_grammar` guard pattern for expert review (cc @Ubospica already pinged) — guards are intentional so per-token state isn't double-applied after the helper, happy to refactor if there's a cleaner shape.

ShangmingCai · 2026-05-31T16:20:18Z

/rerun-test test/registered/disaggregation/test_disaggregation_basic.py

github-actions · 2026-05-31T16:20:38Z

Results for /rerun-test test/registered/disaggregation/test_disaggregation_basic.py:

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_basic.py

Ubospica · 2026-06-04T09:41:08Z

Thanks for your contribution @ashishdatta and the review request from @ShangmingCai @Jiminator! This fix is definitely useful. The fix generally looks great to me, and I am further looking into the details of it.

# Conflicts: # test/registered/disaggregation/test_disaggregation_basic.py

Align spec grammar handling with the current spec_algorithm API and remove a disagg sync path made redundant by the mainline overlap guard. Record the intentional KV rollback owner and keep the unit fake aligned with the current result processor hook.

Spec V2 is now the default speculative path, and the old SGLANG_ENABLE_SPEC_V2 descriptor was removed on main. Launch the regression test directly so it stays compatible after merging main.

Ubospica

The logic generally looks great to me. Just a few comments

Ubospica · 2026-06-16T10:14:10Z

+            logger.error(
+                f"Grammar accept_token failed for req {req.rid} with token {proposed}: {e}"
+            )
+            self.abort_request(AbortReq(rid=req.rid))


Do we need to call update_finish_state() here?

Ubospica · 2026-06-16T10:15:22Z

+                accept_tokens.append(token_id)
+                self._maybe_update_reasoning_tokens(req, token_id)
+                req.grammar.accept_token(token_id)
+                req.update_finish_state()


Maybe we should accept token first, then append it into output_ids and accept_tokens, in case the token get rejected?

Ubospica · 2026-06-16T10:16:05Z

+        self.assertTrue(req.grammar.finished)
+        # Logprob bookkeeping sees only the retained tokens.
+        self.assertEqual(req.logprob.output_token_logprobs_val, [-0.1, -0.2])
+        self.assertEqual(req.logprob.output_token_logprobs_idx, [101, 102])


shall we check kv_committed_len as well

Ubospica · 2026-06-16T10:17:56Z

+"""Unit test for the Spec V2 + grammar trimming in process_batch_result_decode.
+
+Companion to the GPU regression in
+test/registered/disaggregation/test_disaggregation_spec_grammar.py (PR #24082).


Looks the file is deleted

Accept grammar tokens before mutating request output state, assert KV commit rollback in the CPU regression, and fix the stale disaggregation test reference called out in review.

JustinTong0323 · 2026-06-16T14:48:44Z

@ashishdatta pushed a small cleanup pass on the tests/comments. Main changes: moved the processor stub test to CPU CI, folded the spec_verify_ct check into the existing e2e requests, and trimmed the added comments/docstrings. Could you take a quick look?

Consume the pending FINISH_ABORT after grammar accept_token fails so the decode-result cleanup sees the request as finished in the same pass.

Ubospica

LGTM!

# Conflicts: # test/registered/disaggregation/test_disaggregation_basic.py

Co-authored-by: jimmy.shong <jimmy.shong@radixark.ai> Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com> Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Co-authored-by: Codex <codex@example.com>

ashishdatta requested review from ByronHsu, ShangmingCai, Ying1123, hnyls2002, merrymercy and xiezhq-hermann as code owners April 29, 2026 19:04

Jiminator self-assigned this May 19, 2026

Jiminator mentioned this pull request May 27, 2026

fix(spec): trim Spec V2 grammar output in PD disagg overlap decode (port #24082) #26508

Draft

Jiminator force-pushed the ashish/specdec-disagg-structured-output branch from 3f0fd7c to f0bb40c Compare May 27, 2026 23:26

Merge branch 'main' into ashish/specdec-disagg-structured-output

be0aafa

github-actions Bot added the run-ci label May 28, 2026

Jiminator and others added 4 commits May 28, 2026 00:51

Merge branch 'main' into ashish/specdec-disagg-structured-output

7e2f8df

Merge branch 'main' into ashish/specdec-disagg-structured-output

0fbe351

Jiminator added the run-ci-extra label May 30, 2026

ShangmingCai reviewed May 30, 2026

View reviewed changes

Jiminator added 2 commits May 30, 2026 08:56

Jiminator requested a review from Ubospica May 31, 2026 20:50

Merge branch 'main' into ashish/specdec-disagg-structured-output

1fbca66

Jiminator requested review from HaiShaw and sogalin as code owners June 12, 2026 07:57

Jiminator added 4 commits June 14, 2026 02:10

Merge remote-tracking branch 'origin/main' into pr-24082

5a227c1

# Conflicts: # test/registered/disaggregation/test_disaggregation_basic.py

Fix stale spec grammar CI checks

9cbd2b4

Align spec grammar handling with the current spec_algorithm API and remove a disagg sync path made redundant by the mainline overlap guard. Record the intentional KV rollback owner and keep the unit fake aligned with the current result processor hook.

Merge branch 'main' into ashish/specdec-disagg-structured-output

ca49a3d

Drop stale Spec V2 env in disagg test

0afd389

Spec V2 is now the default speculative path, and the old SGLANG_ENABLE_SPEC_V2 descriptor was removed on main. Launch the regression test directly so it stays compatible after merging main.

Ubospica reviewed Jun 16, 2026

View reviewed changes

Jiminator and others added 2 commits June 16, 2026 11:15

Address spec grammar review comments

72b5c05

Accept grammar tokens before mutating request output state, assert KV commit rollback in the CPU regression, and fix the stale disaggregation test reference called out in review.

Clean up spec grammar regression tests

6900cdf

JustinTong0323 and others added 2 commits June 16, 2026 15:51

Merge branch 'main' into ashish/specdec-disagg-structured-output

8b57b95

Finish spec grammar aborts

b5ffafa

Consume the pending FINISH_ABORT after grammar accept_token fails so the decode-result cleanup sees the request as finished in the same pass.

Ubospica approved these changes Jun 18, 2026

View reviewed changes

Codex and others added 2 commits June 18, 2026 07:20

Merge remote-tracking branch 'origin/main' into pr-24082

0c9819d

# Conflicts: # test/registered/disaggregation/test_disaggregation_basic.py

Merge branch 'main' into ashish/specdec-disagg-structured-output

56bb1dc

hnyls2002 approved these changes Jun 18, 2026

View reviewed changes

hnyls2002 merged commit 9fc9d37 into sgl-project:main Jun 18, 2026
168 of 180 checks passed

hnyls2002 mentioned this pull request Jun 18, 2026

[Spec] Unify speculative grammar token-accept path in decode processing #28682

Merged

Uh oh!

Conversation

ashishdatta commented Apr 29, 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

CI States

Uh oh!

gemini-code-assist Bot commented Apr 29, 2026

Uh oh!

Jiminator commented May 28, 2026

Uh oh!

Jiminator commented May 28, 2026

Uh oh!

ShangmingCai May 30, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ShangmingCai May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jiminator May 31, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jiminator May 30, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai left a comment

Choose a reason for hiding this comment

Uh oh!

ShangmingCai commented May 31, 2026

Uh oh!

github-actions Bot commented May 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Ubospica commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Ubospica left a comment

Choose a reason for hiding this comment

Uh oh!

Ubospica Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

Ubospica Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

Ubospica Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

Ubospica Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

JustinTong0323 commented Jun 16, 2026

Uh oh!

Ubospica left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ashishdatta commented Apr 29, 2026 •

edited by github-actions Bot

Loading

ShangmingCai May 30, 2026 •

edited

Loading

ShangmingCai May 30, 2026 •

edited

Loading

ShangmingCai May 30, 2026 •

edited

Loading

github-actions Bot commented May 31, 2026 •

edited

Loading

Ubospica commented Jun 4, 2026 •

edited

Loading