From 132561a151924307cbf40f32f0080281eaecb1bd Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sat, 27 Jun 2026 20:36:52 +0800 Subject: [PATCH] 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 --- .../constrained/reasoner_grammar_backend.py | 1 + python/sglang/srt/managers/schedule_batch.py | 69 +++++++++++++++---- .../test_reasoner_grammar_backend.py | 69 +++++++++++++++++++ .../managers/test_stop_str_speculative.py | 61 ++++++++++++++++ 4 files changed, 188 insertions(+), 12 deletions(-) create mode 100644 test/registered/unit/constrained/test_reasoner_grammar_backend.py create mode 100644 test/registered/unit/managers/test_stop_str_speculative.py diff --git a/python/sglang/srt/constrained/reasoner_grammar_backend.py b/python/sglang/srt/constrained/reasoner_grammar_backend.py index d204bdd9e..3deb6392d 100644 --- a/python/sglang/srt/constrained/reasoner_grammar_backend.py +++ b/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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index b92c1502a..73ab3b4e7 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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): diff --git a/test/registered/unit/constrained/test_reasoner_grammar_backend.py b/test/registered/unit/constrained/test_reasoner_grammar_backend.py new file mode 100644 index 000000000..f14271646 --- /dev/null +++ b/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() diff --git a/test/registered/unit/managers/test_stop_str_speculative.py b/test/registered/unit/managers/test_stop_str_speculative.py new file mode 100644 index 000000000..6079e309b --- /dev/null +++ b/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()