minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Fix spec v2 stop output boundary (#25980)

Browse Source 
Co-authored-by: gss <2783977641@qq.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
gq112
2026-06-09 00:32:21 -07:00
committed by GitHub
co-authored by gss hnyls2002 Liangsheng Yin
parent d981b7b9c4
commit 2218622f50
4 changed files with 207 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/detokenizer_manager.py
+7 -2
View File
@@ -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
python/sglang/srt/managers/schedule_batch.py
+49 -12
View File
@@ -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
test/registered/unit/managers/test_stop_str_speculative.py
+83 -13
View File
@@ -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()
test/registered/unit/managers/test_trim_matched_stop.py
+68
View File
@@ -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()