fix

Author: hiworldwzj

lightllm/server/router/model_infer/infer_batch.py

@@ -53,7 +53,8 @@ def get_overlap_stream(self) -> torch.cuda.Stream:
             self.overlap_stream = torch.cuda.Stream()
         return self.overlap_stream
 
-    def add_reqs(self, requests: List[Tuple[int, int, Any, int]]):
+    def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: bool = True) -> List["InferReq"]:
+        req_objs = []
         request_ids = []
         for r in requests:
             r_id, r_index, multimodal_params, _ = r
@@ -64,12 +65,14 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]]):
                 shm_index=r_index,
                 multimodal_params=multimodal_params,
                 vocab_size=self.vocab_size,
+                init_prefix_cache=init_prefix_cache,
             )
             self.requests_mapping[r_id] = r_obj
             request_ids.append(r_id)
+            req_objs.append(r_obj)
 
         self.infer_req_ids.extend(request_ids)
-        return
+        return req_objs
 
     def free_a_req_mem(self, free_token_index: List, req: "InferReq", is_group_finished: bool):
         if self.radix_cache is None:
@@ -261,6 +264,7 @@ def __init__(
         shm_index: int,
         multimodal_params=None,
         vocab_size: int = -1,
+        init_prefix_cache: bool = True,
     ):
         self.req_id = req_id
         self.req_idx = req_idx
@@ -285,7 +289,8 @@ def __init__(
         self.mtp_gen_token_ids: List[int] = []
 
         self._init_all_state()
-        self._match_radix_cache()
+        if init_prefix_cache:
+            self._match_radix_cache()
         return
 
     def _init_all_state(self):

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

@@ -252,7 +252,6 @@ def _try_read_new_reqs(self):
                             req.infer_aborted = True
                 else:
                     self._init_reqs(reqs=cmds)
-                    self.chunked_prefill_state.need_prefill_count += 1
         return
 
     # 一些可以复用的通用功能函数

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_impl.py

@@ -3,10 +3,9 @@
 import torch.multiprocessing as mp
 import torch.distributed as dist
 import threading
-from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend
-from lightllm.server.router.model_infer.mode_backend.continues_batch.impl import ContinuesBatchBackend
+from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ChunkedPrefillBackend
 from typing import List, Tuple
-from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq
+from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq, g_infer_state_lock
 from lightllm.server.core.objs import FinishStatus
 from lightllm.utils.log_utils import init_logger
 from rpyc.utils.server import ThreadedServer
@@ -19,7 +18,7 @@
 logger = init_logger(__name__)
 
 
-class ContinuesBatchBackendForDecodeNode(ModeBackend):
+class DecodeNode(ChunkedPrefillBackend):
     def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
         super().__init__()
         self.info_queue: mp.Queue = info_queue
@@ -48,23 +47,23 @@ def init_custom(self):
 
         return
 
-    def decode(self):
-        uninit_reqs, aborted_reqs, ok_finished_reqs, prefill_reqs, decode_reqs = self._get_classed_reqs(
-            g_infer_context.infer_req_ids,
-            no_decode=False,
-        )
-        # p d 分离模式下， decode 节点不可能存在需要prefill操作的请求
-        assert len(prefill_reqs) == 0
+    def _init_reqs(self, reqs: List[Tuple]):
+        """
+        替换请求初始化操作，替换为 Decode 节点独有的一些特殊初始化流程
+        """
+        if self.dp_size_in_node != 1:
+            dp_rank_in_node = self.dp_rank_in_node
+            reqs = [req for req in reqs if req[3] == dp_rank_in_node]
 
-        self._filter_reqs(aborted_reqs)
+        g_infer_state_lock.acquire()
 
-        if decode_reqs:
-            ContinuesBatchBackend.normal_decode(
-                self, decode_reqs=decode_reqs, uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs
-            )
+        uninit_reqs = g_infer_context.add_reqs(reqs, init_prefix_cache=False)
+        # 匹配radix cache，并更新一些资源的管理。
+        self._post_init_reqs(uninit_reqs=uninit_reqs)
 
-        self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
-        return
+        g_infer_state_lock.release()
+        req_ids = [e[0] for e in reqs]
+        return req_ids
 
     def _post_init_reqs(self, uninit_reqs: List[InferReq]):
         """
@@ -85,12 +84,11 @@ def _post_init_reqs(self, uninit_reqs: List[InferReq]):
                 req_all_len = len(task.input_tokens) + task.decode_node.max_new_tokens
                 remove_count += req_all_len
                 estimated_peak_token_count += req_all_len
-                req_obj.init_all()
+                req_obj._match_radix_cache()
             else:
                 # 对于不合法的请求，直接模拟将其finished掉
-                req_obj.init_all()
-                req_obj.set_next_gen_token_id(0, 0.0)
                 req_obj.cur_output_len += 1
+                req_obj.set_next_gen_token_id(0, 0.0, 1)
                 req_obj.finish_status.set_status(FinishStatus.FINISHED_STOP)
 
                 if self.is_master_in_dp:

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_impl_for_dp.py

@@ -1,37 +1,28 @@
-import torch
 import torch.multiprocessing as mp
-import torch.distributed as dist
-from typing import List, Tuple
 from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq
 from lightllm.utils.log_utils import init_logger
-from lightllm.server.router.model_infer.mode_backend.pre import padded_prepare_prefill_inputs
-from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
-from .decode_impl import ContinuesBatchBackendForDecodeNode
+from typing import List, Tuple
 from lightllm.server.router.model_infer.mode_backend.dp_backend.impl import DPChunkedPrefillBackend
+from .decode_impl import DecodeNode
 
 logger = init_logger(__name__)
 
 
-class DPForDecodeNode(ContinuesBatchBackendForDecodeNode):
+class DPForDecodeNode(DPChunkedPrefillBackend):
     def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
-        super().__init__(info_queue, mem_queue)
-        self.enable_decode_microbatch_overlap = get_env_start_args().enable_decode_microbatch_overlap
+        super().__init__()
+        self.info_queue: mp.Queue = info_queue
+        self.mem_queue: mp.Queue = mem_queue
         return
 
-    def decode(self):
-        uninit_reqs, aborted_reqs, ok_finished_reqs, prefill_reqs, decode_reqs = self._get_classed_reqs(
-            g_infer_context.infer_req_ids
-        )
-        assert len(prefill_reqs) == 0
-
-        self._filter_reqs(aborted_reqs)
+    def init_custom(self):
+        DecodeNode.init_custom(self)
+        return
 
-        max_decode_num = self._dp_all_reduce_decode_req_num(decode_reqs=decode_reqs)
-        if max_decode_num != 0:
-            if not self.enable_decode_microbatch_overlap:
-                DPChunkedPrefillBackend.normal_decode(self, decode_reqs, max_decode_num, uninit_reqs, ok_finished_reqs)
-            else:
-                DPChunkedPrefillBackend.overlap_decode(self, decode_reqs, max_decode_num, uninit_reqs, ok_finished_reqs)
+    def _init_reqs(self, reqs: List[Tuple]):
+        DecodeNode._init_reqs(self, reqs=reqs)
+        return
 
-        self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
+    def _post_init_reqs(self, uninit_reqs: List[InferReq]):
+        DecodeNode._post_init_reqs(self, uninit_reqs=uninit_reqs)
         return