fix(scheduler): match stop strings with incremental decoding

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

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

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: