fix names

Author: wangzaijun

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

@@ -369,7 +369,7 @@ def _read_reqs_buffer_and_init_reqs(self):
             if init_reqs:
                 req_ids = self._init_reqs(reqs=init_reqs)
                 if self.args.enable_cpu_cache and req_ids:
-                    self._fill_cpu_cache_to_reqs(req_ids=req_ids)
+                    self._load_cpu_cache_to_reqs(req_ids=req_ids)
         return
 
     def _read_nixl_trans_io_buffer_and_update_req_status(self):
@@ -424,10 +424,10 @@ 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):
+    def _load_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_module.fill_cpu_cache_to_reqs(reqs=req_objs)
+        self.multi_level_cache_module.load_cpu_cache_to_reqs(reqs=req_objs)
         g_infer_state_lock.release()
         return
 
@@ -536,6 +536,12 @@ def _get_classed_reqs(
                         req_obj.wait_pause = True
                         wait_pause_count += 1
             else:
+                # 在 diverse mode 模式下，prefill 只会使用 master 状态的请求，slave 请求依靠后续
+                # 的推理代码中将master请求的状态复制到slave请求中去， 所以这里 slave 状态的请求，不
+                # 放入到 prefill reqs 队列中，在其他模式下，所有请求都是 master状态，所以也不受影响
+                if req_obj.is_slave_req():
+                    continue
+
                 token_num = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill)
                 if prefill_tokens + token_num > self.batch_max_tokens:
                     continue

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

@@ -38,6 +38,60 @@ def wait(self):
         if attach_shm_handle is not None:
             attach_shm_handle.wait()
 
+    def load_cpu_cache_to_reqs(self, reqs: List[InferReq]):
+        idle_token_num = g_infer_context.get_can_alloc_token_num()
+        token_page_size = self.args.cpu_cache_token_page_size
+        all_page_list = []
+        is_master_in_dp = self.backend.is_master_in_dp
+        for req in reqs:
+            page_list = req.shm_req.cpu_cache_match_page_indexes.get_all()
+            match_tokens = len(page_list) * token_page_size
+            # 更新命中的 cpu kv cache 长度.
+            if is_master_in_dp:
+                req.shm_req.cpu_prompt_cache_len = match_tokens
+
+            need_token_num = match_tokens - req.cur_kv_len
+            # 多匹配了一定数量的token 才进行复制操作，不然操作效率不高
+            if need_token_num >= 64:
+                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)
+
+                    # 计算需要加载的页面（只加载未匹配的部分）
+                    cur_kv_pages = req.cur_kv_len // token_page_size
+                    need_pages = page_list[cur_kv_pages:]  # 只取需要的页面
+                    actual_need_tokens = len(need_pages) * token_page_size
+
+                    mem_indexes = g_infer_context.req_manager.mem_manager.alloc(need_size=actual_need_tokens)
+
+                    # 将 cpu page 的内容拷贝到 gpu 页面中
+                    load_cpu_kv_to_gpu(
+                        mem_indexes=mem_indexes,
+                        gpu_kv_cache=self.backend.model.mem_manager.kv_buffer,
+                        cpu_kv_cache=self.cpu_cache_client.cpu_kv_cache_tensor,
+                        page_indexes=torch.tensor(need_pages, dtype=torch.int32, device="cpu").cuda(non_blocking=True),
+                    )
+
+                torch.cuda.current_stream().synchronize()
+
+                idle_token_num -= actual_need_tokens
+                g_infer_context.req_manager.req_to_token_indexs[
+                    req.req_idx, req.cur_kv_len : (req.cur_kv_len + actual_need_tokens)
+                ] = mem_indexes
+                req.cur_kv_len = req.cur_kv_len + actual_need_tokens
+                if self.backend.is_master_in_dp:
+                    req.shm_req.shm_cur_kv_len = req.cur_kv_len
+
+            all_page_list.extend(page_list)
+
+        dist.barrier(group=self.init_sync_group)
+
+        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
+
     def handle_finished_reqs(self, finished_reqs: List[InferReq]) -> List[InferReq]:
         """
         将满足cpu kv cache 卸载条件的请求进行处理，并返回需要真正退出的请求列表。
@@ -181,60 +235,6 @@ def update_cpu_cache_task_states(self):
                 task.req_obj.cpu_cache_task_status = InferReq._CpuCacheTaskStatus.FINISHED
         return
 
-    def fill_cpu_cache_to_reqs(self, reqs: List[InferReq]):
-        idle_token_num = g_infer_context.get_can_alloc_token_num()
-        token_page_size = self.args.cpu_cache_token_page_size
-        all_page_list = []
-        is_master_in_dp = self.backend.is_master_in_dp
-        for req in reqs:
-            page_list = req.shm_req.cpu_cache_match_page_indexes.get_all()
-            match_tokens = len(page_list) * token_page_size
-            # 更新命中的 cpu kv cache 长度.
-            if is_master_in_dp:
-                req.shm_req.cpu_prompt_cache_len = match_tokens
-
-            need_token_num = match_tokens - req.cur_kv_len
-            # 多匹配了一定数量的token 才进行复制操作，不然操作效率不高
-            if need_token_num >= 64:
-                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)
-
-                    # 计算需要加载的页面（只加载未匹配的部分）
-                    cur_kv_pages = req.cur_kv_len // token_page_size
-                    need_pages = page_list[cur_kv_pages:]  # 只取需要的页面
-                    actual_need_tokens = len(need_pages) * token_page_size
-
-                    mem_indexes = g_infer_context.req_manager.mem_manager.alloc(need_size=actual_need_tokens)
-
-                    # 将 cpu page 的内容拷贝到 gpu 页面中
-                    load_cpu_kv_to_gpu(
-                        mem_indexes=mem_indexes,
-                        gpu_kv_cache=self.backend.model.mem_manager.kv_buffer,
-                        cpu_kv_cache=self.cpu_cache_client.cpu_kv_cache_tensor,
-                        page_indexes=torch.tensor(need_pages, dtype=torch.int32, device="cpu").cuda(non_blocking=True),
-                    )
-
-                torch.cuda.current_stream().synchronize()
-
-                idle_token_num -= actual_need_tokens
-                g_infer_context.req_manager.req_to_token_indexs[
-                    req.req_idx, req.cur_kv_len : (req.cur_kv_len + actual_need_tokens)
-                ] = mem_indexes
-                req.cur_kv_len = req.cur_kv_len + actual_need_tokens
-                if self.backend.is_master_in_dp:
-                    req.shm_req.shm_cur_kv_len = req.cur_kv_len
-
-            all_page_list.extend(page_list)
-
-        dist.barrier(group=self.init_sync_group)
-
-        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: