Support disk radix cache

lightllm/server/api_cli.py

@@ -205,6 +205,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument("--disable_dynamic_prompt_cache", action="store_true", help="disable dynamic prompt cache")
 
     parser.add_argument("--chunked_prefill_size", type=int, default=4096, help="chunked prefill size")
+    parser.add_argument("--use_hi_dynamic_prompt_cache", action="store_true", help="enable hierachy prompt cache")
     parser.add_argument("--disable_chunked_prefill", action="store_true", help="whether to disable chunked prefill")
     parser.add_argument("--diverse_mode", action="store_true", help="diversity generation mode")
     parser.add_argument("--token_healing_mode", action="store_true", help="code model infer mode")
@@ -311,7 +312,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         "--visual_infer_batch_size", type=int, default=1, help="number of images to process in each inference batch"
     )
     parser.add_argument(
-        "--visual_gpu_ids", nargs="+", type=int, default=None, help="List of GPU IDs to use, e.g., 0 1 2"
+        "--visual_gpu_ids", nargs="+", type=int, default=[0, 1, 2, 3, 4, 5, 6, 7], help="List of GPU IDs to use, e.g., 0 1 2"
     )
     parser.add_argument("--visual_tp", type=int, default=1, help="number of tensort parallel instances for ViT")
     parser.add_argument("--visual_dp", type=int, default=1, help="number of data parallel instances for ViT")

lightllm/server/api_start.py

@@ -158,6 +158,10 @@ def normal_or_p_d_start(args):
             args.batch_max_tokens >= args.chunked_prefill_size
         ), "chunked prefill mode, batch_max_tokens must >= chunked_prefill_size"
 
+    # if use_hi_dynamic_prompt_cache, then use_dynamic_prompt_cache must be True
+    if args.use_hi_dynamic_prompt_cache:
+        assert not args.disable_dynamic_prompt_cache, "use_hi_dynamic_prompt_cache must be used with use_dynamic_prompt_cache"
+
     # help to manage data stored on Ceph
     if "s3://" in args.model_dir:
         from lightllm.utils.petrel_helper import s3_model_prepare

lightllm/server/core/objs/start_args_type.py

@@ -45,6 +45,7 @@ class StartArgs:
     router_max_wait_tokens: int = field(default=6)
     disable_aggressive_schedule: bool = field(default=False)
     disable_dynamic_prompt_cache: bool = field(default=False)
+    use_hi_dynamic_prompt_cache: bool = field(default=False)
     chunked_prefill_size: int = field(default=8192)
     disable_chunked_prefill: bool = field(default=False)
     diverse_mode: bool = field(default=False)

lightllm/server/router/dynamic_prompt/hiradix_cache.py

@@ -0,0 +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
+
+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 LocalCacheManager:
+
+    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)
+
+        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 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)
+
+    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(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.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/dynamic_prompt/radix_cache.py

@@ -116,6 +116,8 @@ def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager: Memo
         )
         self.tree_total_tokens_num.arr[0] = 0
 
+        self.is_hi_radix_cache = False
+
     def insert(self, key, value=None):
         if value is None:
             value = key

lightllm/server/router/manager.py

@@ -167,6 +167,7 @@ async def wait_to_model_ready(self):
             "return_all_prompt_logprobs": self.args.return_all_prompt_logprobs,
             "use_reward_model": self.args.use_reward_model,
             "disable_dynamic_prompt_cache": self.args.disable_dynamic_prompt_cache,
+            "use_hi_dynamic_prompt_cache": self.args.use_hi_dynamic_prompt_cache,
             "data_type": self.args.data_type,
             "eos_id": self.eos_id,
             "diverse_mode": self.args.diverse_mode,

lightllm/server/router/model_infer/infer_batch.py

@@ -110,6 +110,7 @@ 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

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

@@ -12,6 +12,7 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.models import get_model
 from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache
+from lightllm.server.router.dynamic_prompt.hiradix_cache import HiRadixCache
 from lightllm.server.router.model_infer.infer_batch import InferReq, InferSamplingParams
 from lightllm.server.router.token_load import TokenLoad
 from lightllm.common.basemodel.infer_lock import g_infer_state_lock, InferStateLock
@@ -53,6 +54,7 @@ def init_model(self, kvargs):
         self.chunked_prefill_size = self.args.chunked_prefill_size
         self.return_all_prompt_logprobs = self.args.return_all_prompt_logprobs
         self.use_dynamic_prompt_cache = not self.args.disable_dynamic_prompt_cache
+        self.use_hi_dynamic_prompt_cache = self.args.use_hi_dynamic_prompt_cache
         self.eos_id: List[int] = kvargs.get("eos_id", [2])
         self.disable_cudagraph = self.args.disable_cudagraph
 
@@ -118,7 +120,14 @@ def init_model(self, kvargs):
         self.model: TpPartBaseModel = self.model  # for easy typing
         set_random_seed(2147483647)
         self.radix_cache = (
-            RadixCache(
+            HiRadixCache(
+                get_unique_server_name(),
+                self.model.mem_manager.size,
+                self.rank_in_node,
+                mem_manager=self.model.mem_manager
+            )
+            if self.use_dynamic_prompt_cache and self.use_hi_dynamic_prompt_cache
+            else RadixCache(
                 get_unique_server_name(),
                 self.model.mem_manager.size,
                 self.rank_in_node,

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

@@ -59,6 +59,7 @@ def decode(self):
                     prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
                 )
                 model_output = self.model.forward(model_input)
+                self.store_hicache_after_prefill(run_reqs)
                 self._overlap_req_init_and_filter(
                     uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
                 )

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

@@ -35,6 +35,7 @@ def decode(self):
             )
             model_output = self.model.forward(model_input)
             logits = model_output.logits
+            self.store_hicache_after_prefill(run_reqs)
 
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True

test/server/test_hicache.py

@@ -0,0 +1,155 @@
+# test_hicache.py
+import torch
+import time
+import random
+from threading import Thread, Event
+from queue import Queue
+from lightllm.server.router.dynamic_prompt.cache_controller import (
+    HiCacheController,
+    CacheNode,
+    BLOCK_SIZE,
+    HiHostService,
+    HiHostTask,
+)
+
+
+class MockMemoryManager:
+    """模拟内存管理器，仅返回连续的索引值"""
+
+    def __init__(self):
+        self.current_idx = 0
+        self.kvcache_store = {}
+
+    def alloc(self, size):
+        indices = list(range(self.current_idx, self.current_idx + size))
+        self.current_idx += size
+        self.store(indices, torch.tensor([[random.randint(0, 0xFFFF) for __ in range(512)] for _ in range(size)]))
+        return indices
+
+    def load_index_kv_buffer(self, index, load_tensor_dict):
+        self.kvcache_store[index] = load_tensor_dict["kv_buffer"]
+
+    def get_index_kv_buffer(self, index):
+        return {"kv_buffer": self.kvcache_store[index]}
+
+    def to_kvcache(self, indices):
+        assert all(
+            [idx in self.kvcache_store for idx in indices]
+        ), f"Not all of {indices} are not found in kvcache_store"
+        return torch.tensor([self.kvcache_store[idx].tolist() for idx in indices])
+
+    def store(self, indices, value):
+        print(f"[TEST:MemManager] Storing {value.shape} at {indices}")
+        for idx, value_dim in zip(indices, range(value.shape[0])):
+            self.kvcache_store[idx] = value[value_dim]
+            print(f"[TEST:MemManager] Stored {value[value_dim].shape} at {idx}")
+        return indices
+
+    def free(self, indices):
+        print(f"[TEST:MemManager] Freeing {indices}")
+        for idx in indices:
+            del self.kvcache_store[idx]
+
+
+def setup():
+    mem_manager = MockMemoryManager()
+    service = HiHostService()
+    hicache = HiCacheController(mem_manager)
+    hicache.service = service  # 注入模拟服务
+
+    indices = mem_manager.alloc(5)
+    print(mem_manager.to_kvcache(indices))
+
+    # 预先计算单token大小
+    dummy_indices = mem_manager.alloc(1)
+    kvcache = mem_manager.to_kvcache(dummy_indices[:1])
+    token_size = kvcache.nelement() * kvcache.element_size()
+    print(f"[TEST] Single token KV cache size: {token_size} bytes, Block size: {BLOCK_SIZE}")
+
+    return mem_manager, service, hicache, token_size
+
+
+def test_basic_write_read(mem_manager, hicache, token_size):
+    # 计算每个块可容纳的token数量
+    tokens_per_block = BLOCK_SIZE // token_size
+    print(f"[TEST] Each block can hold {tokens_per_block} tokens")
+
+    # 生成测试数据：刚好占满一个块
+    token_ids = list(range(tokens_per_block))
+    indices = mem_manager.alloc(len(token_ids))
+    kvcache = mem_manager.to_kvcache(indices)
+    print(f"[TEST] Generated KV cache with shape: {kvcache.shape}, type: {kvcache.dtype}")
+
+    # 写入缓存
+    hicache.write(torch.tensor(token_ids), torch.tensor(indices))
+    time.sleep(2)
+
+    # 等待任务完成
+    hicache.service.wait_till_all_finished()
+
+    mem_manager.free(indices)
+
+    # 读取验证
+    result = hicache.read(torch.tensor(token_ids))
+    result = mem_manager.to_kvcache(result.tolist())
+    assert result.eq(kvcache).all(), f"Retrieved kvcache: {result}, Expected kvcache: {kvcache}"
+    print("[TEST] Basic test passed. Retrieved kvcache\n\n")
+
+
+def test_node_splitting(mem_manager, hicache, token_size):
+    tokens_per_block = BLOCK_SIZE // token_size
+    # 生成超过一个块的数据
+    token_ids = list(range(12, 12 + tokens_per_block * 3 + 1))
+    indices = mem_manager.alloc(len(token_ids))
+    kvcache = mem_manager.to_kvcache(indices)
+
+    hicache.write(torch.tensor(token_ids), torch.tensor(indices))
+    time.sleep(2)
+    hicache.service.wait_till_all_finished()
+
+    # 验证根节点应该有子节点
+    root = hicache.root
+    assert len(root.children) > 0
+    print(f"\nRoot node has {len(root.children)} children")
+
+    # 读取完整序列
+    result = hicache.read(torch.tensor(token_ids))
+    result = mem_manager.to_kvcache(result.tolist())
+    assert result.eq(kvcache).all(), f"Retrieved kvcache: {result}, Expected kvcache: {kvcache}"
+    print(f"[TEST] Node splitting test passed. Retrieved kvcache: {result.shape}\n\n")
+
+
+def test_partial_read(mem_manager, hicache):
+    token_ids = [97, 98, 99, 100, 101, 102]
+    indices = mem_manager.alloc(len(token_ids))
+    kvcache = mem_manager.to_kvcache(indices)
+    hicache.write(torch.tensor(token_ids), torch.tensor(indices))
+    time.sleep(2)
+    hicache.service.wait_till_all_finished()
+
+    # 查询存在的部分前缀
+    result = hicache.read(torch.tensor([97, 98, 99]))
+    result = mem_manager.to_kvcache(result.tolist())
+    assert result.eq(kvcache[:3]).all()
+    print("[TEST] Partial read passed")
+
+    # 查询不存在的前缀
+    result = hicache.read(torch.tensor([97, 98, 100]))
+    assert len(result) == 2
+    result = mem_manager.to_kvcache(result.tolist())
+    assert result.eq(kvcache[:2]).all()
+    print(f"[TEST] Non-existent prefix returned: {result.tolist()}")
+
+
+def main():
+    mem_manager, service, hicache, token_size = setup()
+    try:
+        test_basic_write_read(mem_manager, hicache, token_size)
+        test_node_splitting(mem_manager, hicache, token_size)
+        test_partial_read(mem_manager, hicache)
+    finally:
+        service.shutdown()
+
+
+if __name__ == "__main__":
+    main()