Fix stop boundaries for grammar-constrained speculative decoding (#31738)

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.

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()