From 3882ef689c672588eae43c8365f1991878d0f8d5 Mon Sep 17 00:00:00 2001 From: taking-lying-flat <1615405@qq.com> Date: Mon, 14 Sep 2026 09:06:55 +0800 Subject: [PATCH] fix(scheduler): match stop strings with incremental decoding --- python/freetoken/core.py | 2 ++ python/freetoken/scheduler/scheduler.py | 27 ++++++++++++++----- python/freetoken/tokenizer/detokenize.py | 33 ++++++++++++++++-------- 3 files changed, 45 insertions(+), 17 deletions(-) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 269b94ffe..9171354e4 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -12,6 +12,7 @@ from freetoken.kvcache import BaseCacheHandle, BaseKVCachePool from freetoken.kvcache.linear_state_pool import LinearStatePool from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.tokenizer.detokenize import DecodeStatus @dataclass @@ -64,6 +65,7 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + stop_decode_status: DecodeStatus | None = None def __post_init__(self) -> None: assert self.input_ids.is_cpu diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 28520e13d..ce28665b7 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -425,16 +425,31 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: self.send_result(reply) def _match_stop_str(self, req: Req) -> str | None: - """First stop string present in this request's generated tail, else None. Decodes - only a short suffix (bounded by the longest stop string's char length, so a stop of - N chars spans at most N tokens) to keep the per-step cost small.""" + """Match stops against the same incrementally decoded text as the frontend.""" + from freetoken.tokenizer.detokenize import DecodeStatus + stop_strs = req.sampling_params.stop_strs prompt_len = req.max_device_len - req.output_len - if len(req.input_ids) <= prompt_len: + end = len(req.input_ids) + # The frontend omits a terminal EOS, including at the output limit with ignore_eos. + if ( + end > prompt_len + and not req.can_decode + and int(req.input_ids[-1]) in self.eos_token_ids + ): + end -= 1 + if end <= prompt_len: return None + if req.stop_decode_status is None: + req.stop_decode_status = DecodeStatus( + decoded_ids=[], decoded_str="", read_offset=0, surr_offset=0, sent_offset=0 + ) + state = req.stop_decode_status + state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist()) + tail = state.decode(self.tokenizer) max_chars = max(len(s) for s in stop_strs) - tail_start = max(prompt_len, len(req.input_ids) - (max_chars + 1)) - tail = self.tokenizer.decode(req.input_ids[tail_start:].tolist()) + # Bound the text used for matching, retaining token context for incomplete characters. + state.decoded_str = state.decoded_str[-max_chars:] for s in stop_strs: if s in tail: return s diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a6..26f2f4f9e 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -72,6 +72,21 @@ class DecodeStatus: surr_offset: int # length of surr ids sent_offset: int # length of sent out string + def update(self, read_str: str, surr_str: str) -> str: + new_text = read_str[len(surr_str) :] + if len(new_text) > 0 and not new_text.endswith("�"): + self.decoded_str += new_text + self.surr_offset = self.read_offset + self.read_offset = len(self.decoded_ids) + return self.decoded_str + return self.decoded_str + find_printable_text(new_text) + + def decode(self, tokenizer: PreTrainedTokenizerBase) -> str: + return self.update( + tokenizer.decode(self.decoded_ids[self.surr_offset :]), + tokenizer.decode(self.decoded_ids[self.surr_offset : self.read_offset]), + ) + class DetokenizeManager: def __init__( @@ -91,6 +106,9 @@ def discard(self, uid: int) -> None: self.decode_map.pop(uid, None) def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: + # Each message must advance its request's decode state before the next one. + if len({msg.uid for msg in msgs}) != len(msgs): + return [self.detokenize([msg])[0] for msg in msgs] read_ids: List[List[int]] = [] surr_ids: List[List[int]] = [] for msg in msgs: @@ -103,7 +121,9 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: sent_offset=0, ) s = self.decode_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): + if not ( + msg.finished and not msg.matched_stop and msg.next_token in self.eos_token_ids + ): s.decoded_ids.append(msg.next_token) read_ids.append(s.decoded_ids[s.surr_offset :]) surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset]) @@ -114,16 +134,7 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: incremental_strs: List[str] = [] for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): s = self.decode_map[msg.uid] - new_text = read_str[len(surr_str) :] - # Streaming chunk: update the decode status - if len(new_text) > 0 and not new_text.endswith("�"): - output_str = s.decoded_str + new_text - s.decoded_str = output_str - s.surr_offset = s.read_offset - s.read_offset = len(s.decoded_ids) - else: - new_text = find_printable_text(new_text) - output_str = s.decoded_str + new_text + output_str = s.update(read_str, surr_str) prev_sent = s.sent_offset if msg.finished: