fix(cache): reuse historical CPU tail pages for hybrid attention

lightllm/server/core/objs/req.py

@@ -152,6 +152,8 @@ class Req(ctypes.Structure):
         ("token_hash_page_len_list", TokenPageLenList),
         # 用于保存查找匹配到的可以被复用的cpu cache 页面信息。
         ("cpu_cache_match_page_indexes", CpuCachePageList),
+        # 历史碎页的实际前缀终点；0 表示沿用原始页边界。与 offload 的 hash/长度列表分开保存。
+        ("cpu_cache_match_tail_len", ctypes.c_int),
     ]
 
     def get_str(self):
@@ -188,6 +190,7 @@ def init(
         self.candetoken_out_len = 0
         self.prompt_cache_len = 0
         self.cpu_prompt_cache_len = 0
+        self.cpu_cache_match_tail_len = 0
         self.disk_prompt_cache_len = 0
         self.finish_token_index = -1
         self.can_released_mark = False
@@ -528,5 +531,4 @@ def get_decode_need_tokens(self):
         return need_tokens
 
     def get_first_router_need_tokens(self):
-
         return min(self.input_len + self.shm_cur_output_len, self.chunked_prefill_size)

lightllm/server/multi_level_kv_cache/manager.py

@@ -19,6 +19,7 @@
 from lightllm.utils.process_check import start_parent_check_thread
 from lightllm.utils.envs_utils import get_unique_server_name
 from lightllm.utils.shm_port_args import get_shm_port_args
+from lightllm.utils.config_utils import is_hybrid_att_model
 
 logger = init_logger(__name__)
 
@@ -29,6 +30,7 @@ def __init__(
         args: StartArgs,
     ):
         self.args: StartArgs = args
+        self.is_hybrid_att_model = is_hybrid_att_model(args.model_dir)
         ports = get_shm_port_args()
         context = zmq.Context(2)
         self.zmq_recv_socket = context.socket(zmq.PULL)
@@ -137,6 +139,28 @@ def _disk_cache_match(self, token_hash_list: List[int], all_pages: List[int]) ->
         self.cpu_cache_client.lock.release()
         return all_pages, len(new_page_indexes)
 
+    def _match_hybrid_att_tail(self, req: Req, pages: List[int]):
+        """在首个缺失页内回退到最长的历史碎页，命中后不再向后拼接页面。"""
+        if not self.is_hybrid_att_model:
+            return
+        page_lens = req.token_hash_page_len_list.get_all()
+        if len(pages) == len(page_lens):
+            return
+        page_start = len(pages) * self.args.cpu_cache_token_page_size
+        page_end = page_lens[len(pages)]
+        hash_size = self.args.linear_att_hash_page_size
+        hashes = req.hybrid_token_hash_list.get_all()
+        self.cpu_cache_client.lock.acquire_sleep1ms()
+        try:
+            for end in range(page_end - hash_size, page_start, -hash_size):
+                page_index, _ = self.cpu_cache_client.query_one_page(hashes[end // hash_size - 1])
+                if page_index is not None:
+                    pages.append(page_index)
+                    req.cpu_cache_match_tail_len = end
+                    return
+        finally:
+            self.cpu_cache_client.lock.release()
+
     def _handle_group_req_multi_cache_match(self, group_req_indexes: GroupReqIndexes, start_time: float):
         """
         match cpu cache and disk cache pages
@@ -203,6 +227,9 @@ def _handle_group_req_multi_cache_match(self, group_req_indexes: GroupReqIndexes
                     logger.exception(f"calculate disk prompt cache len has exception {str(e)}")
                     raise e
 
+            # 优先保留完整 CPU/disk 页的命中机会，再查首个缺失页内的历史碎页。
+            self._match_hybrid_att_tail(req, finded_page_indexes)
+
             while not self.cpu_cache_client.check_allpages_ready(finded_page_indexes):
                 time.sleep(0.01)
 

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

@@ -75,8 +75,12 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]):
                 continue
 
             page_len_list = req.shm_req.token_hash_page_len_list.get_all()
-            page_len_start_list = [0] + page_len_list
             assert len(page_list) <= len(page_len_list)
+            # 只调整加载视图，不能把后续 offload 的新尾页边界替换成历史边界。
+            page_len_list = page_len_list[: len(page_list)]
+            if page_list and req.shm_req.cpu_cache_match_tail_len:
+                page_len_list[-1] = req.shm_req.cpu_cache_match_tail_len
+            page_len_start_list = [0] + page_len_list
 
             if page_list:
                 match_tokens = page_len_list[len(page_list) - 1]
@@ -226,7 +230,6 @@ def _start_kv_cache_offload_task(
             assert len(token_hash_list) == len(page_len_list)
 
             if self.backend.is_master_in_dp:
-
                 find_index = bisect.bisect_right(page_len_list, req.cur_kv_len)
                 move_block_size = find_index