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/io_objs/__init__.py

@@ -1 +1 @@
-from .group_req import GroupReqIndexes, GroupReqObjs, AbortedReqCmd
+from .group_req import GroupReqIndexes, GroupReqObjs, AbortedReqCmd, StopStrMatchedReqCmd

lightllm/server/core/objs/io_objs/group_req.py

@@ -31,3 +31,8 @@ def to_group_req_index(self):
 @dataclass
 class AbortedReqCmd:
     req_id: int
+
+
+@dataclass
+class StopStrMatchedReqCmd:
+    req_id: int

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队列的索引位置
+        ("is_stop_str_matched", ctypes.c_bool),  # 强制停止所有推理并让客户端返回
     ]
 
     def __init__(self):
         self.data_len = 0
         self.src_index = -1
         self.special = False
         self.count_output_tokens = -1
+        self.is_stop_str_matched = 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, is_stop_str_matched: 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.is_stop_str_matched = is_stop_str_matched
         return
 
     def get(self):
@@ -40,6 +43,7 @@ def get(self):
             self.src_index,
             self.special,
             self.count_output_tokens,
+            self.is_stop_str_matched,
         )
 
 
@@ -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, is_stop_str_matched: 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, is_stop_str_matched)
 
         # 更新尾部
         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"
@@ -97,6 +100,8 @@ class Req(ctypes.Structure):
         ("mtp_accepted_token_num", ctypes.c_int),
         # mtp_step 保存一个mtp使用的常量参数，用于快速访问，不会被外部输入初始化
         ("_mtp_step", ctypes.c_int),
+        # stop_str_matched用于判断停止字符串是否匹配成功
+        ("stop_str_matched", ctypes.c_bool),
     ]
 
     def get_str(self):
@@ -150,6 +155,7 @@ def init(
         self.shm_prompt_ids.arr[0 : len(prompt_ids)] = prompt_ids
         self.mtp_accepted_token_num = 0
         self._mtp_step = get_env_start_args().mtp_step
+        self.stop_str_matched = False
 
         self.post_init()
 
@@ -207,7 +213,7 @@ def can_release(self):
         ref_count_ok = self.ref_count == 1
         can_released_mark = self.can_released_mark
 
-        if self.is_aborted and can_released_mark and ref_count_ok:
+        if (self.is_aborted or self.stop_str_matched) and can_released_mark and ref_count_ok:
             return True
 
         if self.finish_status.is_finished() and can_released_mark and ref_count_ok and self.out_tokens_queue.is_empty():

lightllm/server/detokenization/decode_req.py

@@ -36,7 +36,11 @@ def init_token_healing_prefix_str(self, token_id_to_token: Dict[int, str], token
         return
 
     def need_detoken(self):
-        if (not self.req.is_aborted) and len(self.output_ids) < self.req.candetoken_out_len:
+        if (
+            (not self.req.is_aborted)
+            and (not self.req.stop_str_matched)
+            and len(self.output_ids) < self.req.candetoken_out_len
+        ):
             return True
         return False
 
@@ -55,6 +59,9 @@ def get_decode_tokens(self):
     def can_set_release_mark(self):
         if self.req.is_aborted:
             return True
+        if self.req.stop_str_matched:
+            # 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.stop_str_matched:
+                    continue
+
                 exist_decode = True
                 new_text = decode_token(
                     self.tokenizer,
@@ -131,7 +160,18 @@ 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)
+
+                # 停止字符串匹配
+                is_stop_str_matched = False
+                if not decode_req.req.finish_status.is_stoped() and self._stop_sequences_str_matched(
+                    decode_req, self.tokenizer
+                ):
+                    decode_req.req.stop_str_matched = True
+                    is_stop_str_matched = True
+
+                decode_req.req.out_tokens_queue.push(
+                    new_text, src_index, special, count_output_tokens, is_stop_str_matched
+                )
 
             if decode_req.need_detoken():
                 exist_need_detoken = True

lightllm/server/httpserver/manager.py

@@ -661,7 +661,13 @@ 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,
+                                    is_stop_str_matched,
+                                ) = 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 +685,14 @@ async def handle_loop(self):
 
                                 req.out_tokens_queue.pop_no_ret()
 
-                                if req.finish_token_index != src_index:
+                                if not is_stop_str_matched 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 is_stop_str_matched
+                                        else req.finish_status.status
+                                    )
                                     token_list.append((req_id, text, metadata, finish_status))
                             else:
                                 break

lightllm/server/router/manager.py

@@ -15,7 +15,7 @@
 from .batch import Batch, Req
 from .model_infer.model_rpc import start_model_process, ModelRpcClient
 from .req_queue import build_req_queue
-from lightllm.server.core.objs.io_objs import GroupReqIndexes, AbortedReqCmd
+from lightllm.server.core.objs.io_objs import GroupReqIndexes, AbortedReqCmd, StopStrMatchedReqCmd
 from lightllm.server.core.objs import ShmReqManager, StartArgs
 from .dynamic_prompt.radix_cache import RadixCacheReadOnlyClient
 from .shm_reqs_io_buffer import ShmReqsIOBuffer
@@ -277,8 +277,11 @@ async def _step(self):
 
         self._filter_reqs_from_running_batch()
         aborted_reqs = self._get_aborted_reqs_from_running_batch()
+        stop_str_matched_reqs = self._get_stop_str_reqs_from_running_batch()
         if aborted_reqs:
             await self._aborted_reqs(aborted_reqs=aborted_reqs)
+        if stop_str_matched_reqs:
+            await self._stop_str_matched_reqs(stop_str_matched_reqs=stop_str_matched_reqs)
         return
 
     async def _add_batch(self, batch: Batch):
@@ -301,6 +304,15 @@ async def _aborted_reqs(self, aborted_reqs: List[Req]):
         self.shm_reqs_io_buffer.set_ready()
         return
 
+    async def _stop_str_matched_reqs(self, stop_str_matched_reqs: List[Req]):
+        cmds = [StopStrMatchedReqCmd(req_id=r.request_id) for r in stop_str_matched_reqs]
+        while not self.shm_reqs_io_buffer.is_empty():
+            await asyncio.sleep(0.02)
+
+        self.shm_reqs_io_buffer.write_obj(cmds)
+        self.shm_reqs_io_buffer.set_ready()
+        return
+
     def _add_new_batch_to_running_batch(self, new_batch: Batch):
         if self.running_batch is None:
             self.running_batch = new_batch
@@ -325,6 +337,15 @@ def _get_aborted_reqs_from_running_batch(self) -> List[Req]:
                 ans.append(req)
         return ans
 
+    def _get_stop_str_reqs_from_running_batch(self) -> List[Req]:
+        ans = []
+        if self.running_batch is None:
+            return ans
+        for req in self.running_batch.reqs:
+            if req.stop_str_matched:
+                ans.append(req)
+        return ans
+
     def _get_paused_req_num(self) -> int:
         if self.running_batch is None:
             return 0

lightllm/server/router/model_infer/infer_batch.py

@@ -319,7 +319,6 @@ def _init_all_state(self):
         g_infer_context.req_manager.req_sampling_params_manager.init_req_sampling_params(self)
 
         self.stop_sequences = self.sampling_param.shm_param.stop_sequences.to_list()
-        self.stop_sequences_str = self.sampling_param.shm_param.stop_sequences.to_string()
         # token healing mode 才被使用的管理对象
         if self.shm_req.prefix_token_ids.size != 0:
             self.prefix_token_ids = self.shm_req.prefix_token_ids.get_token_ids()
@@ -380,10 +379,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 +405,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 +483,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 +504,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)