Support preload prompt cache kv buffer

Author: helloyongyang

lightllm/server/router/model_infer/infer_batch.py

@@ -375,6 +375,10 @@ def free_self(self):
         if len(group_mapping) == 0:
             group_mapping.clear()
 
+        # # save prompt cache kv buffer
+        # print(f"prompt_cache_token_id : {list(self.radix_cache.root_node.children.values())[0].token_id_key}")
+        # torch.save(self.radix_cache.mem_manager.kv_buffer, f"prompt_cache_rank_{dist.get_rank()}.pt")
+
         if self.radix_cache is not None:
             logger.debug(
                 f"free a batch state:\n"

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

@@ -213,6 +213,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()
 
@@ -313,3 +316,21 @@ def remove_batch(self, batch_id):
         del batch
         g_infer_state_lock.release()
         return
+
+    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].copy_(prompt_cache_kv_buffer[i])
+        else:
+            self.radix_cache.mem_manager.kv_buffer.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