diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index b5a5bf685..12b429782 100644 --- a/lightllm/server/core/objs/req.py +++ b/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) diff --git a/lightllm/server/multi_level_kv_cache/manager.py b/lightllm/server/multi_level_kv_cache/manager.py index ef5b7369c..a442bf74e 100644 --- a/lightllm/server/multi_level_kv_cache/manager.py +++ b/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) diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 2a366f025..13447189e 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/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