support prefix noquant

Author: helloyongyang

lightllm/server/router/model_infer/infer_batch.py

@@ -106,6 +106,23 @@ def free_a_req_mem(self, free_token_index: List, req: "InferReq", is_group_finis
                     self.radix_cache.dec_node_ref_counter(req.shared_kv_node)
                     req.shared_kv_node = None
 
+        # # save prompt cache kv buffer
+        # prompt_cache_token_id = list(self.radix_cache.root_node.children.values())[0].token_id_key
+        # print(f"prompt_cache_token_id : {prompt_cache_token_id}")
+        # if isinstance(self.radix_cache.mem_manager.kv_buffer, list):
+        #     kv_buffer_list = []
+        #     for i in range(len(self.radix_cache.mem_manager.kv_buffer)):
+        #         kv_buffer_list.append(self.radix_cache.mem_manager.kv_buffer[i][:len(prompt_cache_token_id)])
+        #     torch.save(
+        #         kv_buffer_list,
+        #         f"prompt_cache_rank_{dist.get_rank()}.pt"
+        #     )
+        # else:
+        #     torch.save(
+        #         self.radix_cache.mem_manager.kv_buffer[:, :len(prompt_cache_token_id)],
+        #         f"prompt_cache_rank_{dist.get_rank()}.pt"
+        #     )
+
     @torch.no_grad()
     def filter(self, finished_request_ids: List[int]):
         if len(finished_request_ids) == 0:

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

@@ -212,6 +212,9 @@ def init_model(self, kvargs):
             else None
         )
 
+        if "prompt_cache_kv_buffer" in model_cfg:
+            self.preload_prompt_cache_kv_buffer(model_cfg)
+
         self.logger.info(f"loaded model class {self.model.__class__}")
         self.init_custom()
 
@@ -256,3 +259,28 @@ def _init_reqs(self, reqs: List[Tuple], init_req_obj=True):
         g_infer_state_lock.release()
         req_ids = [e[0] for e in reqs]
         return req_ids
+
+    def preload_prompt_cache_kv_buffer(self, model_cfg):
+        self.logger.info("Preload prompt cache kv buffer.")
+        cur_rank = dist.get_rank()
+        prompt_cache_kv_buffer_path = os.path.join(
+            self.weight_dir, model_cfg["prompt_cache_kv_buffer"][f"rank_{cur_rank}"]
+        )
+        prompt_cache_kv_buffer = torch.load(prompt_cache_kv_buffer_path, weights_only=True, map_location="cpu")
+        if isinstance(self.radix_cache.mem_manager.kv_buffer, list):
+            for i in range(len(self.radix_cache.mem_manager.kv_buffer)):
+                self.radix_cache.mem_manager.kv_buffer[i][: len(model_cfg["prompt_cache_token_ids"])].copy_(
+                    prompt_cache_kv_buffer[i]
+                )
+        else:
+            self.radix_cache.mem_manager.kv_buffer[:, : len(model_cfg["prompt_cache_token_ids"])].copy_(
+                prompt_cache_kv_buffer
+            )
+        self.radix_cache.insert(
+            torch.tensor(model_cfg["prompt_cache_token_ids"], dtype=torch.int64, device="cpu"),
+            torch.tensor(range(len(model_cfg["prompt_cache_token_ids"])), dtype=torch.int32, device="cpu"),
+        )
+        self.radix_cache.mem_manager.mem_state[: len(model_cfg["prompt_cache_token_ids"])] = 1
+        self.radix_cache.match_prefix(
+            torch.tensor(model_cfg["prompt_cache_token_ids"], dtype=torch.int64, device="cpu"), update_refs=True
+        )