diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 065546088..b4f597d2b 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -169,7 +169,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): def trim_matched_stop( self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool ): - if no_stop_trim or not finished_reason: + if not finished_reason: return output matched = finished_reason.get("matched", None) @@ -181,10 +181,15 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): # Trim stop str. if isinstance(matched, str) and isinstance(output, str): pos = output.find(matched) - return output[:pos] if pos != -1 else output + if pos == -1: + return output + end = pos + len(matched) + return output[:end] if no_stop_trim else output[:pos] # Trim stop token. if isinstance(matched, int) and isinstance(output, list): + if no_stop_trim: + return output # 200012 <|call|> is the tool call token and one of eos tokens for gpt-oss model if output[-1] == 200012 and self.is_tool_call_parser_gpt_oss: return output diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c85fcacad..e29d6413b 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1192,6 +1192,17 @@ class Req(ReqDllmMixin): return self.surr_and_decode_ids, self.read_offset - self.surr_offset + def _stop_match_tail_len(self, new_accepted_len: int) -> int: + max_len_tail_str = max( + self.sampling_params.stop_str_max_len + 1, + self.sampling_params.stop_regex_max_len + 1, + ) + # Cover all newly accepted tokens so an early stop string is not missed + # when speculative decoding accepts multiple tokens per step. + return min( + max_len_tail_str + max(new_accepted_len - 1, 0), len(self.output_ids) + ) + def tail_str(self, new_accepted_len: int = 1) -> str: # Check stop strings and stop regex patterns together if ( @@ -1200,17 +1211,7 @@ class Req(ReqDllmMixin): ): return "" - max_len_tail_str = max( - self.sampling_params.stop_str_max_len + 1, - self.sampling_params.stop_regex_max_len + 1, - ) - - # Spec decode accepts multiple tokens per step; widen the window to cover - # the whole accepted chunk so a stop string landing mid-chunk (with more - # tokens accepted after it) is not pushed out of view. - tail_len = min( - max_len_tail_str + max(new_accepted_len - 1, 0), len(self.output_ids) - ) + tail_len = self._stop_match_tail_len(new_accepted_len) return self.tokenizer.decode(self.output_ids[-tail_len:]) def check_match_stop_str_prefix(self) -> bool: @@ -1266,6 +1267,34 @@ class Req(ReqDllmMixin): return False + 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 @@ -1276,8 +1305,13 @@ class Req(ReqDllmMixin): # 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 @@ -1287,6 +1321,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 diff --git a/test/registered/unit/managers/test_stop_str_speculative.py b/test/registered/unit/managers/test_stop_str_speculative.py index ce1bf4fcb..351d84cd1 100644 --- a/test/registered/unit/managers/test_stop_str_speculative.py +++ b/test/registered/unit/managers/test_stop_str_speculative.py @@ -1,6 +1,9 @@ -"""Regression: under speculative decoding (multi-token commits) a stop string -committed mid-chunk must still trigger the finish check, else the request -over-generates. Drives the real `Req.update_finish_state`; pure CPU.""" +"""Regression for stop-string / stop-regex finishing under speculative decoding +(multi-token commits): a stop committed mid-chunk must (1) trigger the finish +check and (2) set finished_len so the emitted output is trimmed at the stop, not +leaking tokens accepted after it. Drives the real `Req.update_finish_state` with +a fake tokenizer; pure CPU. Each test guards a distinct branch of +`_locate_str_stop_finished_len` / `_check_str_based_finish`.""" import unittest from array import array @@ -11,8 +14,19 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") +# Token id -> decoded text; decode() concatenates. Distinct symbols so no +# accidental cross-matches (10-39 are lowercase letters). STOP_ID = 1 -ID_TO_TEXT = {STOP_ID: "STOP", **{i: chr(ord("a") + i % 26) for i in range(10, 40)}} +ID_TO_TEXT = { + STOP_ID: "STOP", + **{i: chr(ord("a") + i % 26) for i in range(10, 40)}, + 60: "a", + 61: ".", + 62: "b", + 63: ".", + 70: "X", + 71: "Y", +} # "STOP" (index 3) sits 6 tokens back: outside the old (stop_str_max_len + 1) # window, inside the one widened by new_accepted_len. @@ -27,8 +41,8 @@ class _FakeTokenizer: return "".join(ID_TO_TEXT[int(i)] for i in ids) -def _make_req(output_ids, stop): - sp = SamplingParams(max_new_tokens=1000, stop=stop) +def _make_req(output_ids, stop=None, stop_regex=None): + sp = SamplingParams(max_new_tokens=1000, stop=stop, stop_regex=stop_regex) sp.normalize(tokenizer=None) # char-based stop_str_max_len req = Req( rid="t", @@ -44,17 +58,73 @@ def _make_req(output_ids, stop): class TestStopStrSpeculative(unittest.TestCase): - def test_stop_str_midchunk_finishes(self): - req = _make_req(MIDCHUNK, stop=["STOP"]) - req.update_finish_state(new_accepted_len=6) - self.assertTrue(req.finished()) - self.assertEqual(req.finished_reason.matched, "STOP") - - def test_no_stop_str_does_not_finish(self): + def test_no_stop_does_not_finish(self): req = _make_req([10, 11, 12, 20, 21, 22, 23, 24], stop=["STOP"]) req.update_finish_state(new_accepted_len=6) self.assertFalse(req.finished()) + # --- _locate_str_stop_finished_len: loop match / fallback / empty loop / span --- + def test_stop_str_midchunk(self): + # Loop match: "STOP" (index 3) finishes; finished_len lands just past it, + # so the 5 trailing tokens are dropped. + req = _make_req(MIDCHUNK, stop=["STOP"]) + req.update_finish_state(new_accepted_len=6) + self.assertTrue(req.finished()) + self.assertEqual(req.finished_reason.matched, "STOP") + self.assertEqual(req.finished_len, 4) + + def test_stop_str_at_chunk_end_uses_full_len(self): + # Stop is the last token -> loop never matches before the full window -> + # fallback returns len(output_ids). + req = _make_req([10, 11, 12, 20, 21, STOP_ID], stop=["STOP"]) + req.update_finish_state(new_accepted_len=6) + self.assertTrue(req.finished()) + self.assertEqual(req.finished_len, 6) + + def test_stop_str_non_spec_single_token(self): + # new_accepted_len == 1: locate's range is empty -> fallback (non-spec + # path preserved, no extra decode). + req = _make_req([10, 11, STOP_ID], stop=["STOP"]) + req.update_finish_state(new_accepted_len=1) + self.assertTrue(req.finished()) + self.assertEqual(req.finished_len, 3) + + def test_stop_str_spanning_two_tokens(self): + # "XY" completes only once both tokens (70, 71) are decoded -> finished_len + # covers both; trailing tokens dropped. + req = _make_req([10, 11, 70, 71, 20, 21], stop=["XY"]) + req.update_finish_state(new_accepted_len=6) + self.assertTrue(req.finished()) + self.assertEqual(req.finished_len, 4) + + # --- regex matched() branch --- + def test_stop_regex_midchunk(self): + req = _make_req([10, 11, 70, 71, 20, 21], stop_regex=[r"XY"]) + req.update_finish_state(new_accepted_len=6) + self.assertTrue(req.finished()) + self.assertEqual(req.finished_reason.matched, r"XY") + self.assertEqual(req.finished_len, 4) + + def test_stop_regex_end_anchored_trims_at_first_match(self): + # Documents current behavior: text "a.b." matches `\.$` at the chunk end, + # but locate scans growing prefixes and stops at the FIRST period ("a."), + # so finished_len == 2 and "b." is dropped. + req = _make_req([60, 61, 62, 63], stop_regex=[r"\.$"]) + req.update_finish_state(new_accepted_len=4) + self.assertTrue(req.finished()) + self.assertEqual(req.finished_reason.matched, r"\.$") + self.assertEqual(req.finished_len, 2) + + # --- decoded_text-only branch --- + def test_stop_str_only_in_decoded_text_sets_no_finished_len(self): + # Documents current behavior: when the stop is only in decoded_text (not + # the tail), the request finishes but finished_len is left unset. + req = _make_req([10, 11, 12], stop=["STOP"]) # no STOP in output tokens + req.decoded_text = "earlier STOP text" + req.update_finish_state(new_accepted_len=3) + self.assertTrue(req.finished()) + self.assertIsNone(req.finished_len) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_trim_matched_stop.py b/test/registered/unit/managers/test_trim_matched_stop.py new file mode 100644 index 000000000..3fd09e59c --- /dev/null +++ b/test/registered/unit/managers/test_trim_matched_stop.py @@ -0,0 +1,68 @@ +"""Unit test for DetokenizerManager.trim_matched_stop. + +Under speculative decoding the output handed to trim_matched_stop may carry +content after the matched stop. With no_stop_trim the stop string must be kept +but the trailing over-generation still dropped (`output[:end]`); without it the +stop is removed (`output[:pos]`). Pure CPU: calls the method with a stub self, +so no DetokenizerManager.__init__ / IPC / tokenizer.""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.managers.detokenizer_manager import DetokenizerManager +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +GPT_OSS_CALL_TOKEN = 200012 + + +def _trim(output, matched, no_stop_trim, *, gpt_oss=False): + stub = SimpleNamespace(is_tool_call_parser_gpt_oss=gpt_oss) + finished_reason = None if matched is None else {"matched": matched} + return DetokenizerManager.trim_matched_stop( + stub, output, finished_reason, no_stop_trim + ) + + +class TestTrimMatchedStop(unittest.TestCase): + def test_no_finished_reason_returns_output(self): + self.assertEqual(_trim("abc", None, False), "abc") + + def test_no_matched_returns_output(self): + stub = SimpleNamespace(is_tool_call_parser_gpt_oss=False) + self.assertEqual( + DetokenizerManager.trim_matched_stop(stub, "abc", {}, False), "abc" + ) + + # --- stop string --- + def test_str_trim_removes_stop(self): + # no_stop_trim=False: drop the stop string and anything after it. + self.assertEqual(_trim("ans\n\nQuestion: A", "Question", False), "ans\n\n") + + def test_str_no_trim_keeps_stop_but_drops_trailing(self): + # no_stop_trim=True: keep through the stop, drop the over-generated tail. + self.assertEqual( + _trim("ans\n\nQuestion: A", "Question", True), "ans\n\nQuestion" + ) + + def test_str_not_found_returns_output(self): + self.assertEqual(_trim("no stop here", "Question", False), "no stop here") + + # --- stop token --- + def test_token_trim_drops_last(self): + self.assertEqual(_trim([1, 2, 3], 3, False), [1, 2]) + + def test_token_no_trim_keeps_all(self): + self.assertEqual(_trim([1, 2, 3], 3, True), [1, 2, 3]) + + def test_token_gpt_oss_call_kept(self): + # gpt-oss tool-call token is also an eos; keep it even when trimming. + self.assertEqual( + _trim([1, 2, GPT_OSS_CALL_TOKEN], GPT_OSS_CALL_TOKEN, False, gpt_oss=True), + [1, 2, GPT_OSS_CALL_TOKEN], + ) + + +if __name__ == "__main__": + unittest.main()