Cpu KV Cache feature (#997)

Co-authored-by: wangzaijun <[email protected]>
Co-authored-by: niushengxiao <[email protected]>

Author: 3 people

lightllm/common/basemodel/basemodel.py

@@ -61,6 +61,8 @@ def __init__(self, kvargs):
         self.finetune_config = kvargs.get("finetune_config", None)
         self.max_req_num = kvargs.get("max_req_num", 1000)
         self.max_seq_length = kvargs.get("max_seq_length", 1024 * 5)
+        # 用于等待外围的一些模块的初始化完成（如 CPU KV Cache 注册完成）
+        self.wait_events = kvargs.get("wait_events", [])
         # is_token_healing 和 return_all_prompt_logics 是有排斥关系的两个模式，只能单独有一个生效
         # 主要是在prefill阶段返回多少个token的用于后续处理相关。
         self.is_token_healing = kvargs.get("is_token_healing", False)
@@ -110,12 +112,19 @@ def __init__(self, kvargs):
         self._init_inferstate_cls()
         self._autotune_warmup()
         self._init_padded_req()
+        # wait必须在init cudagraph 之前，避免错误捕获
+        self._wait_other_modules_ready()
         self._init_cudagraph()
         self._check_max_len_infer()
         torch.cuda.empty_cache()
         set_model_init_status(True)
         return
 
+    def _wait_other_modules_ready(self):
+        for event in self.wait_events:
+            event.wait()
+        return
+
     def _init_config(self):
         with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
             self.config = json.load(json_file)
@@ -352,8 +361,13 @@ def _prefill(
             alloc_mem_index=infer_state.mem_index,
             max_q_seq_len=infer_state.max_q_seq_len,
         )
+        prefill_mem_indexes_ready_event = torch.cuda.Event()
+        prefill_mem_indexes_ready_event.record()
+
         infer_state.init_some_extra_state(self, model_input.input_ids)
-        return self._context_forward(model_input.input_ids, infer_state)
+        model_output = self._context_forward(model_input.input_ids, infer_state)
+        model_output.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event
+        return model_output
 
     def _decode(
         self,
@@ -505,13 +519,18 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
         )
         infer_state1.init_some_extra_state(self, input_ids1)
 
+        prefill_mem_indexes_ready_event = torch.cuda.Event()
+        prefill_mem_indexes_ready_event.record()
+
         model_output0, model_output1 = self._overlap_tpsp_context_forward(
             input_ids0, infer_state0, input_ids1=input_ids1, infer_state1=infer_state1
         )
 
         # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候
         # 该调用没有实际意义
         dist_group_manager.clear_deepep_buffer()
+        model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event
+        model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event
         return model_output0, model_output1
 
     @torch.no_grad()

lightllm/common/basemodel/batch_objs.py

@@ -58,6 +58,8 @@ def to_cuda(self):
 class ModelOutput:
     # 通用变量
     logits: torch.Tensor
+    # 用于判断 mem_indexes 是否成功写入 req manager 中的事件对象。
+    prefill_mem_indexes_ready_event: torch.Event = None
 
     # 专有变量，用于一些特殊的模型，特殊的模式下, 传递一些特殊
     # 的输出变量。只在特殊的模型模式下才会具体使用和生效。