chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Preserve speculative finish correctness across reasoning and stop strings

Browse Source 
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
This commit is contained in:
laoyao0822
2026-06-27 20:36:52 +08:00
parent f3cf95c0f1
commit 132561a151
4 changed files with 188 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/constrained/reasoner_grammar_backend.py
View File

@@ -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)

69
python/sglang/srt/managers/schedule_batch.py
View File

@@ -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):