feat: move stop string matching to detokenization

Author: niushengxiao

lightllm/server/api_openai.py

@@ -539,6 +539,7 @@ async def _collect_generation_results(
                         earliest_stop_index = actual_stop_index
 
             if earliest_stop_index < len(final_text):
+                logger.info(f"removed stop sequence in tail: '{final_text[earliest_stop_index:]}'")
                 final_text = final_text[:earliest_stop_index]
 
     return {

lightllm/server/core/objs/out_token_circlequeue.py

@@ -14,15 +14,17 @@ class QueueItem(ctypes.Structure):
         ("special", ctypes.c_bool),
         ("count_output_tokens", ctypes.c_int),
         ("src_index", ctypes.c_int),  # 在源token队列的索引位置
+        ("force_stop", ctypes.c_bool),  # 强制停止所有推理并让客户端返回
     ]
 
     def __init__(self):
         self.data_len = 0
         self.src_index = -1
         self.special = False
         self.count_output_tokens = -1
+        self.force_stop = False
 
-    def set(self, token_str: str, src_index: int, special: bool, count_output_tokens: int):
+    def set(self, token_str: str, src_index: int, special: bool, count_output_tokens: int, force_stop: bool):
         str_bytes = token_str.encode("utf-8")
         assert (
             len(str_bytes) <= LIGHTLLM_TOKEN_MAX_BYTES
@@ -32,6 +34,7 @@ def set(self, token_str: str, src_index: int, special: bool, count_output_tokens
         self.src_index = src_index
         self.special = special
         self.count_output_tokens = count_output_tokens
+        self.force_stop = force_stop
         return
 
     def get(self):
@@ -40,6 +43,7 @@ def get(self):
             self.src_index,
             self.special,
             self.count_output_tokens,
+            self.force_stop,
         )
 
 
@@ -62,13 +66,13 @@ def is_empty(self):
     def is_full(self):
         return (self.tail + 1) % LIGHTLLM_OUT_TOKEN_QUEUE_SIZE == self.head
 
-    def push(self, token_str: str, src_index: int, special: bool, count_output_tokens: int):
+    def push(self, token_str: str, src_index: int, special: bool, count_output_tokens: int, force_stop: bool):
         if self.is_full():
             raise Exception("Queue is full")
 
         # 添加元素
         item: QueueItem = self.items[self.tail]
-        item.set(token_str, src_index, special, count_output_tokens)
+        item.set(token_str, src_index, special, count_output_tokens, force_stop)
 
         # 更新尾部
         self.tail = (self.tail + 1) % LIGHTLLM_OUT_TOKEN_QUEUE_SIZE
@@ -85,7 +89,7 @@ def pop(self) -> Tuple[str, int, bool, int]:
         self.head = (self.head + 1) % LIGHTLLM_OUT_TOKEN_QUEUE_SIZE
         return result
 
-    def peek(self) -> Tuple[str, int, bool, int]:
+    def peek(self) -> Tuple[str, int, bool, int, bool]:
         if self.is_empty():
             raise Exception("Queue is empty")
 

lightllm/server/core/objs/req.py

@@ -32,6 +32,9 @@ def get_status(self):
     def is_finished(self):
         return self.FINISHED_STOP <= self.status <= self.FINISHED_LENGTH
 
+    def is_stoped(self):
+        return self.status == self.FINISHED_STOP
+
     def get_finish_reason(self):
         if self.status == self.FINISHED_STOP:
             return "stop"

lightllm/server/detokenization/decode_req.py

@@ -54,7 +54,8 @@ def get_decode_tokens(self):
 
     def can_set_release_mark(self):
         if self.req.is_aborted:
-            return True
+            # httpserver那里必须先处理完请求, 这里才能释放
+            return self.req.out_tokens_queue.is_empty()
         if (
             self.req.finish_status.is_finished()
             and self.req.candetoken_out_len == len(self.output_ids)

lightllm/server/detokenization/manager.py

@@ -101,6 +101,32 @@ def handle_loop(self):
             logger.exception(f"detoken process has exception {str(e)}")
         return
 
+    def _stop_sequences_str_matched(self, decode_req, tokenizer):
+        stop_sequences_str = (
+            decode_req.req.sample_params.stop_sequences.to_string()
+            if decode_req.req.sample_params.stop_sequences
+            else []
+        )
+        if not stop_sequences_str or tokenizer is None:
+            return False
+
+        max_stop_str_len = max(len(stop_str) for stop_str in stop_sequences_str) if stop_sequences_str else 0
+        if max_stop_str_len == 0:
+            return False
+
+        output_len = len(decode_req.output_ids)
+        tail_token_len = min(decode_req.req.input_len + output_len, max_stop_str_len + 10)  # +10 for safety
+        if tail_token_len > 0:
+            tail_token_ids = decode_req.req.shm_prompt_ids.arr[
+                (decode_req.req.input_len + output_len - tail_token_len) : (decode_req.req.input_len + output_len)
+            ]
+            tail_str = tokenizer.decode(tail_token_ids, skip_special_tokens=False)
+            for stop_str in stop_sequences_str:
+                if stop_str in tail_str:
+                    logger.info(f"Found stop sequence in tail: stop_str='{stop_str}', " f"tail_str='{tail_str}'")
+                    return True
+        return False
+
     def gen_token_out(self):
         exist_need_detoken = False
         exist_decode = False
@@ -111,6 +137,9 @@ def gen_token_out(self):
                 special = new_token_id in self.all_special_ids
                 count_output_tokens = len(decode_req.output_ids)
 
+                if decode_req.req.is_aborted:
+                    continue
+
                 exist_decode = True
                 new_text = decode_token(
                     self.tokenizer,
@@ -131,7 +160,16 @@ def gen_token_out(self):
                         logger.error(
                             f"error token healing state, prefix_str {decode_req.prefix_str} new_text {new_text}"
                         )
-                decode_req.req.out_tokens_queue.push(new_text, src_index, special, count_output_tokens)
+
+                # 停止字符串匹配
+                force_stop = False
+                if not decode_req.req.finish_status.is_stoped() and self._stop_sequences_str_matched(
+                    decode_req, self.tokenizer
+                ):
+                    decode_req.req.is_aborted = True
+                    force_stop = True
+
+                decode_req.req.out_tokens_queue.push(new_text, src_index, special, count_output_tokens, force_stop)
 
             if decode_req.need_detoken():
                 exist_need_detoken = True

lightllm/server/httpserver/manager.py

@@ -661,7 +661,7 @@ async def handle_loop(self):
                         for _ in range(read_token_count):
                             if not req.out_tokens_queue.is_empty():
 
-                                text, src_index, special, count_output_tokens = req.out_tokens_queue.peek()
+                                text, src_index, special, count_output_tokens, force_stop = req.out_tokens_queue.peek()
                                 req.cumlogprob += float(req.shm_logprobs.arr[src_index])
                                 metadata = {
                                     "id": int(req.shm_prompt_ids.arr[src_index]),
@@ -679,10 +679,12 @@ async def handle_loop(self):
 
                                 req.out_tokens_queue.pop_no_ret()
 
-                                if req.finish_token_index != src_index:
+                                if not force_stop and req.finish_token_index != src_index:
                                     token_list.append((req_id, text, metadata, FinishStatus()))
                                 else:
-                                    finish_status = FinishStatus(req.finish_status.status)
+                                    finish_status = FinishStatus(
+                                        req.finish_status.FINISHED_STOP if force_stop else req.finish_status.status
+                                    )
                                     token_list.append((req_id, text, metadata, finish_status))
                             else:
                                 break

lightllm/server/router/model_infer/infer_batch.py

@@ -380,10 +380,8 @@ def update_mtp_accepted_token_num(self, accept_token_num: int):
     def get_last_gen_token(self):
         return self.shm_req.shm_prompt_ids.arr[self.shm_req.input_len + self.cur_output_len - 1]
 
-    def update_finish_status(self, eos_ids, output_len: int, tokenizer=None):
-        if self._stop_sequences_matched(output_len=output_len) or self._stop_sequences_str_matched(
-            tokenizer, output_len
-        ):
+    def update_finish_status(self, eos_ids, output_len: int):
+        if self._stop_sequences_matched(output_len=output_len):
             self.finish_status.set_status(FinishStatus.FINISHED_STOP)
         elif (
             output_len > 0
@@ -408,26 +406,6 @@ def _stop_sequences_matched(self, output_len: int):
                         return True
         return False
 
-    def _stop_sequences_str_matched(self, tokenizer, output_len):
-        if not self.stop_sequences_str or tokenizer is None:
-            return False
-
-        max_stop_str_len = max(len(stop_str) for stop_str in self.stop_sequences_str) if self.stop_sequences_str else 0
-        if max_stop_str_len == 0:
-            return False
-
-        tail_token_len = min(self.shm_req.input_len + output_len, max_stop_str_len + 10)  # +10 for safety
-        if tail_token_len > 0:
-            tail_token_ids = self.shm_req.shm_prompt_ids.arr[
-                (self.shm_req.input_len + output_len - tail_token_len) : (self.shm_req.input_len + output_len)
-            ]
-            tail_str = tokenizer.decode(tail_token_ids, skip_special_tokens=False)
-            for stop_str in self.stop_sequences_str:
-                if stop_str in tail_str:
-                    logger.info(f"Found stop sequence in tail: stop_str='{stop_str}', tail_str='{tail_str}'")
-                    return True
-        return False
-
     def prefill_need_token_num(self, is_chuncked_prefill: bool):
         if is_chuncked_prefill:
             input_token_ids = self.get_chuncked_input_token_ids()
@@ -506,7 +484,6 @@ def handle(
         eos_ids: List[int],
         extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]],
         is_master_in_dp: bool,
-        tokenizer=None,
     ):
         if self.output_len <= 0:
             return
@@ -528,7 +505,7 @@ def handle(
             return
 
         # 更新判断请求的 finished 状态
-        req_obj.update_finish_status(eos_ids=eos_ids, output_len=self.output_len, tokenizer=tokenizer)
+        req_obj.update_finish_status(eos_ids=eos_ids, output_len=self.output_len)
 
         if extra_post_req_handle_func is not None:
             extra_post_req_handle_func(req_obj, next_token_id, next_token_logprob)

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -26,12 +26,11 @@
 from lightllm.utils.dist_utils import get_dp_world_size, get_global_dp_rank, get_current_rank_in_dp
 from lightllm.utils.dist_utils import get_current_device_id, get_current_rank_in_node, get_node_world_size
 from lightllm.utils.dist_utils import get_dp_rank_in_node, create_new_group_for_current_node
-from lightllm.utils.envs_utils import get_env_start_args, enable_stop_string_match
+from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.distributed import dist_group_manager
 from lightllm.server.router.shm_reqs_io_buffer import ShmReqsIOBuffer
 from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager, OverlapEventPack
 from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel
-from lightllm.server.tokenizer import get_tokenizer
 
 
 class ModeBackend:
@@ -507,14 +506,6 @@ def _post_handle(
         extra_post_req_handle_func 用于提供在一个请求确定输出的时候，给出额外的后处理操作，主要是用于
         约束输出等模式，设置自己请求内部的状态机的状态，并添加额外的停止判定条件等。
         """
-        if enable_stop_string_match():
-            if not hasattr(self, "tokenizer"):
-                self.tokenizer = get_tokenizer(
-                    self.args.model_dir, self.args.tokenizer_mode, trust_remote_code=self.args.trust_remote_code
-                )
-        else:
-            self.tokenizer = None
-
         for req_obj, next_token_id, next_token_logprob, pack in zip(
             run_reqs, next_token_ids, next_token_logprobs, run_reqs_update_packs
         ):
@@ -526,7 +517,6 @@ def _post_handle(
                 eos_ids=self.eos_id,
                 extra_post_req_handle_func=extra_post_req_handle_func,
                 is_master_in_dp=self.is_master_in_dp,
-                tokenizer=self.tokenizer,
             )
 
         g_infer_context.req_manager.req_sampling_params_manager.update_reqs_token_counter(

lightllm/utils/envs_utils.py

@@ -68,11 +68,6 @@ def get_lightllm_gunicorn_keep_alive():
     return int(os.getenv("LIGHTLMM_GUNICORN_KEEP_ALIVE", 10))
 
 
-@lru_cache(maxsize=None)
-def enable_stop_string_match():
-    return os.getenv("ENABLE_STOP_STRING_MATCH", "False").upper() in ["ON", "TRUE", "1"]
-
-
 @lru_cache(maxsize=None)
 def get_lightllm_websocket_max_message_size():
     """