Skip to content

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

Draft
Jiminator wants to merge 1 commit into
mainfrom
codex/port-pr-24082-specv2-disagg-grammar
Draft

Jiminator wants to merge 1 commit into
mainfrom
codex/port-pr-24082-specv2-disagg-grammar

Conversation

@Jiminator

@Jiminator Jiminator commented May 27, 2026

Copy link
Copy Markdown
Collaborator

Ports @ashishdatta's #24082 onto current main (post scheduler-component refactor) and adds focused regression tests. Original fix and approach by @ashishdatta (credited as co-author); this PR re-applies it on the refactored layout since #24082 can't rebase cleanly (scheduler_output_processor_mixin.py was removed). Supersedes #24082.

Changes (both gated on is_spec_v2)

  1. disaggregation/decode.py — in event_loop_overlap_disagg_decode, process the previously-queued result (advancing the grammar) before launching the next Spec V2 grammar decode batch, and skip the duplicate last-batch pop. Without it, the next batch's draft proposes against a stale grammar.
  2. scheduler_components/batch_result_processor.py — in process_batch_result_decode, for Spec V2 + grammar, accept proposed tokens one at a time and stop at grammar completion, trimming output_ids / grammar FSM / reasoning / logprobs to the accepted prefix instead of emitting past the finish.

The v0.5.11 fix (#21723) covered only the Spec V1 path (eagle_info.py); Spec V2 and the disagg overlap loop were left unaddressed. This PR touches neither eagle_info.py nor the V1 path.

Tests

  • test/registered/unit/managers/test_batch_result_processor_spec_grammar.py — CPU unit test driving the real process_batch_result_decode trimming.
  • test/registered/disaggregation/test_disaggregation_spec_grammar.py — PD disagg + overlap + Spec V2 (topk=1) + xgrammar regression.

The PD test sets SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 (test-setup only): the EAGLE3 draft derives a 2048 ctx vs the target's 131072, which the Spec V2 draft worker otherwise rejects at startup.

Validation (H200)

Scenario Fix applied Fix reverted
CPU unit pass fail
PD single-node (2 GPU, mooncake/RDMA) 3/3 2 fail
PD cross-node (prefill+decode on separate hosts) 4/4 bug reproduced (corrupt JSON + abort)

Per-change isolation — both independently necessary, each guarded by a different test:

  • Revert decode.py sync only → PD test fails (Tokens not accepted → abort); trimming alone can't help (bad token is proposed before the result processor runs).
  • Revert trimming only → PD test passes (topk=1 greedy rarely accepts across the grammar boundary) but the unit test fails. So the trimming is guarded by the unit test, the sync by the PD test.

CI States

Latest PR Test (Base): ❌ Run #26544656185
Latest PR Test (Extra): ❌ Run #26544656085
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

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

@Jiminator
Jiminator force-pushed the codex/port-pr-24082-specv2-disagg-grammar branch from 01affa6 to df11184 Compare May 27, 2026 22:31
@Jiminator
Jiminator marked this pull request as draft May 27, 2026 23:04
…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
Jiminator force-pushed the codex/port-pr-24082-specv2-disagg-grammar branch from df11184 to fad591b Compare May 27, 2026 23:24
@Jiminator Jiminator closed this Sep 14, 2026
@Jiminator
Jiminator deleted the codex/port-pr-24082-specv2-disagg-grammar branch September 14, 2026 04:45
@alexnails
alexnails restored the codex/port-pr-24082-specv2-disagg-grammar branch September 14, 2026 05:41
@hnyls2002 hnyls2002 reopened this Sep 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants