support prefix noquant (#722)

Co-authored-by: hiworldwzj <30762946+hiworldwzj@users.noreply.github.com>

Author: helloyongyang

lightllm/server/router/model_infer/infer_batch.py

@@ -106,6 +106,27 @@ 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
 
+    def _save_promptcache_kvbuffer(self):
+        """
+        save prompt cache kv buffer
+        这个接口是用于保存非量化的缓存prompt cache资源，是定制场景使用的接口，当前代码中不会有调用。
+        其保存的 kv 会配合量化推理模式, 加载到量化推理的prompt cache中, 提升量化推理的精度。
+        like paper:
+        https://arxiv.org/abs/2403.01241
+        """
+        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,10 @@ def init_model(self, kvargs):
             else None
         )
 
+        if "prompt_cache_kv_buffer" in model_cfg:
+            assert self.use_dynamic_prompt_cache
+            self.preload_prompt_cache_kv_buffer(model_cfg)
+
         self.logger.info(f"loaded model class {self.model.__class__}")
         self.init_custom()
 
@@ -256,3 +260,25 @@ 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")
+        intact_kv_len = len(model_cfg["prompt_cache_token_ids"])
+        intact_kv_index = self.radix_cache.mem_manager.alloc(intact_kv_len)
+        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][intact_kv_index].copy_(prompt_cache_kv_buffer[i])
+        else:
+            self.radix_cache.mem_manager.kv_buffer[:, intact_kv_index].copy_(prompt_cache_kv_buffer)
+        self.radix_cache.insert(
+            torch.tensor(model_cfg["prompt_cache_token_ids"], dtype=torch.int64, device="cpu"),
+            intact_kv_index,
+        )
+        self.radix_cache.match_prefix(
+            torch.tensor(model_cfg["prompt_cache_token_ids"], dtype=torch.int64, device="cpu"), update_refs=True
+        )

lightllm/server/router/req_queue/base_queue.py

@@ -3,6 +3,7 @@
 from ..batch import Batch, Req
 from lightllm.server.core.objs import FinishStatus
 from lightllm.common.basemodel.infer_lock import g_router_lock
+from lightllm.utils.config_utils import get_fixed_kv_len
 
 
 class BaseQueue:
@@ -13,7 +14,11 @@ def __init__(self, args, router, dp_index, dp_size) -> None:
         from lightllm.server.router.manager import RouterManager
 
         self.router: RouterManager = router
-        self.max_total_tokens = args.max_total_token_num
+        # max_total_token_num - get_fixed_kv_len() 是为了减去被特定
+        # 推理模式预先占用了部分token kv 资源，这会导致整体可用的kv 资源
+        # 在极端情况下减少，在非特定模式下，get_fixed_kv_len() 返回的都是
+        # 0， 不会有任何影响。
+        self.max_total_tokens = args.max_total_token_num - get_fixed_kv_len()
         assert args.batch_max_tokens is not None
         self.batch_max_tokens = args.batch_max_tokens
         self.running_max_req_size = args.running_max_req_size  # 最大并非请求数量

lightllm/utils/config_utils.py

@@ -1,5 +1,7 @@
 import json
 import os
+from functools import lru_cache
+from .envs_utils import get_env_start_args
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -46,3 +48,13 @@ def get_dtype(model_path: str):
     except:
         logger.warning("torch_dtype not in config.json, use float16 as default")
         return "float16"
+
+
+@lru_cache(maxsize=None)
+def get_fixed_kv_len():
+    start_args = get_env_start_args()
+    model_cfg = get_config_json(start_args.model_dir)
+    if "prompt_cache_token_ids" in model_cfg:
+        return len(model_cfg["prompt_cache_token_ids"])
+    else:
+        return 0