refactor hicache

Author: yujinbiao

lightllm/server/router/dynamic_prompt/hiradix_cache.py

@@ -1,117 +1,128 @@
 import torch
+import time
+import tempfile
+import numpy as np
 import torch.distributed as dist
+from os.path import join
 from .radix_cache import RadixCache, TreeNode, match
 from typing import Tuple, Dict, Set, List
 from lightllm.common.mem_manager import MemoryManager
 from lightllm.utils.log_utils import init_logger
 from threading import Lock
 from enum import Enum
+from .shared_arr import SharedArray
 from kvcache.python.jit import PyLocalCacheService
-import time
 
 logger = init_logger(__name__)
 
+def wait_until_ready(task, timeout=10.0, check_interval=0.01):
+    start_time = time.time()
+    while not task.ready():
+        time.sleep(check_interval)
+        if time.time() - start_time > timeout:
+            logger.error("Current kv cache task not ready in time")
+            return False
+    return True
 
-class HiRadixCache(RadixCache):
-    def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager, max_seq_length):
-        super().__init__(unique_name, total_token_num, rank_in_node, mem_manager)
-        logger.info("Initializing HiRadixCache")
-        self.rank_in_node = rank_in_node
-        try:
-            # TODO: determine by model type && dp, tp
-            store_once = True  # Deepseek -> True, Llama -> False
-            self.do_store = store_once and self.rank_in_node == 0
-            self.is_hi_radix_cache = True
-            all_buffers = self.mem_manager.kv_buffer
-            all_buffers = all_buffers.view(all_buffers.shape[0], all_buffers.shape[1], -1)
-            self.py_cache_service = (
-                PyLocalCacheService(
-                    file="cache/cache_file",
-                    storage_size=128 * (1024 ** 3),
-                    num_shard=32,
-                    kvcache_tensor=all_buffers,
-                    num_worker=32,
-                )
-                if self.do_store
-                else None
-            )
-            self.working_tasks = {}
-        except Exception as e:
-            logger.error(f"error alloc hi cache buffer {e}, fallback to normal radix cache")
-            self.hi_cache_kv_buffer = None
-            self.is_hi_radix_cache = False
+class LocalCacheManager:
 
-    def insert_disk(self, req_id, key, value):
-        if not self.do_store:
-            return
-        if req_id in self.working_tasks:
-            self.abort_req_store_task(req_id)
-        self.working_tasks[req_id] = self.py_cache_service.create(tokens=key, kv_page_indexer=value, mode="w")
-        logger.info(f"Created store task for req {req_id}.")
+    def __init__(self, unique_name: str, rank_in_node: int, mem_manager):
+        tmp_dir = tempfile.mkdtemp(prefix=f"cache_{unique_name}_{rank_in_node}")
+        self.cache_file = join(tmp_dir, "cache_file")
+        all_buffers = mem_manager.kv_buffer
+        all_buffers = all_buffers.view(all_buffers.shape[0], all_buffers.shape[1], -1)
 
-    def abort_req_store_task(self, req_id):
-        if not self.do_store or req_id not in self.working_tasks:
-            return
-        if self.working_tasks[req_id].ready():
-            logger.info(f"Calling abort for req {req_id}, but is finished.")
-            return
-        logger.info(f"Aborting req {req_id} unfinished.")
-        self.py_cache_service.az5(self.working_tasks[req_id])
+        self.py_cache_service = PyLocalCacheService(
+            file=self.cache_file,
+            storage_size=128 * (1024 ** 3),  # 128GB
+            num_shard=32,
+            kvcache_tensor=all_buffers,
+            num_worker=8
+        )
 
-    def match_prefix(self, key, update_refs=False):
-        assert len(key) != 0
-        ans_value_list = []
-        pull_hi_cache_tensor = torch.tensor([0], dtype=torch.int64).cuda(self.rank_in_node)
-        if self.do_store:
-            tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=False)
-            max_len = self._query_hi_cache(key)  # x64
-            logger.info(f"Matched {sum(len(s) for s in ans_value_list)} from gpu and {max_len} from disk.")
-            pull_hi_cache_tensor[0] = max_len if (max_len > sum(len(s) for s in ans_value_list)) else 0
-        dist.broadcast(pull_hi_cache_tensor, src=0)
-        pull_hi_cache = False
+    def insert(self, tokens, kv_page_indexer, start_pos=0):
+        t = self.py_cache_service.create(
+                tokens=tokens, 
+                kv_page_indexer=kv_page_indexer, 
+                mode="w",
+                start_pos=start_pos)
+        res = wait_until_ready(t)
+        if not res:
+            self.py_cache_service.az5(t)
 
-        if pull_hi_cache_tensor[0] == 0:
-            ans_value_list = []
-            tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
-        elif pull_hi_cache_tensor[0] > 0:
-            pull_hi_cache = True
-            max_len = pull_hi_cache_tensor[0]
-            try:
-                self.free_radix_cache_to_get_enough_token(max_len)
-            except:
-                logger.info(f"Unable to free on rank {self.rank_in_node}")
-                pull_hi_cache_tensor[0] = 0
-                pull_hi_cache = False
-                ans_value_list = []
-                tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
-        if pull_hi_cache:
-            buffers = self.mem_manager.alloc(max_len)
-            if self.do_store:
-                read_task = self.py_cache_service.create(tokens=key[:max_len], kv_page_indexer=buffers, mode="r")
-                while not read_task.ready():
-                    time.sleep(0.05)
-            dist.broadcast(self.mem_manager.get_index_kv_buffer(buffers)["kv_buffer"], src=0)
-            logger.info(f"HiCache pulled one cache with len = {max_len}")
-            self._insert_helper(self.root_node, key, buffers)
-            ans_value_list = []
-            tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
-        if tree_node != self.root_node:
-            if len(ans_value_list) != 0:
-                value = torch.concat(ans_value_list)
-            else:
-                assert False, "can not run to here"
-            return tree_node, len(value), value
-        else:
-            self.dec_node_ref_counter(self.root_node)
-            return None, 0, None
+    def read(self, tokens, kv_page_indexer, start_pos=0):
+        t = self.py_cache_service.create(
+                tokens=tokens, 
+                kv_page_indexer=kv_page_indexer, 
+                mode="r",
+                start_pos=start_pos)
+        res = wait_until_ready(t)
+        return res
 
-    def _query_hi_cache(self, key) -> bool:
-        query_result = self.py_cache_service.query(key)
-        # query_result is a list of bool, find out the max len true continuous from start
+    def query(self, tokens):
+        query_result = self.py_cache_service.query(tokens)
         max_len = 0
         for result in query_result:
             if result:
                 max_len += 1
             else:
                 break
-        return max_len * self.py_cache_service.tokens_per_block
+        return max_len * self.block_size
+
+    @property
+    def block_size(self,):
+        return self.py_cache_service.tokens_per_block
+
+class HiRadixCache(RadixCache):
+    def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager):
+        super().__init__(unique_name, total_token_num, rank_in_node, mem_manager)
+        self.rank_in_node = rank_in_node
+        self.local_cache_manager = LocalCacheManager(
+            unique_name,
+            rank_in_node,
+            mem_manager,
+        )
+        self.is_hi_radix_cache = True
+        self.disk_cache_match_count = SharedArray(f"{unique_name}_disk_cache_match_count_{rank_in_node}", (1,), dtype=np.int64)
+        self.disk_cache_match_count.arr[0] = 0
+        self.total_match_count = SharedArray(f"{unique_name}_total_match_count_{rank_in_node}", (1,), dtype=np.int64)
+        self.total_match_count.arr[0] = 0
+        self.disk_cache_match_ratio = SharedArray(f"{unique_name}_disk_cache_match_ratio_{rank_in_node}", (1,), dtype=np.float32)
+        self.disk_cache_match_ratio.arr[0] = 0.0
+        logger.info(f"Initializing HiRadixCache {rank_in_node}")
+
+    def insert(self, key, value=None):
+        share_len = super().insert(key, value)
+        if share_len == 0:
+            return 0
+        self.local_cache_manager.insert(key, value)
+        return share_len
+
+    def match_prefix(self, key, update_refs=False):
+        assert len(key) != 0
+        self.total_match_count.arr[0] += 1
+        ans_value_list = []
+        ans_value = None
+        tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=False)
+        if tree_node.node_prefix_total_len != 0:
+            ans_value = torch.concat(ans_value_list)
+        max_len = 0
+        if tree_node.node_prefix_total_len < len(key):
+            max_len = self.local_cache_manager.query(key)
+        if max_len > tree_node.node_prefix_total_len:
+            pull_len = max_len - tree_node.node_prefix_total_len
+            self.disk_cache_match_count.arr[0] += 1
+            self.disk_cache_match_ratio.arr[0] = self.disk_cache_match_count.arr[0] / self.total_match_count.arr[0]
+            self.free_radix_cache_to_get_enough_token(pull_len)
+            buffers = self.mem_manager.alloc(pull_len)
+            start_pos = 0
+            if ans_value is not None:
+                buffers = torch.concat([ans_value, buffers])
+                start_pos = (tree_node.node_prefix_total_len - 1) // self.local_cache_manager.block_size * self.local_cache_manager.block_size
+            logger.debug(f"HiCache current match ratio {self.disk_cache_match_ratio.arr[0]}, pulled cache len {pull_len} from disk")
+            res = self.local_cache_manager.read(tokens=key[:max_len], kv_page_indexer=buffers, start_pos=start_pos)
+            if res:
+                super().insert(key[:max_len], buffers)
+            else:
+                self.mem_manager.free(buffers[tree_node.node_prefix_total_len:])
+        return super().match_prefix(key, update_refs=update_refs)

lightllm/server/router/model_infer/infer_batch.py

@@ -109,8 +109,6 @@ 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
 
-            if self.radix_cache.is_hi_radix_cache:
-                self.radix_cache.abort_req_store_task(req.req_id)
 
     def _save_promptcache_kvbuffer(self):
         """

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

@@ -121,8 +121,7 @@ def init_model(self, kvargs):
                 get_unique_server_name(),
                 self.model.mem_manager.size,
                 self.rank_in_node,
-                mem_manager=self.model.mem_manager,
-                max_seq_length=kvargs.get("max_seq_length", 1024 * 5),
+                mem_manager=self.model.mem_manager
             )
             if self.use_dynamic_prompt_cache and self.use_hi_dynamic_prompt_cache
             else RadixCache(
@@ -347,23 +346,6 @@ def _overlap_req_init_and_filter(
 
         return
 
-    def _overlap_store_prefill_reqs(self, run_reqs: List[InferReq]):
-        if run_reqs:
-            with torch.cuda.stream(g_infer_context.get_overlap_stream()):
-                if self.use_hi_dynamic_prompt_cache and self.radix_cache is not None:
-                    for req in run_reqs:
-                        if req.cur_output_len > 1:
-                            continue
-                        key = torch.tensor(
-                            req.get_input_token_ids()[0 : req.cur_kv_len], dtype=torch.int64, device="cpu"
-                        )
-                        value = self.model.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].detach().cpu()
-                        self.radix_cache.insert_disk(req.req_id, key, value)
-
-            torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream())
-
-        return
-
     # 一些可以复用的通用功能函数
     def _post_init_reqs(self, uninit_reqs: List[InferReq]):
         """

lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py

@@ -43,7 +43,6 @@ def decode(self):
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )
-            self._overlap_store_prefill_reqs(run_reqs=run_reqs)
             next_token_ids, next_token_probs = sample(logits, run_reqs, self.eos_id)
             next_token_ids = next_token_ids.detach().cpu().numpy()
             next_token_logprobs = torch.log(next_token_probs).detach().cpu().numpy()