Preserve speculative finish correctness across reasoning and stop strings

ReasonerGrammarObject wraps an inner grammar, but disaggregated prebuilt replay checks the wrapper current_token to avoid accepting an already accepted token twice. Track the token on the wrapper itself so reasoning grammar follows the same contract as XGrammar.

Speculative decode can accept a stop string and EOS in one step. Check stop strings before token-based EOS after sanitizing invalid token ids, and set finished_len at the matched stop position so trailing accepted tokens do not leak.

Constraint: Current branch predates upstream helper methods for locating string stop positions, so the stop-string fix is manually ported instead of cherry-picked.

Rejected: Direct cherry-pick of bbc853df46 | current schedule_batch.py lacks the upstream helper context.

Confidence: high

Scope-risk: moderate

Directive: Keep vocab-boundary sanitization before string decoding; do not move token-based EOS ahead of stop-string checks without a same-step speculative regression test.

Tested: RED/GREEN remote pytest in cjy-glm5-new for constrained current_token and stop-string speculative tests

Tested: git diff --check; py_compile for reasoner_grammar_backend.py and schedule_batch.py

Not-tested: Full scheduler/disaggregation integration suite

python/sglang/srt/constrained/reasoner_grammar_backend.py

@@ -50,6 +50,7 @@ class ReasonerGrammarObject(BaseGrammarObject):
             self.tokens_after_think_end -= 1
 
     def accept_token(self, token: int):
+        self.current_token = token
         if self.tokens_after_think_end >= 0:
             self.grammar.accept_token(token)
         self.transfer_state(token)

python/sglang/srt/managers/schedule_batch.py

@@ -980,20 +980,25 @@ class Req(ReqDllmMixin):
 
         return self.surr_and_decode_ids, self.read_offset - self.surr_offset
 
-    def tail_str(self) -> str:
+    def _stop_match_tail_len(self, new_accepted_len: int = 1) -> int:
         # Check stop strings and stop regex patterns together
         if (
             len(self.sampling_params.stop_strs) == 0
             and len(self.sampling_params.stop_regex_strs) == 0
         ):
-            return ""
+            return 0
 
         max_len_tail_str = max(
-            self.sampling_params.stop_str_max_len + 1,
-            self.sampling_params.stop_regex_max_len + 1,
+            self.sampling_params.stop_str_max_len + new_accepted_len,
+            self.sampling_params.stop_regex_max_len + new_accepted_len,
         )
 
-        tail_len = min(max_len_tail_str, len(self.output_ids))
+        return min(max_len_tail_str, len(self.output_ids))
+
+    def tail_str(self, new_accepted_len: int = 1) -> str:
+        tail_len = self._stop_match_tail_len(new_accepted_len)
+        if tail_len == 0:
+            return ""
         return self.tokenizer.decode(self.output_ids[-tail_len:])
 
     def check_match_stop_str_prefix(self) -> bool:
@@ -1049,18 +1054,51 @@ class Req(ReqDllmMixin):
 
         return False
 
-    def _check_str_based_finish(self):
+    def _locate_str_stop_finished_len(
+        self,
+        new_accepted_len: int,
+        *,
+        stop_str: Optional[str] = None,
+        stop_regex: Optional[str] = None,
+    ) -> int:
+        """Map a matched stop string/regex to output_ids length (stop included)."""
+
+        def matched(text: str) -> bool:
+            if stop_str is not None:
+                return stop_str in text
+            return re.search(stop_regex, text) is not None
+
+        tail_len = self._stop_match_tail_len(new_accepted_len)
+        start = len(self.output_ids) - tail_len
+        token_window = self.output_ids[start:]
+
+        # Old prefixes were checked in the previous step.
+        for token_count in range(
+            max(1, len(token_window) - new_accepted_len + 1), len(token_window)
+        ):
+            if matched(self.tokenizer.decode(token_window[:token_count])):
+                return start + token_count
+
+        # The full tail window is already known to match by the caller.
+        return len(self.output_ids)
+
+    def _check_str_based_finish(self, new_accepted_len: int = 1):
         if (
             len(self.sampling_params.stop_strs) > 0
             or len(self.sampling_params.stop_regex_strs) > 0
         ):
-            tail_str = self.tail_str()
+            tail_str = self.tail_str(new_accepted_len)
 
             # Check stop strings
             if len(self.sampling_params.stop_strs) > 0:
                 for stop_str in self.sampling_params.stop_strs:
-                    if stop_str in tail_str or stop_str in self.decoded_text:
+                    stop_str_in_tail = stop_str in tail_str
+                    if stop_str_in_tail or stop_str in self.decoded_text:
                         self.finished_reason = FINISH_MATCHED_STR(matched=stop_str)
+                        if stop_str_in_tail:
+                            self.finished_len = self._locate_str_stop_finished_len(
+                                new_accepted_len, stop_str=stop_str
+                            )
                         return True
 
             # Check stop regex
@@ -1070,6 +1108,9 @@ class Req(ReqDllmMixin):
                         self.finished_reason = FINISHED_MATCHED_REGEX(
                             matched=stop_regex_str
                         )
+                        self.finished_len = self._locate_str_stop_finished_len(
+                            new_accepted_len, stop_regex=stop_regex_str
+                        )
                         return True
 
         return False
@@ -1113,13 +1154,17 @@ class Req(ReqDllmMixin):
 
         new_accepted_tokens = self.output_ids[-new_accepted_len:]
 
-        if self._check_token_based_finish(new_accepted_tokens):
-            return
-
+        # Sanitize out-of-range / NaN token ids before any decode.
         if self._check_vocab_boundary_finish(new_accepted_tokens):
             return
 
-        if self._check_str_based_finish():
+        # Stop string beats EOS/stop-token matched in the same step. Speculative
+        # decoding can accept multiple tokens, and token-based finish would trim
+        # at EOS while leaking a stop string accepted earlier in the same step.
+        if self._check_str_based_finish(new_accepted_len):
+            return
+
+        if self._check_token_based_finish(new_accepted_tokens):
             return
 
     def reset_for_retract(self):

test/registered/unit/constrained/test_reasoner_grammar_backend.py

@@ -0,0 +1,69 @@
+import unittest
+
+from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
+
+
+class _InnerGrammar:
+    def __init__(self):
+        self.accepted_tokens = []
+
+    def accept_token(self, token: int):
+        self.accepted_tokens.append(token)
+
+    def is_terminated(self):
+        return False
+
+    def rollback(self, k: int):
+        self.accepted_tokens = self.accepted_tokens[:-k]
+
+    def allocate_vocab_mask(self, vocab_size: int, batch_size: int, device):
+        return None
+
+    def fill_vocab_mask(self, vocab_mask, idx: int):
+        return None
+
+    def move_vocab_mask(self, vocab_mask, device):
+        return vocab_mask
+
+    @property
+    def apply_vocab_mask(self):
+        return None
+
+    def copy(self):
+        ret = _InnerGrammar()
+        ret.accepted_tokens = list(self.accepted_tokens)
+        return ret
+
+    @property
+    def finished(self):
+        return False
+
+
+class TestReasonerGrammarCurrentToken(unittest.TestCase):
+    def test_wrapper_tracks_current_token_for_disagg_prebuilt_dedup(self):
+        inner = _InnerGrammar()
+        grammar = ReasonerGrammarObject(inner, think_end_id=7)
+        grammar.maybe_init_reasoning(True)
+
+        grammar.accept_token(10)  # still in thinking; inner grammar is untouched
+        self.assertEqual(grammar.current_token, 10)
+        self.assertEqual(inner.accepted_tokens, [])
+
+        grammar.accept_token(7)  # exits thinking
+        self.assertEqual(grammar.current_token, 7)
+        self.assertEqual(inner.accepted_tokens, [])
+
+        grammar.accept_token(58)  # generation token accepted by inner grammar
+        self.assertEqual(grammar.current_token, 58)
+        self.assertEqual(inner.accepted_tokens, [58])
+
+        # Mirrors disaggregation/decode_schedule_batch_mixin.py:process_prebuilt.
+        # Once current_token is tracked on the wrapper, a re-prebuilt request
+        # skips re-accepting the already accepted token.
+        if grammar.current_token is None:
+            grammar.accept_token(58)
+        self.assertEqual(inner.accepted_tokens, [58])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/managers/test_stop_str_speculative.py

@@ -0,0 +1,61 @@
+import unittest
+from array import array
+
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.sampling.sampling_params import SamplingParams
+
+
+STOP_ID = 1
+EOS_ID = 2
+ID_TO_TEXT = {
+    STOP_ID: "STOP",
+    EOS_ID: "",
+    **{i: chr(ord("a") + i % 26) for i in range(10, 40)},
+}
+
+
+class _FakeTokenizer:
+    eos_token_id = EOS_ID
+    additional_stop_token_ids = None
+
+    def decode(self, ids):
+        return "".join(ID_TO_TEXT[int(i)] for i in ids)
+
+
+class _NormalizeTokenizer:
+    def encode(self, text, add_special_tokens=False):
+        return list(range(len(text)))
+
+
+def _make_req(output_ids, *, stop=None, eos_token_ids=frozenset()):
+    sampling_params = SamplingParams(max_new_tokens=1000, stop=stop)
+    sampling_params.normalize(tokenizer=_NormalizeTokenizer())
+    req = Req(
+        rid="r",
+        origin_input_text="",
+        origin_input_ids=array("q", [0]),
+        sampling_params=sampling_params,
+        eos_token_ids=set(eos_token_ids),
+        vocab_size=10_000,
+    )
+    req.tokenizer = _FakeTokenizer()
+    req.output_ids = array("q", output_ids)
+    return req
+
+
+class TestStopStrSpeculative(unittest.TestCase):
+    def test_stop_str_wins_over_eos_in_same_spec_step(self):
+        # A speculative step may accept both the stop string and EOS. The stop
+        # string must finish first so finished_len trims at STOP instead of EOS.
+        req = _make_req([10, 11, STOP_ID, EOS_ID], stop=["STOP"], eos_token_ids={EOS_ID})
+
+        req.check_finished(new_accepted_len=4)
+
+        self.assertTrue(req.finished())
+        self.assertEqual(req.finished_reason.matched, "STOP")
+        self.assertEqual(req.finished_len, 3)
+        self.assertEqual(list(req.output_ids_through_stop), [10, 11, STOP_ID])
+
+
+if __name__ == "__main__":
+    unittest.main()