Fix spec decoding with grammar in disagg

python/sglang/srt/disaggregation/decode.py

@@ -1636,6 +1636,22 @@ def event_loop_overlap_disagg_decode(self: Scheduler):
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
 
+            # Spec V2 grammar decode depends on the grammar state produced by the
+            # previous batch. With overlap scheduling that result is still sitting
+            # in result_queue, so process it now (before run_batch) to advance the
+            # grammar; otherwise the new batch would propose tokens against a stale
+            # grammar and could emit output past grammar completion.
+            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)
+
             # Launch the current batch
             if batch:
                 batch_result = self.run_batch(batch)
@@ -1645,8 +1661,9 @@ def event_loop_overlap_disagg_decode(self: Scheduler):
 
             # Process the last batch
             if self.last_batch:
-                tmp_batch, tmp_result = self.result_queue.popleft()
-                self.process_batch_result(tmp_batch, tmp_result)
+                if not need_grammar_sync:
+                    tmp_batch, tmp_result = self.result_queue.popleft()
+                    self.process_batch_result(tmp_batch, tmp_result)
             elif batch is None:
                 self.on_idle()
 

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -653,18 +653,29 @@ def process_batch_result_decode(
             # Non-spec and V2: full post-processing
             next_token_id = next_token_ids[i]
             new_accepted_len = 1
+            # Spec V2 + grammar advances output_ids, the grammar FSM, reasoning,
+            # and finish state token-by-token inside the helper (stopping at
+            # grammar completion), so the shared reasoning/finish/grammar calls
+            # below are skipped for that path.
+            is_spec_v2_grammar = batch.is_spec_v2 and req.grammar is not None
             if batch.spec_algorithm.is_none():
                 req.output_ids.append(next_token_id)
+            elif is_spec_v2_grammar:
+                next_token_id = self._accept_spec_v2_grammar_tokens(req, next_token_id)
+                next_token_ids[i] = next_token_id
+                new_accepted_len = len(next_token_id)
             else:
                 req.output_ids.extend(next_token_id)
                 new_accepted_len = len(next_token_id)
 
-            self._maybe_update_reasoning_tokens(req, next_token_id)
+            if not is_spec_v2_grammar:
+                self._maybe_update_reasoning_tokens(req, next_token_id)
 
             # Update Mamba last track seqlen
             self._mamba_prefix_cache_update(req, batch, result, i)
             req.time_stats.set_last_decode_finish_time()
-            req.update_finish_state(new_accepted_len)
+            if not is_spec_v2_grammar:
+                req.update_finish_state(new_accepted_len)
 
             self._handle_finished_req(req, i, logits_output)
 
@@ -684,9 +695,13 @@ def process_batch_result_decode(
                 )
 
             if req.grammar is not None:
-                self._apply_decode_grammar(
-                    req=req, next_token_id=next_token_id, batch=batch
-                )
+                if is_spec_v2_grammar:
+                    # Already advanced token-by-token above; just sync terminal flag.
+                    req.grammar.finished = req.finished()
+                else:
+                    self._apply_decode_grammar(
+                        req=req, next_token_id=next_token_id, batch=batch
+                    )
 
         self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
         self.token_to_kv_pool_allocator.free_group_end()
@@ -777,6 +792,48 @@ def _apply_decode_logprobs(
                     logits_output.next_token_token_ids_logprobs_idx[flat_idx]
                 )
 
+    def _accept_spec_v2_grammar_tokens(
+        self, req: Req, proposed: List[int]
+    ) -> List[int]:
+        """Accept Spec V2 proposed tokens one at a time, stopping at grammar
+        completion, and return the retained prefix.
+
+        The verify phase proposes several tokens at once, but the grammar may
+        terminate partway through that list. Tokens proposed after termination
+        are dropped so output_ids, the grammar FSM, reasoning state, and (via the
+        returned prefix) logprob bookkeeping don't advance past completion; the
+        over-allocated KV/draft state for the request is released on finish rather
+        than reused. Mutates req (output_ids, reasoning, grammar, finish state).
+        """
+        accept_tokens = []
+        try:
+            for token_id in proposed:
+                req.output_ids.append(token_id)
+                accept_tokens.append(token_id)
+                self._maybe_update_reasoning_tokens(req, token_id)
+                req.grammar.accept_token(token_id)
+                req.update_finish_state()
+                if req.finished():
+                    break
+        except ValueError as e:
+            # accept_token raises ValueError if the token is not in the grammar
+            # (misconfigured grammar or invalid token); abort the request.
+            logger.error(
+                f"Grammar accept_token failed for req {req.rid} with token {proposed}: {e}"
+            )
+            self.abort_request(AbortReq(rid=req.rid))
+
+        # _resolve_spec_overlap_tokens already added (len(proposed) - 1) to
+        # req.kv_committed_len based on the full spec-accepted count. If grammar
+        # terminated mid-list, the slots for the dropped tokens are committed but
+        # unused; roll them back so the KV pool free-list stays balanced when the
+        # request is released (otherwise SchedulerInvariantChecker._report_leak
+        # catches the imbalance — see PR comment / CI failure).
+        dropped = len(proposed) - len(accept_tokens)
+        if dropped > 0:
+            req.kv_committed_len -= dropped
+        return accept_tokens
+
     def _apply_decode_grammar(
         self,
         *,

test/registered/disaggregation/test_disaggregation_spec_grammar.py

@@ -0,0 +1,162 @@
+"""Regression test for PR #24082.
+
+Covers the specific trigger that the v0.5.11 Spec V1 grammar-finish fix does not
+reach: PD disaggregation + overlap scheduling + EAGLE Spec V2 (topk=1) + a grammar
+constraint.
+
+Spec V2 proposes several tokens per decode step. With a grammar constraint, the
+request can reach grammar completion partway through an accepted list, so:
+
+  * the decode result processor must accept the proposed tokens one at a time and
+    stop at grammar completion (no tokens emitted past the closing of the grammar),
+    trimming the over-proposed tokens and aligning logprob bookkeeping; and
+  * the disaggregated overlap decode loop must process the previous batch result
+    (advancing the grammar) before launching the next Spec V2 grammar decode batch.
+
+Notes for whoever runs this on GPU hardware:
+  * Spec V2 is gated behind ``SGLANG_ENABLE_SPEC_V2`` and only supports
+    ``--speculative-eagle-topk 1``. The env override below is inherited by the
+    launched prefill/decode subprocesses.
+  * Overlap scheduling is on by default, so the decode side runs
+    ``event_loop_overlap_disagg_decode`` (the loop modified by this PR).
+"""
+
+import json
+import unittest
+
+import requests
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import (
+    DEFAULT_DRAFT_MODEL_EAGLE3,
+    DEFAULT_TARGET_MODEL_EAGLE3,
+)
+
+register_cuda_ci(est_time=420, stage="base-b", runner_config="2-gpu-large")
+
+
+class TestDisaggregationSpecV2Grammar(PDDisaggregationServerBase):
+    # Minimal delta from the known-good PD spec config (TestDisaggregationMooncakeSpec):
+    # switch topk 4 -> 1 and enable Spec V2 so the Spec-V2 grammar path is exercised.
+    model = DEFAULT_TARGET_MODEL_EAGLE3
+    draft_model = DEFAULT_DRAFT_MODEL_EAGLE3
+    spec_algorithm = "EAGLE"
+    spec_steps = 3
+    spec_topk = 1  # Spec V2 only supports topk=1
+    spec_draft_tokens = 4
+    grammar_backend = "xgrammar"
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        spec_args = [
+            "--speculative-algorithm",
+            cls.spec_algorithm,
+            "--speculative-draft-model-path",
+            cls.draft_model,
+            "--speculative-num-steps",
+            str(cls.spec_steps),
+            "--speculative-eagle-topk",
+            str(cls.spec_topk),
+            "--speculative-num-draft-tokens",
+            str(cls.spec_draft_tokens),
+            "--grammar-backend",
+            cls.grammar_backend,
+            "--cuda-graph-max-bs",
+            "8",
+            "--dtype=float16",
+            # Cap context to the EAGLE3 draft's native length so the target and
+            # draft ModelConfigs agree (the draft derives 2048 vs the Llama-3.1
+            # target's 131072, which the Spec V2 draft worker otherwise rejects).
+            # 2048 is far above this test's output length.
+            "--context-length",
+            "2048",
+        ]
+        cls.extra_prefill_args = spec_args
+        cls.extra_decode_args = spec_args
+        with envs.SGLANG_ENABLE_SPEC_V2.override(True):
+            cls.launch_all()
+
+    @staticmethod
+    def _json_schema() -> str:
+        return json.dumps(
+            {
+                "type": "object",
+                "properties": {
+                    "name": {"type": "string", "pattern": "^[\\w]+$"},
+                    "population": {"type": "integer"},
+                    "country": {"type": "string", "pattern": "^[\\w ]+$"},
+                    "capital": {"type": "string", "pattern": "^[\\w ]+$"},
+                },
+                "required": ["name", "population", "country", "capital"],
+            }
+        )
+
+    def _generate(self, return_logprob: bool):
+        # max_new_tokens is generous so completion is driven by grammar termination,
+        # not the length cap, and the output spans multiple decode iterations.
+        response = requests.post(
+            f"{self.lb_url}/generate",
+            json={
+                "text": "Here is the information of the capital of France in the JSON format.\n",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 256,
+                    "json_schema": self._json_schema(),
+                },
+                "return_logprob": return_logprob,
+                "logprob_start_len": 0,
+            },
+        )
+        self.assertEqual(response.status_code, 200, response.text)
+        return response.json()
+
+    def test_structured_output_no_trailing_tokens(self):
+        """Output is valid JSON with nothing emitted past grammar completion."""
+        out = self._generate(return_logprob=False)
+        text = out["text"]
+
+        # json.loads rejects trailing non-whitespace content, so a clean parse of
+        # the raw text means no stray tokens leaked after the grammar terminated.
+        parsed = json.loads(text)
+        for key in ("name", "population", "country", "capital"):
+            self.assertIn(key, parsed)
+
+        # Belt and suspenders: the decoded text should end exactly at the JSON
+        # object close, not be followed by extra generated content.
+        self.assertTrue(
+            text.strip().endswith("}"), f"unexpected trailing tokens: {text!r}"
+        )
+
+    def test_spec_v2_actually_ran(self):
+        """The accepted-length stat confirms Spec V2 verification took place."""
+        out = self._generate(return_logprob=False)
+        spec_verify_ct = out["meta_info"]["spec_verify_ct"]
+        self.assertGreater(
+            spec_verify_ct,
+            0,
+            f"expected Spec V2 to run (spec_verify_ct > 0), got {spec_verify_ct}",
+        )
+
+    def test_logprob_count_matches_completion_tokens(self):
+        """Trimmed Spec V2 tokens must keep logprob count == completion token count."""
+        out = self._generate(return_logprob=True)
+        meta = out["meta_info"]
+        completion_tokens = meta["completion_tokens"]
+        output_logprobs = meta["output_token_logprobs"]
+        self.assertEqual(
+            len(output_logprobs),
+            completion_tokens,
+            "output logprobs must align with retained (trimmed) tokens: "
+            f"got {len(output_logprobs)} logprobs vs {completion_tokens} completion tokens",
+        )
+        # And the constrained output is still valid structured JSON.
+        json.loads(out["text"])
+
+
+if __name__ == "__main__":
+    unittest.main()