diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 27eb699b2c3a..b958ab6e1fa1 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1628,6 +1628,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) @@ -1637,8 +1653,10 @@ 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) + # Skip if need_grammar_sync already drained the queued result above. + 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() 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 76f5cc13a624..4d37b6fe0ee1 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -653,18 +653,55 @@ 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: the verify phase proposes several tokens at once, + # but the grammar may terminate partway through that list. Accept the + # proposed tokens one at a time and stop as soon as the request finishes + # so we don't advance output_ids, the grammar FSM, reasoning state, or + # logprob bookkeeping past grammar completion. grammar_advanced records + # that the grammar/finish state was already handled inline below. + grammar_advanced = False if batch.spec_algorithm.is_none(): req.output_ids.append(next_token_id) + elif batch.is_spec_v2 and req.grammar is not None: + accept_tokens = [] + try: + for token_id in next_token_id: + 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: + # Grammar accept_token can raise ValueError if the token is not + # in the grammar. This can happen if the grammar is not set + # correctly or the token is invalid. + logger.error( + f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" + ) + self.abort_request(AbortReq(rid=req.rid)) + + # Drop speculative tokens proposed after the grammar reached a + # terminal state. The request is finished, so the over-advanced + # KV/draft state is released instead of reused, and downstream + # logprob handling only sees the retained prefix. + next_token_id = accept_tokens + next_token_ids[i] = accept_tokens + new_accepted_len = len(accept_tokens) + grammar_advanced = True 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 grammar_advanced: + 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 grammar_advanced: + req.update_finish_state(new_accepted_len) self._handle_finished_req(req, i, logits_output) @@ -684,9 +721,14 @@ 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 grammar_advanced: + # Grammar was already advanced token-by-token above; just sync + # the terminal flag without re-accepting the trimmed tokens. + 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() diff --git a/test/registered/disaggregation/test_disaggregation_spec_grammar.py b/test/registered/disaggregation/test_disaggregation_spec_grammar.py new file mode 100644 index 000000000000..b850a5a297ba --- /dev/null +++ b/test/registered/disaggregation/test_disaggregation_spec_grammar.py @@ -0,0 +1,165 @@ +"""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 +from types import SimpleNamespace + +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", + ] + cls.extra_prefill_args = spec_args + cls.extra_decode_args = spec_args + with ( + envs.SGLANG_ENABLE_SPEC_V2.override(True), + # The EAGLE3 draft model config derives a 2048 context length, which is + # shorter than the Llama-3.1 target's 131072. The Spec V2 draft worker + # (eagle_worker_v2.py) builds its own ModelConfig and rejects this + # mismatch unless overriding longer context is allowed. Outputs here are + # well under 2048 tokens, so allowing the override is safe. + envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.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() 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 000000000000..44efc2f0cf4f --- /dev/null +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -0,0 +1,186 @@ +"""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). + +This drives the real decode result processor on CPU with plain-Python stubs for +the GPU/IO-bound helpers. The core fix: when Spec V2 proposes several tokens but +the grammar terminates partway through the list, only the accepted prefix is +retained in output_ids, the grammar FSM, and logprob bookkeeping. +""" + +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_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large") + + +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.is_spec_v2 = True + self.spec_algorithm = _FakeSpecAlgorithm() + + def batch_size(self) -> int: + return len(self.reqs) + + +class _TrimmingProcessor(SchedulerBatchResultProcessor): + """Overrides GPU/IO-bound helpers so the Spec V2 grammar trimming branch of + process_batch_result_decode can run on CPU. + + _normalize_decode_outputs returns the test-provided token/logprob lists + directly (the real one round-trips CUDA tensors and feeds the spec worker). + The returned next_token_ids list is the same object the production code + mutates in place via ``next_token_ids[i] = accept_tokens``, so the trim is + observable on ``result.test_next_token_ids`` after the call. + """ + + 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_finished_req(self, req, 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): + # Grammar terminates after the 2nd accepted token; the 3rd proposed token + # must be dropped everywhere. + req = self._make_req(terminate_after=2) + proc = _make_processor() + result = _make_result([101, 102, 103], [-0.1, -0.2, -0.3]) + batch = _FakeBatch([req], return_logprob=True) + + proc.process_batch_result_decode(batch, result) + + self.assertTrue(req.finished()) + # output_ids may be a list or array('q', ...) depending on sglang version. + self.assertEqual(list(req.output_ids), [101, 102]) + # next_token_ids[i] was trimmed in place for downstream consumers. + self.assertEqual(result.test_next_token_ids[0], [101, 102]) + # Grammar advanced exactly over the retained prefix and is marked terminal. + self.assertEqual(req.grammar.accepted, [101, 102]) + 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]) + + def test_keeps_all_tokens_when_grammar_not_terminated(self): + # Grammar never terminates within this list: all proposed tokens retained. + req = self._make_req(terminate_after=99) + proc = _make_processor() + result = _make_result([201, 202, 203], [-0.5, -0.6, -0.7]) + 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]) + + +if __name__ == "__main__": + unittest.main()