fix

Author: hiworldwzj

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

@@ -296,7 +296,7 @@ def _get_classed_reqs(
         将请求分类返回:
         1. wait_pause_reqs 因为推理资源不够，等待被暂停的请求。
         2. paused_reqs 已经被暂停的请求，可能会被恢复。
-        3. finished_reqs 需要释放的请求
+        3. finished_reqs 需要释放的请求, 包含正常结束和aborted退出的请求。
         4. prefill_reqs 需要进行prefill操作的请求
         5. decode_reqs 需要进行decode操作的请求
         """
@@ -369,14 +369,22 @@ def _get_classed_reqs(
 
         g_infer_state_lock.release()
 
+        self._pre_handle_finished_reqs(finished_reqs=finished_reqs)
         g_infer_context.filter_reqs(finished_reqs=finished_reqs)
+
         g_infer_context.pause_reqs(wait_pause_reqs)
 
         if recover_paused:
             g_infer_context.recover_paused_reqs(paused_reqs=paused_reqs)
 
         return prefill_reqs, decode_reqs
 
+    def _pre_handle_finished_reqs(self, finished_reqs: List[InferReq]):
+        """
+        给 PD 分离模式下，prefill node 使用的继承钩子函数，用于发起 kv 传输任务。
+        """
+        pass
+
     # 一些可以复用的通用功能函数
     def _pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: bool) -> List[InferReqUpdatePack]:
         update_func_objs: List[InferReqUpdatePack] = []

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl.py

@@ -5,22 +5,21 @@
 import torch.multiprocessing as mp
 import torch.distributed as dist
 from typing import List, Tuple
-from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend
-from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context
+from lightllm.server.router.model_infer.infer_batch import InferReq
 from lightllm.server.pd_io_struct import KVMoveTask, DecodeNodeInfo
 from lightllm.utils.log_utils import init_logger
 from lightllm.common.basemodel.infer_lock import g_router_lock, g_infer_state_lock
-from lightllm.server.router.model_infer.mode_backend.continues_batch.impl import ContinuesBatchBackend
 from rpyc.utils.server import ThreadedServer
 from .prefill_task_cache import g_kv_move_task_cache
 from lightllm.utils.device_utils import kv_trans_use_p2p
 from lightllm.utils.envs_utils import get_unique_server_name
 from lightllm.utils.dist_utils import create_new_group_for_current_dp
+from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ChunkedPrefillBackend
 
 logger = init_logger(__name__)
 
 
-class ChunckedPrefillForPrefillNode(ModeBackend):
+class ChunckedPrefillForPrefillNode(ChunkedPrefillBackend):
     def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
         super().__init__()
         self.info_queue: mp.Queue = info_queue
@@ -49,36 +48,23 @@ def init_custom(self):
 
         return
 
-    def decode(self):
-        uinit_reqs, aborted_reqs, ok_finished_reqs, prefill_reqs, decode_reqs = self._get_classed_reqs(
-            g_infer_context.infer_req_ids,
-            no_decode=True,
-        )
-        assert len(decode_reqs) == 0
-
-        if aborted_reqs:
-            self._filter_reqs(aborted_reqs)
-
-        if ok_finished_reqs:
-            self.prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(ok_finished_reqs)
-            self._filter_reqs(ok_finished_reqs)
-            ok_finished_reqs.clear()
-
-        if prefill_reqs:
-            ContinuesBatchBackend.normal_prefill_reqs(
-                self, prefill_reqs=prefill_reqs, uninit_reqs=uinit_reqs, ok_finished_reqs=ok_finished_reqs
-            )
-
-        self._overlap_req_init_and_filter(uninit_reqs=uinit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
+    def _pre_handle_finished_reqs(self, finished_reqs):
+        self._prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(run_reqs=finished_reqs)
         return
 
-    def prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(self, run_reqs: List[InferReq]):
-        # 提前在radix cache中回收相关的信息，并添加引用信息
+    def _prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(self, run_reqs: List[InferReq]):
+        # 提前在radix cache中回收相关的信息，并添加引用进行锁定，方便传输进程传输kv。
         if self.is_master_in_dp:
             logger.info("prefill_req_handle_and_frozen_tokens")
+
         g_infer_state_lock.acquire()
         try:
             for req in run_reqs:
+
+                # 区分abort 和 正常结束的请求，正常结束的请求才发起kv传输任务。
+                if not req.finish_status.is_finished():
+                    continue
+
                 req: InferReq = req
                 key = req.get_input_token_ids()[0 : req.cur_kv_len]
                 key = torch.tensor(key, dtype=torch.int64, device="cpu")

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl_for_dp.py

@@ -0,0 +1,27 @@
+import torch.multiprocessing as mp
+from typing import List, Tuple
+from lightllm.server.router.model_infer.infer_batch import InferReq
+from lightllm.utils.log_utils import init_logger
+from .prefill_impl import ChunckedPrefillForPrefillNode
+from lightllm.server.router.model_infer.mode_backend.dp_backend.impl import DPChunkedPrefillBackend
+
+logger = init_logger(__name__)
+
+
+class DPChunkedForPrefillNode(DPChunkedPrefillBackend):
+    def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
+        super().__init__()
+        self.info_queue: mp.Queue = info_queue
+        self.mem_queue: mp.Queue = mem_queue
+
+    def init_custom(self):
+        ChunckedPrefillForPrefillNode.init_custom(self)
+        return
+
+    def _pre_handle_finished_reqs(self, finished_reqs):
+        self._prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(run_reqs=finished_reqs)
+        return
+
+    def _prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(self, run_reqs: List[InferReq]):
+        DPChunkedForPrefillNode._prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(self, run_reqs=run_reqs)
+        return