From 85b31d9ff4e69a45a571a85ba62398faf025ea15 Mon Sep 17 00:00:00 2001 From: xythink Date: Mon, 27 Apr 2026 12:43:11 +0800 Subject: [PATCH 1/3] fix: expand tail_str window to cover all accepted tokens in speculative decoding The tail_str() window for stop-string matching was sized only by stop_str_max_len, ignoring how many tokens were accepted in one speculative decoding step. When multiple tokens were accepted at once, an early stop string could fall outside the tail window and never be detected. This fix passes new_accepted_len through to tail_str() so the window always covers all newly accepted tokens plus the stop-string length. --- python/sglang/srt/managers/schedule_batch.py | 16 +- .../managers/test_stop_str_speculative.py | 159 ++++++++++++++++++ 2 files changed, 170 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/managers/test_stop_str_speculative.py diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 6c3e1f35066e..3b85b0e386a0 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1081,7 +1081,7 @@ def init_incremental_detokenize(self): return self.surr_and_decode_ids, self.read_offset - self.surr_offset - def tail_str(self) -> str: + def tail_str(self, new_accepted_len: int = 1) -> str: # Check stop strings and stop regex patterns together if ( len(self.sampling_params.stop_strs) == 0 @@ -1094,7 +1094,13 @@ def tail_str(self) -> str: self.sampling_params.stop_regex_max_len + 1, ) - tail_len = min(max_len_tail_str, len(self.output_ids)) + # Ensure the window covers all newly accepted tokens plus the + # stop-string length so that multi-token speculative acceptance + # never pushes an early stop string out of view. + tail_len = min( + max(max_len_tail_str, new_accepted_len + max_len_tail_str), + len(self.output_ids), + ) return self.tokenizer.decode(self.output_ids[-tail_len:]) def check_match_stop_str_prefix(self) -> bool: @@ -1150,12 +1156,12 @@ def _check_token_based_finish(self, new_accepted_tokens: List[int]) -> bool: return False - def _check_str_based_finish(self): + 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: @@ -1220,7 +1226,7 @@ def check_finished(self, new_accepted_len: int = 1): if self._check_vocab_boundary_finish(new_accepted_tokens): return - if self._check_str_based_finish(): + if self._check_str_based_finish(new_accepted_len): return def reset_for_retract(self): 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 000000000000..b4372f29dd50 --- /dev/null +++ b/test/registered/unit/managers/test_stop_str_speculative.py @@ -0,0 +1,159 @@ +"""Test that stop-string checking covers all newly accepted tokens, +not just the tail window — critical for speculative decoding. + +This test uses pure mocks so it can run without sglang dependencies.""" + +import unittest +from unittest.mock import MagicMock + + +def tail_str_new(self, new_accepted_len: int = 1) -> str: + """Fixed version of tail_str that expands window for speculative decoding.""" + if ( + len(self.sampling_params.stop_strs) == 0 + and len(self.sampling_params.stop_regex_strs) == 0 + ): + return "" + + max_len_tail_str = max( + self.sampling_params.stop_str_max_len + 1, + self.sampling_params.stop_regex_max_len + 1, + ) + + # Ensure the window covers all newly accepted tokens plus the + # stop-string length so that multi-token speculative acceptance + # never pushes an early stop string out of view. + tail_len = min( + max(max_len_tail_str, new_accepted_len + max_len_tail_str), + len(self.output_ids), + ) + return self.tokenizer.decode(self.output_ids[-tail_len:]) + + +def tail_str_old(self) -> str: + """Original (buggy) version of tail_str for comparison.""" + if ( + len(self.sampling_params.stop_strs) == 0 + and len(self.sampling_params.stop_regex_strs) == 0 + ): + return "" + + max_len_tail_str = max( + self.sampling_params.stop_str_max_len + 1, + self.sampling_params.stop_regex_max_len + 1, + ) + + tail_len = min(max_len_tail_str, len(self.output_ids)) + return self.tokenizer.decode(self.output_ids[-tail_len:]) + + +class TestStopStrSpeculative(unittest.TestCase): + """Verify tail_str window expands to cover multi-token acceptance.""" + + def _make_req(self, stop_strs, output_ids, stop_str_max_len=None): + """Create a minimal mock that looks like a Req for tail_str.""" + req = MagicMock() + req.output_ids = output_ids + req.sampling_params.stop_strs = stop_strs + req.sampling_params.stop_regex_strs = [] + req.sampling_params.stop_str_max_len = ( + stop_str_max_len if stop_str_max_len is not None + else max(len(s) for s in stop_strs) + ) + req.sampling_params.stop_regex_max_len = 0 + + # Simple tokenizer mock: decode token ids to characters + tok = MagicMock() + def fake_decode(ids): + return "".join(chr(i) for i in ids) + tok.decode = fake_decode + req.tokenizer = tok + + return req + + def test_default_window_near_old_behavior(self): + """With new_accepted_len=1 (default), window is very close to old logic. + + The new formula adds +1 token to the window when new_accepted_len=1, + which is negligible and safe for non-speculative paths. + """ + stop_strs = ["hello"] # stop_str_max_len = 5 + output_ids = list(range(65, 85)) # 20 tokens (ASCII A-T) + req = self._make_req(stop_strs, output_ids) + + # Old: tail_len = min(5+1, 20) = 6 + # New with new_accepted_len=1: tail_len = min(max(6, 1+6), 20) = 7 + old_result = tail_str_old(req) + new_result = tail_str_new(req, new_accepted_len=1) + # New window is 1 token larger (7 vs 6), which is safe + self.assertEqual(len(old_result), 6) + self.assertEqual(len(new_result), 7) + # New result contains old result as suffix + self.assertTrue(new_result.endswith(old_result)) + + def test_expanded_window_with_speculative(self): + """With new_accepted_len > max_len_tail_str, window expands.""" + stop_strs = ["hi"] # stop_str_max_len = 2 + output_ids = list(range(65, 85)) # 20 tokens + req = self._make_req(stop_strs, output_ids) + + # max_len_tail_str = max(2+1, 0+1) = 3 + # With new_accepted_len=10: tail_len = min(max(3, 10+3), 20) = 13 + result = tail_str_new(req, new_accepted_len=10) + self.assertEqual(len(result), 13) + + def test_window_clamped_to_output_len(self): + """Window never exceeds len(output_ids).""" + stop_strs = ["hi"] + output_ids = list(range(65, 70)) # only 5 tokens + req = self._make_req(stop_strs, output_ids) + + # new_accepted_len=100 would want window=103, but clamped to 5 + result = tail_str_new(req, new_accepted_len=100) + self.assertEqual(len(result), 5) + + def test_stop_str_missed_by_old_window(self): + """Old window misses stop string at the beginning of a large batch.""" + stop_strs = ["AB"] # stop_str_max_len = 2 + # "AB" at positions 0-1, then 18 filler chars = 20 total + output_ids = [65, 66] + list(range(67, 87)) + req = self._make_req(stop_strs, output_ids) + + # Old: tail_len = min(2+1, 20) = 3 → last 3 chars only + old_tail = tail_str_old(req) + self.assertNotIn("AB", old_tail, "Old window should miss the stop string") + + def test_stop_str_found_by_expanded_window(self): + """Expanded window finds stop string at the beginning of a large batch.""" + stop_strs = ["AB"] # stop_str_max_len = 2 + output_ids = [65, 66] + list(range(67, 87)) + req = self._make_req(stop_strs, output_ids) + + # New: tail_len = min(max(3, 20+3), 20) = 20 → all chars decoded + new_tail = tail_str_new(req, new_accepted_len=20) + self.assertIn("AB", new_tail, "Expanded window should find the stop string") + + def test_no_stop_strs_returns_empty(self): + """Returns empty string when no stop strings configured.""" + output_ids = list(range(65, 85)) + req = self._make_req(["x"], output_ids) # placeholder for mock setup + req.sampling_params.stop_strs = [] + + result = tail_str_new(req) + self.assertEqual(result, "") + + def test_regex_max_len_also_expands_window(self): + """When stop_regex_max_len is larger, it also benefits from expansion.""" + stop_strs = ["hi"] + output_ids = list(range(65, 85)) + req = self._make_req(stop_strs, output_ids) + req.sampling_params.stop_regex_max_len = 5 + + # max_len_tail_str = max(2+1, 5+1) = 6 + # With new_accepted_len=10: tail_len = min(max(6, 10+6), 20) = 16 + result = tail_str_new(req, new_accepted_len=10) + self.assertEqual(len(result), 16) + + +if __name__ == "__main__": + unittest.main() From 840ec537f0ddc01d06413fff1e3d2418b09b0730 Mon Sep 17 00:00:00 2001 From: xythink Date: Mon, 27 Apr 2026 12:57:31 +0800 Subject: [PATCH 2/3] style: add CI registry and fix formatting for stop-str test --- .../unit/managers/test_stop_str_speculative.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/test/registered/unit/managers/test_stop_str_speculative.py b/test/registered/unit/managers/test_stop_str_speculative.py index b4372f29dd50..f372514390f0 100644 --- a/test/registered/unit/managers/test_stop_str_speculative.py +++ b/test/registered/unit/managers/test_stop_str_speculative.py @@ -6,6 +6,10 @@ import unittest from unittest.mock import MagicMock +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + def tail_str_new(self, new_accepted_len: int = 1) -> str: """Fixed version of tail_str that expands window for speculative decoding.""" @@ -57,15 +61,18 @@ def _make_req(self, stop_strs, output_ids, stop_str_max_len=None): req.sampling_params.stop_strs = stop_strs req.sampling_params.stop_regex_strs = [] req.sampling_params.stop_str_max_len = ( - stop_str_max_len if stop_str_max_len is not None + stop_str_max_len + if stop_str_max_len is not None else max(len(s) for s in stop_strs) ) req.sampling_params.stop_regex_max_len = 0 # Simple tokenizer mock: decode token ids to characters tok = MagicMock() + def fake_decode(ids): return "".join(chr(i) for i in ids) + tok.decode = fake_decode req.tokenizer = tok From 5f969d1fbf64d505dbd1c945e4e6cc8fc2f931ab Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 8 Jun 2026 23:12:06 -0700 Subject: [PATCH 3/3] tighten stop-str window; use real-Req unit test --- python/sglang/srt/managers/schedule_batch.py | 9 +- .../managers/test_stop_str_speculative.py | 186 ++++-------------- 2 files changed, 44 insertions(+), 151 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index d7a40ff2cd9c..c85fcacad480 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1205,12 +1205,11 @@ def tail_str(self, new_accepted_len: int = 1) -> str: self.sampling_params.stop_regex_max_len + 1, ) - # Ensure the window covers all newly accepted tokens plus the - # stop-string length so that multi-token speculative acceptance - # never pushes an early stop string out of view. + # 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(max_len_tail_str, new_accepted_len + max_len_tail_str), - len(self.output_ids), + max_len_tail_str + max(new_accepted_len - 1, 0), len(self.output_ids) ) return self.tokenizer.decode(self.output_ids[-tail_len:]) diff --git a/test/registered/unit/managers/test_stop_str_speculative.py b/test/registered/unit/managers/test_stop_str_speculative.py index f372514390f0..ce1bf4fcb2f3 100644 --- a/test/registered/unit/managers/test_stop_str_speculative.py +++ b/test/registered/unit/managers/test_stop_str_speculative.py @@ -1,165 +1,59 @@ -"""Test that stop-string checking covers all newly accepted tokens, -not just the tail window — critical for speculative decoding. - -This test uses pure mocks so it can run without sglang dependencies.""" +"""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.""" import unittest -from unittest.mock import MagicMock +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=1, suite="stage-a-test-cpu") +register_cpu_ci(est_time=5, suite="base-a-test-cpu") +STOP_ID = 1 +ID_TO_TEXT = {STOP_ID: "STOP", **{i: chr(ord("a") + i % 26) for i in range(10, 40)}} -def tail_str_new(self, new_accepted_len: int = 1) -> str: - """Fixed version of tail_str that expands window for speculative decoding.""" - if ( - len(self.sampling_params.stop_strs) == 0 - and len(self.sampling_params.stop_regex_strs) == 0 - ): - return "" +# "STOP" (index 3) sits 6 tokens back: outside the old (stop_str_max_len + 1) +# window, inside the one widened by new_accepted_len. +MIDCHUNK = [10, 11, 12, STOP_ID, 20, 21, 22, 23, 24] - max_len_tail_str = max( - self.sampling_params.stop_str_max_len + 1, - self.sampling_params.stop_regex_max_len + 1, - ) - # Ensure the window covers all newly accepted tokens plus the - # stop-string length so that multi-token speculative acceptance - # never pushes an early stop string out of view. - tail_len = min( - max(max_len_tail_str, new_accepted_len + max_len_tail_str), - len(self.output_ids), - ) - return self.tokenizer.decode(self.output_ids[-tail_len:]) +class _FakeTokenizer: + eos_token_id = -1 + additional_stop_token_ids = None + def decode(self, ids): + return "".join(ID_TO_TEXT[int(i)] for i in ids) -def tail_str_old(self) -> str: - """Original (buggy) version of tail_str for comparison.""" - if ( - len(self.sampling_params.stop_strs) == 0 - and len(self.sampling_params.stop_regex_strs) == 0 - ): - return "" - max_len_tail_str = max( - self.sampling_params.stop_str_max_len + 1, - self.sampling_params.stop_regex_max_len + 1, +def _make_req(output_ids, stop): + sp = SamplingParams(max_new_tokens=1000, stop=stop) + sp.normalize(tokenizer=None) # char-based stop_str_max_len + req = Req( + rid="t", + origin_input_text="", + origin_input_ids=array("q", [0]), + sampling_params=sp, + eos_token_ids=set(), + vocab_size=10_000, ) - - tail_len = min(max_len_tail_str, len(self.output_ids)) - return self.tokenizer.decode(self.output_ids[-tail_len:]) + req.tokenizer = _FakeTokenizer() + req.output_ids = array("q", output_ids) + return req class TestStopStrSpeculative(unittest.TestCase): - """Verify tail_str window expands to cover multi-token acceptance.""" - - def _make_req(self, stop_strs, output_ids, stop_str_max_len=None): - """Create a minimal mock that looks like a Req for tail_str.""" - req = MagicMock() - req.output_ids = output_ids - req.sampling_params.stop_strs = stop_strs - req.sampling_params.stop_regex_strs = [] - req.sampling_params.stop_str_max_len = ( - stop_str_max_len - if stop_str_max_len is not None - else max(len(s) for s in stop_strs) - ) - req.sampling_params.stop_regex_max_len = 0 - - # Simple tokenizer mock: decode token ids to characters - tok = MagicMock() - - def fake_decode(ids): - return "".join(chr(i) for i in ids) - - tok.decode = fake_decode - req.tokenizer = tok - - return req - - def test_default_window_near_old_behavior(self): - """With new_accepted_len=1 (default), window is very close to old logic. - - The new formula adds +1 token to the window when new_accepted_len=1, - which is negligible and safe for non-speculative paths. - """ - stop_strs = ["hello"] # stop_str_max_len = 5 - output_ids = list(range(65, 85)) # 20 tokens (ASCII A-T) - req = self._make_req(stop_strs, output_ids) - - # Old: tail_len = min(5+1, 20) = 6 - # New with new_accepted_len=1: tail_len = min(max(6, 1+6), 20) = 7 - old_result = tail_str_old(req) - new_result = tail_str_new(req, new_accepted_len=1) - # New window is 1 token larger (7 vs 6), which is safe - self.assertEqual(len(old_result), 6) - self.assertEqual(len(new_result), 7) - # New result contains old result as suffix - self.assertTrue(new_result.endswith(old_result)) - - def test_expanded_window_with_speculative(self): - """With new_accepted_len > max_len_tail_str, window expands.""" - stop_strs = ["hi"] # stop_str_max_len = 2 - output_ids = list(range(65, 85)) # 20 tokens - req = self._make_req(stop_strs, output_ids) - - # max_len_tail_str = max(2+1, 0+1) = 3 - # With new_accepted_len=10: tail_len = min(max(3, 10+3), 20) = 13 - result = tail_str_new(req, new_accepted_len=10) - self.assertEqual(len(result), 13) - - def test_window_clamped_to_output_len(self): - """Window never exceeds len(output_ids).""" - stop_strs = ["hi"] - output_ids = list(range(65, 70)) # only 5 tokens - req = self._make_req(stop_strs, output_ids) - - # new_accepted_len=100 would want window=103, but clamped to 5 - result = tail_str_new(req, new_accepted_len=100) - self.assertEqual(len(result), 5) - - def test_stop_str_missed_by_old_window(self): - """Old window misses stop string at the beginning of a large batch.""" - stop_strs = ["AB"] # stop_str_max_len = 2 - # "AB" at positions 0-1, then 18 filler chars = 20 total - output_ids = [65, 66] + list(range(67, 87)) - req = self._make_req(stop_strs, output_ids) - - # Old: tail_len = min(2+1, 20) = 3 → last 3 chars only - old_tail = tail_str_old(req) - self.assertNotIn("AB", old_tail, "Old window should miss the stop string") - - def test_stop_str_found_by_expanded_window(self): - """Expanded window finds stop string at the beginning of a large batch.""" - stop_strs = ["AB"] # stop_str_max_len = 2 - output_ids = [65, 66] + list(range(67, 87)) - req = self._make_req(stop_strs, output_ids) - - # New: tail_len = min(max(3, 20+3), 20) = 20 → all chars decoded - new_tail = tail_str_new(req, new_accepted_len=20) - self.assertIn("AB", new_tail, "Expanded window should find the stop string") - - def test_no_stop_strs_returns_empty(self): - """Returns empty string when no stop strings configured.""" - output_ids = list(range(65, 85)) - req = self._make_req(["x"], output_ids) # placeholder for mock setup - req.sampling_params.stop_strs = [] - - result = tail_str_new(req) - self.assertEqual(result, "") - - def test_regex_max_len_also_expands_window(self): - """When stop_regex_max_len is larger, it also benefits from expansion.""" - stop_strs = ["hi"] - output_ids = list(range(65, 85)) - req = self._make_req(stop_strs, output_ids) - req.sampling_params.stop_regex_max_len = 5 - - # max_len_tail_str = max(2+1, 5+1) = 6 - # With new_accepted_len=10: tail_len = min(max(6, 10+6), 20) = 16 - result = tail_str_new(req, new_accepted_len=10) - self.assertEqual(len(result), 16) + 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): + 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()) if __name__ == "__main__":