From 9462c303a5c5349488e84b2c346e9b5a066de21c Mon Sep 17 00:00:00 2001 From: Yang Liu Date: Mon, 20 Jul 2026 17:55:27 -0700 Subject: [PATCH] Fix stop boundaries for grammar-constrained speculative decoding (#31738) --- python/sglang/srt/managers/schedule_batch.py | 9 ++- .../managers/test_grammar_stop_speculative.py | 69 +++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/managers/test_grammar_stop_speculative.py diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 14dbcfefe..2dee4ec64 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1489,11 +1489,6 @@ class Req(ReqDllmMixin): self.finished_len = self.sampling_params.max_new_tokens return - if self.grammar is not None: - if self.grammar.is_terminated(): - self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1]) - return - new_accepted_tokens = self.output_ids[-new_accepted_len:] # Sanitize out-of-range / NaN token ids before any decode. @@ -1509,6 +1504,10 @@ class Req(ReqDllmMixin): if self._check_token_based_finish(new_accepted_tokens): return + if self.grammar is not None and self.grammar.is_terminated(): + self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1]) + return + def reset_for_retract(self): # Increment retraction count before resetting other state. We should not reset this # since we are tracking the total number of retractions for each request. diff --git a/test/registered/unit/managers/test_grammar_stop_speculative.py b/test/registered/unit/managers/test_grammar_stop_speculative.py new file mode 100644 index 000000000..564489ece --- /dev/null +++ b/test/registered/unit/managers/test_grammar_stop_speculative.py @@ -0,0 +1,69 @@ +"""Regression tests for grammar termination during speculative decoding.""" + +import unittest +from array import array + +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +EOS_TOKEN_ID = 2 +STOP_TOKEN_ID = 17 + + +class _FakeTokenizer: + eos_token_id = -1 + additional_stop_token_ids = None + + +class _TerminatedGrammar: + @staticmethod + def is_terminated(): + return True + + +def _make_req(stop_token_ids=None): + sampling_params = SamplingParams( + max_new_tokens=1_000, + stop_token_ids=stop_token_ids, + ) + sampling_params.normalize(tokenizer=_FakeTokenizer()) + req = Req( + rid="grammar-stop", + origin_input_text="", + origin_input_ids=array("q", [0]), + sampling_params=sampling_params, + eos_token_ids={EOS_TOKEN_ID}, + vocab_size=100, + ) + req.tokenizer = _FakeTokenizer() + req.grammar = _TerminatedGrammar() + req.output_ids = array("q", [11, 13, STOP_TOKEN_ID, EOS_TOKEN_ID]) + return req + + +class TestGrammarStopSpeculative(unittest.TestCase): + def test_requested_stop_token_wins_over_trailing_eos(self): + req = _make_req(stop_token_ids=[STOP_TOKEN_ID]) + + req.update_finish_state(new_accepted_len=4) + + self.assertEqual(req.finished_reason.matched, STOP_TOKEN_ID) + self.assertEqual(req.finished_len, 3) + self.assertEqual(list(req.output_ids_through_stop), [11, 13, STOP_TOKEN_ID]) + + def test_tokenizer_stop_token_wins_over_trailing_eos(self): + req = _make_req() + req.tokenizer.additional_stop_token_ids = {STOP_TOKEN_ID} + + req.update_finish_state(new_accepted_len=4) + + self.assertEqual(req.finished_reason.matched, STOP_TOKEN_ID) + self.assertEqual(req.finished_len, 3) + self.assertEqual(list(req.output_ids_through_stop), [11, 13, STOP_TOKEN_ID]) + + +if __name__ == "__main__": + unittest.main()