From 9fc9d37f6d10eb1fba46d4eaed2abdad5c530807 Mon Sep 17 00:00:00 2001 From: Ashish <1856117+ashishdatta@users.noreply.github.com> Date: Thu, 18 Jun 2026 14:58:32 -0700 Subject: [PATCH] Fix spec decoding with grammar in disagg (#24082) Co-authored-by: jimmy.shong Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Co-authored-by: Xinyuan Tong Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Co-authored-by: Codex --- .../batch_result_processor.py | 57 +++++- .../test_disaggregation_basic.py | 100 ++++++++++- ...est_batch_result_processor_spec_grammar.py | 169 ++++++++++++++++++ .../spec/test_decode_bookkeeping_ownership.py | 7 + 4 files changed, 327 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/managers/test_batch_result_processor_spec_grammar.py diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index a636bf201..77221cd37 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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, *, diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index 74509a6d0..ca9e250c3 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/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): diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py new file mode 100644 index 000000000..ac1949c95 --- /dev/null +++ b/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() diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 09f967780..a0f83dd15 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/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",