Fix spec decoding with grammar in disagg (#24082)

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>

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

@@ -648,16 +648,25 @@ class SchedulerBatchResultProcessor:
             # Non-spec and V2: full post-processing
             next_token_id = next_token_ids[i]
             new_accepted_len = 1
+            needs_tokenwise_grammar_accept = (
+                not batch.spec_algorithm.is_none() and req.grammar is not None
+            )
             if batch.spec_algorithm.is_none():
                 req.output_ids.append(next_token_id)
+            elif needs_tokenwise_grammar_accept:
+                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 needs_tokenwise_grammar_accept:
+                self._maybe_update_reasoning_tokens(req, next_token_id)
 
             req.time_stats.set_last_decode_finish_time()
-            req.update_finish_state(new_accepted_len)
+            if not needs_tokenwise_grammar_accept:
+                req.update_finish_state(new_accepted_len)
 
             self._handle_finish_state_updated_req(req, batch, result, i, logits_output)
 
@@ -689,9 +698,13 @@ class SchedulerBatchResultProcessor:
                     )
 
             if req.grammar is not None:
-                self._apply_decode_grammar(
-                    req=req, next_token_id=next_token_id, batch=batch
-                )
+                if needs_tokenwise_grammar_accept:
+                    # 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 +790,40 @@ class SchedulerBatchResultProcessor:
                     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 speculative grammar tokens until the request finishes.
+
+        Returns the retained prefix and rolls back KV commits for dropped suffix
+        tokens.
+        """
+        accept_tokens = []
+        try:
+            for token_id in proposed:
+                req.grammar.accept_token(token_id)
+                req.output_ids.append(token_id)
+                accept_tokens.append(token_id)
+                self._maybe_update_reasoning_tokens(req, 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))
+            req.update_finish_state()
+
+        # _resolve_spec_v2_tokens committed the full proposed list; rollback the
+        # suffix that grammar termination dropped.
+        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_basic.py

@@ -27,7 +27,7 @@ from sglang.test.test_utils import (
     DEFAULT_TARGET_MODEL_EAGLE3,
 )
 
-register_cuda_ci(est_time=700, stage="base-b", runner_config="2-gpu-large")
+register_cuda_ci(est_time=890, stage="base-b", runner_config="2-gpu-large")
 
 
 class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServerBase):
@@ -260,6 +260,104 @@ class TestDisaggregationMooncakeSpec(JSONConstrainedMixin, PDDisaggregationServe
         self.assertGreater(metrics["score"], 0.74)
 
 
+class TestDisaggregationSpecV2Grammar(PDDisaggregationServerBase):
+    """Regression for PD disagg + EAGLE Spec V2 + grammar structured output."""
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model = DEFAULT_TARGET_MODEL_EAGLE3
+        spec_args = [
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-draft-model-path",
+            DEFAULT_DRAFT_MODEL_EAGLE3,
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",  # Spec V2 only supports topk=1
+            "--speculative-num-draft-tokens",
+            "4",
+            "--grammar-backend",
+            "xgrammar",
+            "--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
+        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):
+        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)
+        out = response.json()
+        self.assertGreater(
+            out["meta_info"]["spec_verify_ct"],
+            0,
+            "expected Spec V2 to run (spec_verify_ct > 0)",
+        )
+        return out
+
+    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"]
+        parsed = json.loads(text)
+        for key in ("name", "population", "country", "capital"):
+            self.assertIn(key, parsed)
+        self.assertTrue(
+            text.strip().endswith("}"), f"unexpected trailing tokens: {text!r}"
+        )
+
+    def test_logprob_count_matches_completion_tokens(self):
+        """Trimmed Spec V2 tokens 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",
+        )
+        json.loads(out["text"])
+
+
 class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
     @classmethod
     def setUpClass(cls):

test/registered/unit/managers/test_batch_result_processor_spec_grammar.py

@@ -0,0 +1,169 @@
+"""Unit tests for Spec V2 grammar trimming in process_batch_result_decode."""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.managers.scheduler_components.batch_result_processor import (
+    SchedulerBatchResultProcessor,
+)
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class _FakeGrammar:
+    """Grammar stub that reports termination after `terminate_after` tokens."""
+
+    def __init__(self, terminate_after: int):
+        self.accepted = []
+        self.finished = False
+        self._terminate_after = terminate_after
+
+    def accept_token(self, token_id: int):
+        self.accepted.append(token_id)
+
+    def is_terminated(self) -> bool:
+        return len(self.accepted) >= self._terminate_after
+
+
+class _FakeSpecAlgorithm:
+    def is_none(self) -> bool:
+        return False
+
+
+class _FakeBatch:
+    def __init__(self, reqs, return_logprob: bool):
+        self.reqs = reqs
+        self.return_logprob = return_logprob
+        self.spec_algorithm = _FakeSpecAlgorithm()
+
+    def batch_size(self) -> int:
+        return len(self.reqs)
+
+
+class _TrimmingProcessor(SchedulerBatchResultProcessor):
+    """Runs the real result processor with GPU/IO-bound helpers stubbed."""
+
+    def _normalize_decode_outputs(
+        self, *, batch, result, logits_output, next_token_ids
+    ):
+        return result.test_next_token_ids, result.test_next_token_logprobs
+
+    def _mamba_prefix_cache_update(self, req, batch, result, i):
+        pass
+
+    def _handle_finish_state_updated_req(self, req, batch, result, i, logits_output):
+        pass
+
+
+def _make_processor() -> _TrimmingProcessor:
+    metrics_reporter = SimpleNamespace(
+        num_generated_tokens=0,
+        forward_ct_decode=0,
+        update_spec_metrics=lambda *a, **k: None,
+        report_decode_stats=lambda *a, **k: None,
+    )
+    allocator = SimpleNamespace(
+        free_group_begin=lambda: None,
+        free_group_end=lambda: None,
+    )
+    output_streamer = SimpleNamespace(stream_output=lambda *a, **k: None)
+    return _TrimmingProcessor(
+        is_generation=True,
+        disaggregation_mode=None,
+        enable_overlap=False,
+        enable_overlap_mlx=False,
+        server_args=SimpleNamespace(enable_metrics=False),
+        model_config=SimpleNamespace(think_end_id=None),
+        token_to_kv_pool_allocator=allocator,
+        tree_cache=None,
+        hisparse_coordinator=None,
+        req_to_token_pool=None,
+        decode_offload_manager=None,
+        metrics_collector=None,
+        metrics_reporter=metrics_reporter,
+        draft_worker=None,
+        model_worker=None,
+        logprob_result_processor=None,
+        output_streamer=output_streamer,
+        abort_request=lambda *a, **k: None,
+    )
+
+
+def _make_result(accept_tokens, logprobs):
+    return SimpleNamespace(
+        copy_done=None,
+        routed_experts_output=None,
+        indexer_topk_output=None,
+        logits_output=SimpleNamespace(hidden_states=None, customized_info=None),
+        next_token_ids=None,
+        can_run_cuda_graph=False,
+        num_correct_drafts=len(accept_tokens),
+        test_next_token_ids=[list(accept_tokens)],
+        test_next_token_logprobs=[list(logprobs)],
+    )
+
+
+class TestSpecV2GrammarTrimming(CustomTestCase):
+    def _make_req(self, terminate_after: int) -> Req:
+        sp = SamplingParams(max_new_tokens=256, temperature=0)
+        sp.normalize(None)
+        req = Req(
+            rid="r0",
+            origin_input_text="",
+            origin_input_ids=[1, 2, 3],
+            sampling_params=sp,
+        )
+        req.vocab_size = 32000
+        req.return_logprob = True
+        req.logprob.output_token_logprobs_val = []
+        req.logprob.output_token_logprobs_idx = []
+        req.grammar = _FakeGrammar(terminate_after=terminate_after)
+        return req
+
+    def test_trims_tokens_after_grammar_completion(self):
+        req = self._make_req(terminate_after=2)
+        proc = _make_processor()
+        result = _make_result([101, 102, 103], [-0.1, -0.2, -0.3])
+        req.kv_committed_len = len(req.origin_input_ids) + len(
+            result.test_next_token_ids[0]
+        )
+        batch = _FakeBatch([req], return_logprob=True)
+
+        proc.process_batch_result_decode(batch, result)
+
+        self.assertTrue(req.finished())
+        self.assertEqual(list(req.output_ids), [101, 102])
+        self.assertEqual(result.test_next_token_ids[0], [101, 102])
+        self.assertEqual(req.grammar.accepted, [101, 102])
+        self.assertTrue(req.grammar.finished)
+        self.assertEqual(req.logprob.output_token_logprobs_val, [-0.1, -0.2])
+        self.assertEqual(req.logprob.output_token_logprobs_idx, [101, 102])
+        self.assertEqual(req.kv_committed_len, len(req.origin_input_ids) + 2)
+
+    def test_keeps_all_tokens_when_grammar_not_terminated(self):
+        req = self._make_req(terminate_after=99)
+        proc = _make_processor()
+        result = _make_result([201, 202, 203], [-0.5, -0.6, -0.7])
+        req.kv_committed_len = len(req.origin_input_ids) + len(
+            result.test_next_token_ids[0]
+        )
+        batch = _FakeBatch([req], return_logprob=True)
+
+        proc.process_batch_result_decode(batch, result)
+
+        self.assertFalse(req.finished())
+        self.assertEqual(list(req.output_ids), [201, 202, 203])
+        self.assertEqual(result.test_next_token_ids[0], [201, 202, 203])
+        self.assertEqual(req.grammar.accepted, [201, 202, 203])
+        self.assertFalse(req.grammar.finished)
+        self.assertEqual(req.logprob.output_token_logprobs_val, [-0.5, -0.6, -0.7])
+        self.assertEqual(req.logprob.output_token_logprobs_idx, [201, 202, 203])
+        self.assertEqual(req.kv_committed_len, len(req.origin_input_ids) + 3)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/spec/test_decode_bookkeeping_ownership.py

@@ -44,6 +44,10 @@ _RESOLVE = (
     "managers/scheduler_components/batch_result_processor.py",
     "SchedulerBatchResultProcessor._resolve_spec_v2_tokens",
 )
+_GRAMMAR_ACCEPT = (
+    "managers/scheduler_components/batch_result_processor.py",
+    "SchedulerBatchResultProcessor._accept_spec_v2_grammar_tokens",
+)
 _SS = "session/streaming_session.py"
 _OWNER_SITES = {
     # non-spec scheduler
@@ -64,6 +68,9 @@ _OWNER_SITES = {
     # pre-claim in prepare_for_decode, unlike the EAGLE mixin).
     (*_RESOLVE, "kv_committed_len"): 3,
     (*_RESOLVE, "spec_verify_ct"): 1,
+    # Spec grammar trim rolls back KV slots for tokens dropped after grammar
+    # completion; resolve already committed the full accepted list.
+    (*_GRAMMAR_ACCEPT, "kv_committed_len"): 1,
     (
         "speculative/dflash_info_v2.py",
         "DFlashDraftInputV2.prepare_for_decode",