fix

Author: hiworldwzj

lightllm/server/router/model_infer/infer_batch.py

@@ -106,8 +106,14 @@ 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
-    def save_promptcache_kvbuffer(self):
+    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):

lightllm/server/router/req_queue/base_queue.py

@@ -14,6 +14,10 @@ def __init__(self, args, router, dp_index, dp_size) -> None:
         from lightllm.server.router.manager import RouterManager
 
         self.router: RouterManager = router
+        # 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

lightllm/utils/config_utils.py

@@ -2,7 +2,6 @@
 import os
 from functools import lru_cache
 from .envs_utils import get_env_start_args
-from transformers.configuration_utils import PretrainedConfig
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -54,7 +53,7 @@ def get_dtype(model_path: str):
 @lru_cache(maxsize=None)
 def get_fixed_kv_len():
     start_args = get_env_start_args()
-    model_cfg, _ = PretrainedConfig.get_config_dict(start_args.model_dir)
+    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: