fix

Author: liujiacheng

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

@@ -157,6 +157,7 @@ def init_model(self, kvargs):
 
         self.logger.info(f"loaded model class {self.model.__class__}")
         g_infer_context.register(
+            backend=self,
             req_manager=self.model.req_manager,
             radix_cache=self.radix_cache,
             shm_req_manager=self.shm_req_manager,
@@ -326,7 +327,10 @@ def _read_reqs_buffer_and_init_reqs(self):
                         req: InferReq = g_infer_context.requests_mapping[obj.req_id]
                         req.infer_aborted = True
             else:
-                self._init_reqs(reqs=cmds)
+                req_ids = self._init_reqs(reqs=cmds)
+                if self.args.enable_cpu_cache:
+                    self._fill_cpu_cache_to_reqs(req_ids=req_ids)
+
         return
 
     # 一些可以复用的通用功能函数
@@ -348,6 +352,13 @@ def _init_reqs(self, reqs: List[Tuple]):
         req_ids = [e[0] for e in reqs]
         return req_ids
 
+    def _fill_cpu_cache_to_reqs(self, req_ids):
+        req_objs: List[InferReq] = [g_infer_context.requests_mapping[req_id] for req_id in req_ids]
+        g_infer_state_lock.acquire()
+        self.multi_level_cache_manager.fill_cpu_cache_to_reqs(reqs=req_objs)
+        g_infer_state_lock.release()
+        return
+
     # 一些可以复用的通用功能函数
     def _get_classed_reqs(
         self,
@@ -374,6 +385,8 @@ def _get_classed_reqs(
         4. prefill_reqs 需要进行prefill操作的请求
         5. decode_reqs 需要进行decode操作的请求
         """
+        if self.args.enable_cpu_cache:
+            self.multi_level_cache_manager.update_kv_cache_offload_task_states()
 
         if req_ids is None:
             req_ids = g_infer_context.infer_req_ids
@@ -486,7 +499,7 @@ def _cpu_kv_cache_task_handle(self, finished_reqs: List[InferReq]) -> List[Infer
                     else:
                         # 将请求的 kv cache 卸载到 cpu cache 中
                         multi_level_cache_manager = self.multi_level_cache_manager
-                        trans_task = multi_level_cache_manager.req_to_cpu_cache_task(
+                        trans_task = multi_level_cache_manager.start_kv_cache_offload_task(
                             req=req, cpu_kv_cache_stream=g_infer_context.get_cpu_kv_cache_stream()
                         )
                         if trans_task is not None:

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

@@ -9,6 +9,7 @@
 from ..infer_batch import InferReq
 from lightllm.utils.dist_utils import create_new_group_for_current_dp
 from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu
+from lightllm.server.router.model_infer.infer_batch import g_infer_context
 
 
 class MultiLevelCacheManager(object):
@@ -20,11 +21,14 @@ def __init__(self, backend):
         self.gloo_group = create_new_group_for_current_dp("gloo")
         self.filter_group = create_new_group_for_current_dp("gloo")
         self.sync_group = create_new_group_for_current_dp("nccl")
+        self.init_sync_group = create_new_group_for_current_dp("nccl")
 
         self.cpu_cache_handle_queue = deque()
         self.cpu_cache_client = CpuKvCacheClient(init_shm_data=False)
 
-    def req_to_cpu_cache_task(self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream) -> Optional["TransTask"]:
+    def start_kv_cache_offload_task(
+        self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream
+    ) -> Optional["TransTask"]:
         with torch.cuda.stream(cpu_kv_cache_stream):
             all_token_hash_list = req.shm_req.token_hash_list.get_all()
             block_size = req.cur_kv_len // self.args.cpu_cache_token_chuncked_size
@@ -79,7 +83,7 @@ def req_to_cpu_cache_task(self, req: InferReq, cpu_kv_cache_stream: torch.cuda.S
 
         return trans_task
 
-    def handle_task_queue(self):
+    def update_kv_cache_offload_task_states(self):
         if self.backend.is_master_in_dp:
             trans_ok_reqs = []
             while len(self.cpu_cache_handle_queue) != 0:
@@ -110,6 +114,37 @@ def handle_task_queue(self):
                 req.req_obj.cpu_cache_task_finished = True
         return
 
+    def fill_cpu_cache_to_reqs(self, reqs: List[InferReq]):
+        idle_token_num = g_infer_context.get_can_alloc_token_num()
+        token_chuncked_size = self.args.cpu_cache_token_chuncked_size
+        all_page_list = []
+        for req in reqs:
+            page_list = req.shm_req.cpu_cache_match_page_indexes.get_all()
+            match_tokens = len(page_list) * token_chuncked_size
+            need_token_num = match_tokens - req.cur_kv_len
+            # 多匹配了一定数量的token 才进行复制操作，不然操作效率不高
+            if need_token_num > 256:
+                if need_token_num <= idle_token_num:
+                    if self.backend.radix_cache is not None:
+                        g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(need_token_num=need_token_num)
+
+                mem_indexes = g_infer_context.req_manager.mem_manager.alloc(need_size=need_token_num)
+                idle_token_num -= need_token_num
+                g_infer_context.req_manager.req_to_token_indexs[
+                    req.req_idx, req.cur_kv_len : (req.cur_kv_len + need_token_num)
+                ] = mem_indexes
+                req.cur_kv_len = req.cur_kv_len + need_token_num
+                if self.backend.is_master_in_dp:
+                    req.shm_req.shm_cur_kv_len = req.cur_kv_len
+
+            all_page_list.extend(page_list)
+
+        if self.backend.is_master_in_dp:
+            self.cpu_cache_client.lock.acquire_sleep1ms()
+            self.cpu_cache_client.deref_pages(page_list=all_page_list)
+            self.cpu_cache_client.lock.release()
+        return
+
 
 @dataclasses.dataclass
 class TransTask: