From 0986bed8e279ce67659e4e473300fc3c5c717d6a Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Tue, 24 Mar 2026 11:02:50 +0800 Subject: [PATCH] [HiCache][HybridModel]: Support mamba state offloading & HybridCacheController (#20457) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: pansicheng Co-authored-by: 晟海 Co-authored-by: ispobock --- benchmark/hicache/bench_multiturn.py | 11 + python/sglang/srt/disaggregation/decode.py | 3 + python/sglang/srt/managers/schedule_batch.py | 3 - python/sglang/srt/managers/schedule_policy.py | 1 + python/sglang/srt/managers/scheduler.py | 6 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 7 +- .../srt/mem_cache/hi_mamba_radix_cache.py | 832 +++++++++++++----- .../sglang/srt/mem_cache/hicache_storage.py | 259 +++++- .../hybrid_cache/hybrid_cache_controller.py | 500 +++++++++++ .../sglang/srt/mem_cache/mamba_radix_cache.py | 11 + python/sglang/srt/mem_cache/memory_pool.py | 29 +- .../sglang/srt/mem_cache/memory_pool_host.py | 530 ++++++++++- .../4-gpu-models/test_qwen35_models.py | 40 +- 13 files changed, 2012 insertions(+), 220 deletions(-) create mode 100644 python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py diff --git a/benchmark/hicache/bench_multiturn.py b/benchmark/hicache/bench_multiturn.py index d821bbc7b..49483fe84 100644 --- a/benchmark/hicache/bench_multiturn.py +++ b/benchmark/hicache/bench_multiturn.py @@ -529,6 +529,8 @@ class WorkloadGenerator: f"({expected} clients), releasing {len(next_round_reqs)} " f"requests for round {current_barrier_round + 1}" ) + self._send_heartbeat(input_len=100, output_len=100) + time.sleep(10) for req in next_round_reqs: self.ready_queue.append(req) next_round_reqs = [] @@ -541,6 +543,15 @@ class WorkloadGenerator: print(f"Error processing response for client {client_id}: {e}") continue + def _send_heartbeat(self, input_len=100, output_len=20): + """Send a small heartbeat request to the server.""" + heartbeat_input = [1] * input_len + payload = gen_payload(heartbeat_input, output_len, self.lora_path) + try: + requests.post(self.url, json=payload, timeout=30) + except Exception as e: + print(f"Heartbeat request failed: {e}") + def run(self): request_thread = threading.Thread(target=self.request_sender, daemon=True) response_thread = threading.Thread(target=self.response_handler, daemon=True) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index b64fa4ff9..ce7ae1557 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -195,6 +195,9 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): effective_mamba_size = ( mamba_size if mamba_size is not None else size ) + pre_alloc_size + # TODO: Support PP + self.start_layer = 0 + self.layer_transfer_counter = None self._init_mamba_pool( size=effective_mamba_size, mamba_spec_state_size=size + pre_alloc_size, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 670cc1c85..9657a72b1 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -641,7 +641,6 @@ class Req(ReqDllmMixin): self.extend_logprob_start_len = 0 self.last_node: Any = None self.last_host_node: Any = None - self.last_host_backup_node: Any = None self.host_hit_length = 0 # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 @@ -915,14 +914,12 @@ class Req(ReqDllmMixin): self.prefix_indices, self.last_node, self.last_host_node, - self.last_host_backup_node, self.host_hit_length, self.mamba_branching_seqlen, ) = ( match_result.device_indices, match_result.last_device_node, match_result.last_host_node, - match_result.last_host_backup_node, match_result.host_hit_length, match_result.mamba_branching_seqlen, ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index dbce81617..2cce30600 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -777,6 +777,7 @@ class PrefillAdder: InitLoadBackParams( last_host_node=req.last_host_node, host_hit_length=req.host_hit_length, + req=req, ) ) req.prefix_indices = torch.cat([req.prefix_indices, new_indices]) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f31303e8b..2c52825a4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1818,11 +1818,7 @@ class Scheduler( def _prefetch_kvcache(self, req: Req): if self.enable_hicache_storage: req.init_next_round_input(self.tree_cache, cow_mamba=False) - last_host_node = ( - req.last_host_backup_node - if req.last_host_backup_node is not None - else req.last_host_node - ) + last_host_node = req.last_host_node if last_host_node.backuped or last_host_node is self.tree_cache.root_node: last_hash = last_host_node.get_last_hash_value() matched_len = len(req.prefix_indices) + req.host_hit_length diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 0efc27d44..be219339c 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -129,8 +129,10 @@ class MatchResult(NamedTuple): last_host_node : The last TreeNode on the host that was matched. Note that if HiCache is not enabled, this **must** be the same as `last_device_node`. - last_host_backup_node: The deepest backuped node for prefetch from storage. - host_hit_length : Length of the KV cache hit on the host, if applicable. + host_hit_length : Length of the host cache hit. For pure-KV caches this is the + number of evicted KV tokens on CPU. For hybrid Mamba models this + is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only + host hit still triggers load-back without adding a separate field. 0 if HiCache is not enabled. mamba_branching_seqlen: The mamba radix cache branching point, which is the longest page-aligned position that could've been cache hit if there @@ -140,7 +142,6 @@ class MatchResult(NamedTuple): device_indices: torch.Tensor last_device_node: Any last_host_node: Any - last_host_backup_node: Any = None host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None cache_protected_len: Optional[int] = None diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 13398b430..f96c5ac8b 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -12,7 +12,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch -from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, DecLockRefResult, @@ -23,13 +22,25 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, + PrefetchOperation, +) from sglang.srt.mem_cache.mamba_radix_cache import ( + LRUList, MambaRadixCache, TreeNode, get_last_access_time, ) -from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool -from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool +from sglang.srt.mem_cache.memory_pool_host import ( + HostPoolGroup, + MambaPoolHost, + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + PoolEntry, +) from sglang.srt.mem_cache.radix_cache import ( RadixKey, compute_node_hash_values, @@ -45,12 +56,43 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -class HiMambaRadixCache(MambaRadixCache): - """Hierarchical cache for hybrid Mamba models. +class HostLRUList(LRUList): + def __init__(self): + super().__init__(mamba=True) + self.prv = "host_mamba_prev" + self.nxt = "host_mamba_next" + setattr(self.head, self.nxt, self.tail) + setattr(self.tail, self.prv, self.head) - Only the Full (attention) KV cache is backed up to L2 (host) / L3 (storage). - Mamba states remain on device and are managed by LRU eviction. - """ + def reset_node_mru(self, node): + assert node.id in self.cache, f"Resetting node {node.id=} not in host mamba lru" + assert ( + node.mamba_host_value is not None + ), f"Resetting host mamba tombstone node in lru list: {node.id=}" + self._remove_node(node) + self._add_node(node) + + def insert_mru(self, node): + assert ( + node.mamba_host_value is not None + ), f"Inserting host mamba tombstone node in lru list: {node.id=}" + assert ( + node.id not in self.cache + ), f"Inserting node {node.id=} already in host mamba lru list" + self.cache[node.id] = node + self._add_node(node) + + def remove_node(self, node: TreeNode): + assert node.id in self.cache, f"Removing node {node.id=} not in host mamba lru" + assert ( + node.mamba_host_value is not None + ), f"Removing host mamba tombstone node from lru list: {node.id=}" + del self.cache[node.id] + self._remove_node(node) + + +class HiMambaRadixCache(MambaRadixCache): + """Hierarchical cache for hybrid Mamba models.""" def __init__(self, params: CacheInitParams, server_args: ServerArgs): self._enable_metrics_flag = params.enable_metrics @@ -66,13 +108,23 @@ class HiMambaRadixCache(MambaRadixCache): bind_to_closest_numa_node_cuda() self.page_size = params.page_size - kvcache = params.token_to_kv_pool_allocator.get_kvcache() + self.hybrid_kv_cache = params.token_to_kv_pool_allocator.get_kvcache() + if not isinstance(self.hybrid_kv_cache, HybridLinearKVPool): + raise ValueError( + "HiMambaRadixCache requires HybridLinearKVPool for hybrid SSM models." + ) + if not isinstance(params.req_to_token_pool, HybridReqToTokenPool): + raise ValueError( + "HiMambaRadixCache requires HybridReqToTokenPool for hybrid SSM models." + ) - if isinstance(kvcache, HybridLinearKVPool): - kvcache = kvcache.full_kv_pool - self.kvcache = kvcache - - self.full_kv_pool_host = MHATokenToKVPoolHost( + self.kvcache = self.hybrid_kv_cache.full_kv_pool + kv_host_pool_cls = ( + MLATokenToKVPoolHost + if self.hybrid_kv_cache.use_mla + else MHATokenToKVPoolHost + ) + self.full_kv_pool_host = kv_host_pool_cls( self.kvcache, server_args.hicache_ratio, server_args.hicache_size, @@ -80,6 +132,53 @@ class HiMambaRadixCache(MambaRadixCache): server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) + self.mamba_pool_host = MambaPoolHost( + params.req_to_token_pool.mamba_pool, + server_args.hicache_ratio, + server_args.hicache_size, + allocator_type=server_args.hicache_storage_backend, + layout=server_args.hicache_mem_layout, + ) + + full_layer_ids = sorted( + self.hybrid_kv_cache.full_attention_layer_id_mapping.keys() + ) + mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) + self.transfer_layer_num = len(set(full_layer_ids) | set(mamba_layer_ids)) + + full_layer_mapping = dict(self.hybrid_kv_cache.full_attention_layer_id_mapping) + mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) + transfer_layer_num = self.transfer_layer_num + + def kv_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: + return None + return full_layer_mapping.get(layer_id) + + def mamba_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: + return None + return mamba_layer_mapping.get(layer_id) + + self.host_pool_group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=self.full_kv_pool_host, + device_pool=self.kvcache, + layer_mapper=kv_layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name=PoolName.MAMBA, + host_pool=self.mamba_pool_host, + device_pool=params.req_to_token_pool.mamba_pool, + layer_mapper=mamba_layer_mapper, + host_evict_fn=self.evict_mamba_host, + device_evict_fn=self.evict_mamba, + ), + ] + ) self.tp_group = params.tp_cache_group self.tp_world_size = ( @@ -105,9 +204,9 @@ class HiMambaRadixCache(MambaRadixCache): self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy self.load_cache_event = threading.Event() - self.cache_controller = HiCacheController( + self.cache_controller = HybridCacheController( params.token_to_kv_pool_allocator, - self.full_kv_pool_host, + self.host_pool_group, params.page_size, self.tp_group, load_cache_event=self.load_cache_event, @@ -119,6 +218,13 @@ class HiMambaRadixCache(MambaRadixCache): storage_backend_extra_config=extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + transfer_layer_num=self.transfer_layer_num, + ) + params.req_to_token_pool.register_layer_transfer_counter( + self.cache_controller.layer_done_counter + ) + self.hybrid_kv_cache.register_layer_transfer_counter( + self.cache_controller.layer_done_counter ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, @@ -146,6 +252,7 @@ class HiMambaRadixCache(MambaRadixCache): self.evictable_full_device_leaves: set[TreeNode] = set() self.evictable_full_host_leaves: set[TreeNode] = set() + self.mamba_host_lru_list = HostLRUList() # Detach storage backend automatically on process shutdown atexit.register(self.shutdown) @@ -154,8 +261,10 @@ class HiMambaRadixCache(MambaRadixCache): def reset(self) -> None: TreeNode.counter = 0 + self._flush_pending_storage_backups_before_reset() self.cache_controller.reset() self.full_kv_pool_host.clear() + self.mamba_pool_host.clear() self.ongoing_write_through = {} self.ongoing_load_back = {} self.ongoing_prefetch = {} @@ -163,21 +272,37 @@ class HiMambaRadixCache(MambaRadixCache): self.prefetch_loaded_tokens_by_reqid.clear() self.evictable_full_device_leaves.clear() self.evictable_full_host_leaves.clear() + self.mamba_host_lru_list = HostLRUList() + logger.info( + "HiMambaRadixCache reset completed: host_kv_available=%s host_mamba_available=%s", + self.full_kv_pool_host.available_size(), + self.mamba_pool_host.available_size(), + ) super().reset() def write_backup(self, node: TreeNode, write_back=False): + # If mamba host slot already exists, refresh its LRU position. + if node.mamba_value is not None and node.mamba_host_value is not None: + if self.mamba_host_lru_list.in_list(node): + self.mamba_host_lru_list.reset_node_mru(node) + + extra_pools = self.mamba_backup_transfers(node) host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, + extra_pools=extra_pools, ) if host_indices is None: self.evict_host(len(node.value)) host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, + extra_pools=extra_pools, ) if host_indices is not None: node.host_value = host_indices + if extra_pools is not None: + self.mamba_backup_commit(node, extra_pools) assert len(node.host_value) > 0 self.ongoing_write_through[node.id] = node if not write_back: @@ -189,7 +314,7 @@ class HiMambaRadixCache(MambaRadixCache): return len(host_indices) def load_back( - self, node: TreeNode, mem_quota: Optional[int] = None + self, node: TreeNode, mem_quota: Optional[int] = None, req=None ) -> Optional[torch.Tensor]: """Load full KV back from host.""" last_hit_node = node @@ -202,34 +327,64 @@ class HiMambaRadixCache(MambaRadixCache): else: ancestor_node = node + mamba_restore_nodes = [] + if last_hit_node.mamba_backuped and last_hit_node.mamba_evicted: + mamba_restore_nodes.append(last_hit_node) + result = self.inc_lock_ref(ancestor_node) delta = result.delta - full_host_indices = torch.cat([n.host_value for n in nodes_to_load]) - if (len(full_host_indices) < self.load_back_threshold) or ( - len(full_host_indices) > mem_quota + delta - if mem_quota is not None - else False + if nodes_to_load: + full_host_indices = torch.cat([n.host_value for n in nodes_to_load]) + else: + full_host_indices = torch.empty((0,), dtype=torch.int64, device="cpu") + + if ( + len(full_host_indices) > 0 + and ( + (len(full_host_indices) < self.load_back_threshold) + or ( + len(full_host_indices) > mem_quota + delta + if mem_quota is not None + else False + ) + ) + and len(mamba_restore_nodes) == 0 ): # skip loading back if the total size is too small or exceeding the memory quota self.dec_lock_ref(ancestor_node) return None + logger.debug( + f"Init load back from cpu -> gpu, kv hit length: {len(full_host_indices)}, mamba host hit length: {len(mamba_restore_nodes)}" + ) + mamba_pools = self.mamba_restore_transfers( + last_hit_node, mamba_restore_nodes, req + ) full_device_indices = self.cache_controller.load( host_indices=full_host_indices, node_id=last_hit_node.id, + extra_pools=mamba_pools, ) if full_device_indices is None: - self.evict(EvictParams(num_tokens=len(full_host_indices))) + if len(full_host_indices) > 0: + self.evict(EvictParams(num_tokens=len(full_host_indices))) + + mamba_pools = self.mamba_restore_transfers( + last_hit_node, mamba_restore_nodes, req + ) full_device_indices = self.cache_controller.load( host_indices=full_host_indices, node_id=last_hit_node.id, + extra_pools=mamba_pools, ) self.dec_lock_ref(ancestor_node) if full_device_indices is None: # no sufficient GPU memory to load back KV caches return None + self.mamba_restore_commit(mamba_restore_nodes, mamba_pools) + offset = 0 for n in nodes_to_load: n_len = len(n.host_value) @@ -238,16 +393,15 @@ class HiMambaRadixCache(MambaRadixCache): self.full_lru_list.insert_mru(n) self.full_evictable_size_ += n_len - - if n.mamba_value is not None: - if self.mamba_lru_list.in_list(n): - self.mamba_lru_list.reset_node_mru(n) - else: - self.mamba_lru_list.insert_mru(n) - self.mamba_evictable_size_ += len(n.mamba_value) - self._update_leaf_status(n) + for n in mamba_restore_nodes: + if self.mamba_lru_list.in_list(n): + self.mamba_lru_list.reset_node_mru(n) + else: + self.mamba_lru_list.insert_mru(n) + self.mamba_evictable_size_ += len(n.mamba_value) + self._update_leaf_status(ancestor_node) self.inc_lock_ref(last_hit_node) @@ -261,15 +415,18 @@ class HiMambaRadixCache(MambaRadixCache): ): last_node = params.last_host_node mem_quota = params.mem_quota - if last_node.evicted: - loading_values = self.load_back(last_node, mem_quota) + req = params.req + if last_node.evicted or (last_node.mamba_evicted and last_node.mamba_backuped): + loading_values = self.load_back(last_node, mem_quota, req=req) if loading_values is not None: logger.debug( f"loading back {len(loading_values)} tokens for node {last_node.id}" ) return loading_values, last_node - while last_node.evicted: + while last_node is not self.root_node and ( + last_node.evicted or last_node.mamba_evicted + ): last_node = last_node.parent return ( @@ -400,28 +557,110 @@ class HiMambaRadixCache(MambaRadixCache): return self.evictable_full_host_leaves.add(node) - def _evict_to_host(self, node: TreeNode) -> int: - """Evict full KV to host. Mamba stays on device. Returns num evicted.""" - num_full = len(node.value) + def _free_device_mamba(self, node: TreeNode) -> int: + if node.mamba_value is None: + return 0 + mamba_num = len(node.mamba_value) + self.req_to_token_pool.mamba_pool.free(node.mamba_value) + if node.mamba_lock_ref > 0: + self.mamba_protected_size_ -= mamba_num + node.mamba_lock_ref = 0 + else: + self.mamba_evictable_size_ -= mamba_num + if self.mamba_lru_list.in_list(node): + self.mamba_lru_list.remove_node(node) + node.mamba_value = None + return mamba_num - if not node.backuped: - if self.cache_controller.write_policy == "write_back": - self.write_backup(node, write_back=True) - self.writing_check(write_back=True) + def _evict_to_host(self, node: TreeNode) -> Tuple[int, int]: + # GPU -> CPU demotion: node stays in tree as evicted+backuped + assert not node.evicted, f"already evicted, {node.id=}" + assert node.backuped, f"not backuped, {node.id=}" + + num_full = len(node.value) self.cache_controller.evict_device(node.value) self.full_evictable_size_ -= num_full if self.full_lru_list.in_list(node): self.full_lru_list.remove_node(node) + mamba_num = self._free_device_mamba(node) + node.value = None self._update_leaf_status(node) self._update_full_device_leaf_status(node.parent) - return num_full + return num_full, mamba_num + + def _evict_regular(self, node: TreeNode) -> Tuple[int, int]: + # evict a non-backuped device leaf — free GPU KV + mamba, delete from tree + assert not node.evicted, f"already evicted, {node.id=}" + assert not node.backuped, f"backuped node, {node.id=}" + assert len(node.children) == 0, f"non-leaf, {node.id=}" + + full_num_evicted = len(node.value) + + self.cache_controller.evict_device(node.value) + self.full_evictable_size_ -= full_num_evicted + if self.full_lru_list.in_list(node): + self.full_lru_list.remove_node(node) + + mamba_num_evicted = self._free_device_mamba(node) + + if node.mamba_host_value is not None: + if self.mamba_host_lru_list.in_list(node): + self.mamba_host_lru_list.remove_node(node) + self.mamba_pool_host.free(node.mamba_host_value) + node.mamba_host_value = None + + node.value = None + self._discard_from_leaf_sets(node) + + parent = node.parent + key = self.get_child_key_fn(node.key) + v = parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + + self._update_leaf_status(parent) + _, cascade_full_num_evicted, cascade_mamba_num_evicted = ( + self._iteratively_delete_tombstone_leaf(node) + ) + return ( + full_num_evicted + cascade_full_num_evicted, + mamba_num_evicted + cascade_mamba_num_evicted, + ) + + def _evict_host_leaf(self, node: TreeNode) -> int: + # evict a host-resident leaf: free host KV + mamba, delete from tree, cascade + assert node.evicted, f"not evicted, {node.id=}" + assert node.backuped, f"not backuped, {node.id=}" + assert node.mamba_value is None, f"has device mamba, {node.id=}" + assert ( + node.host_ref_counter == 0 + ), f"in use, {node.id=} {node.host_ref_counter=}" + + full_num_evicted = self.cache_controller.evict_host(node.host_value) + node.host_value = None + + if node.mamba_host_value is not None: + if self.mamba_host_lru_list.in_list(node): + self.mamba_host_lru_list.remove_node(node) + self.mamba_pool_host.free(node.mamba_host_value) + node.mamba_host_value = None + + self._discard_from_leaf_sets(node) + parent = node.parent + key = self.get_child_key_fn(node.key) + v = parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + + self._update_leaf_status(parent) + _, cascade_full_num_evicted, _ = self._iteratively_delete_tombstone_leaf(node) + + return full_num_evicted + cascade_full_num_evicted def _delete_tombstone_leaf(self, node: TreeNode) -> None: - """Remove a tombstone leaf from the tree and free HiCache resources.""" - assert node.mamba_value is None, f"node has mamba value, {node.id=}" + assert node.mamba_value is None, f"has mamba value, {node.id=}" + assert node.mamba_host_value is None, f"has mamba host value, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}" parent = node.parent key = self.get_child_key_fn(node.key) @@ -439,9 +678,6 @@ class HiMambaRadixCache(MambaRadixCache): def _iteratively_delete_tombstone_leaf( self, node: TreeNode ) -> Tuple[TreeNode, int, int]: - """Cascade-delete tombstone parents that became leaves. - Returns (node, full_num_evicted, mamba_num_evicted). - """ full_num_evicted = 0 mamba_num_evicted = 0 @@ -450,6 +686,8 @@ class HiMambaRadixCache(MambaRadixCache): break if node.parent.mamba_value is not None: break + if node.parent.mamba_host_value is not None: + break if node.parent.full_lock_ref > 0 or node.parent.mamba_lock_ref > 0: break @@ -468,6 +706,22 @@ class HiMambaRadixCache(MambaRadixCache): return node, full_num_evicted, mamba_num_evicted + def _evict_device_leaf(self, x: TreeNode) -> Tuple[int, int]: + """Evict a device leaf node, choosing the right strategy: + + - backuped: demote to host via _evict_to_host (node stays in tree) + - not backuped + write_back: write_backup first, then demote + - not backuped + write_through: _evict_regular (delete from tree) + """ + if not x.backuped: + if self.cache_controller.write_policy == "write_back": + self.write_backup(x, write_back=True) + self.writing_check(write_back=True) + return self._evict_to_host(x) + else: + return self._evict_regular(x) + return self._evict_to_host(x) + def evict(self, params: EvictParams) -> EvictResult: if self.disable: return EvictResult() @@ -486,15 +740,16 @@ class HiMambaRadixCache(MambaRadixCache): if x not in self.evictable_full_device_leaves: continue - evicted_full = self._evict_to_host(x) + evicted_full, evicted_mamba = self._evict_device_leaf(x) full_num_evicted += evicted_full + mamba_num_evicted += evicted_mamba parent = x.parent if parent in self.evictable_full_device_leaves: heapq.heappush(eviction_heap, (parent.last_access_time, parent)) if params.mamba_num > 0: - mamba_num_evicted = self.evict_mamba(params.mamba_num) + mamba_num_evicted += self.evict_mamba(params.mamba_num) return EvictResult( num_tokens_evicted=full_num_evicted, @@ -502,27 +757,7 @@ class HiMambaRadixCache(MambaRadixCache): ) def evict_host(self, num_tokens: int): - if self.enable_storage: - host_leaves = list(self.evictable_full_host_leaves) - heap = [(n.last_access_time, n) for n in host_leaves] - heapq.heapify(heap) - - num_evicted = 0 - while num_evicted < num_tokens and heap: - _, x = heapq.heappop(heap) - if x not in self.evictable_full_host_leaves: - continue - - num_evicted += self.cache_controller.evict_host(x.host_value) - x.host_value = None - - self.evictable_full_host_leaves.discard(x) - self._update_full_host_leaf_status(x.parent) - if x.parent in self.evictable_full_host_leaves: - heapq.heappush(heap, (x.parent.last_access_time, x.parent)) - return - - # Non-L3 path: evict host leaves and clean up tree + """Evict host-resident leaf nodes: free host KV + mamba, delete from tree, cascade.""" heap = [(n.last_access_time, n) for n in self.evictable_full_host_leaves] heapq.heapify(heap) @@ -532,28 +767,49 @@ class HiMambaRadixCache(MambaRadixCache): if x not in self.evictable_full_host_leaves: continue - num_evicted += self.cache_controller.evict_host(x.host_value) - x.host_value = None + num_evicted += self._evict_host_leaf(x) - if x.mamba_value is not None: - if self.mamba_lru_list.in_list(x): - self.mamba_lru_list.remove_node(x) - self.mamba_evictable_size_ -= len(x.mamba_value) - self.req_to_token_pool.mamba_pool.free(x.mamba_value) - x.mamba_value = None + if x.parent in self.evictable_full_host_leaves: + heapq.heappush(heap, (x.parent.last_access_time, x.parent)) - self.evictable_full_host_leaves.discard(x) + def evict_mamba_host(self, num_mamba_hosts: int) -> int: + """Evict host mamba states. - parent = x.parent - child_key = self.get_child_key_fn(x.key) - v = parent.children.pop(child_key, None) - assert v == x, f"parent does not have child key, {x.id=}" + Internal host node: free host mamba only (tombstone). + Host leaf node: same as Full host evict — _evict_host_leaf_node frees + host KV + mamba, deletes from tree, cascades. + """ + if self.disable or num_mamba_hosts <= 0: + return 0 - self._update_leaf_status(parent) - if parent in self.evictable_full_host_leaves: - heapq.heappush(heap, (parent.last_access_time, parent)) + x = self.mamba_host_lru_list.get_lru_no_lock() + num_evicted = 0 + while num_evicted < num_mamba_hosts and self.mamba_host_lru_list.in_list(x): + x_next = self.mamba_host_lru_list.get_prev_no_lock(x) + if x.host_ref_counter > 0: + x = x_next + continue + + if x in self.evictable_full_host_leaves: + self._evict_host_leaf(x) + num_evicted += 1 + else: + # internal host node: free host mamba only (tombstone) + self.mamba_host_lru_list.remove_node(x) + self.mamba_pool_host.free(x.mamba_host_value) + x.mamba_host_value = None + num_evicted += 1 + + x = x_next + return num_evicted def evict_mamba(self, mamba_num: int) -> int: + """Evict mamba states. + + Internal node: tombstone — free GPU mamba only, KV stays on GPU. + Leaf node: same as Full evict — _evict_to_host moves KV+mamba to host, + node stays in tree, then cascade tombstone parent device leaves. + """ if self.disable or mamba_num <= 0: return 0 @@ -563,64 +819,26 @@ class HiMambaRadixCache(MambaRadixCache): assert x.mamba_value is not None, f"node has no mamba value, {x.id=}" assert x != self.root_node, f"root node is not evictable, {x.id=}" assert x.mamba_lock_ref == 0, f"node is in use, {x.id=}" + assert ( + not x.evicted + ), f"evicted node should not be in mamba_lru_list, {x.id=}" - if x.evicted: - self.req_to_token_pool.mamba_pool.free(x.mamba_value) - mamba_num_evicted += len(x.mamba_value) + if len(x.children) > 0: + # Internal: free device mamba only, KV stays on device (tombstone) x_next = self.mamba_lru_list.get_prev_no_lock(x) - self.mamba_lru_list.remove_node(x) - self.mamba_evictable_size_ -= len(x.mamba_value) - x.mamba_value = None - - if len(x.children) == 0: - self._delete_tombstone_leaf(x) - _, _, cascade_mamba = self._iteratively_delete_tombstone_leaf(x) - mamba_num_evicted += cascade_mamba - elif len(x.children) > 0: - self.req_to_token_pool.mamba_pool.free(x.mamba_value) mamba_num_evicted += len(x.mamba_value) - x_next = self.mamba_lru_list.get_prev_no_lock(x) + self.req_to_token_pool.mamba_pool.free(x.mamba_value) self.mamba_lru_list.remove_node(x) self._tombstone_internal_node(x) else: + # Leaf: evict KV + mamba atomically assert ( x.full_lock_ref == 0 ), f"evict leaf node invalid with {x.id=} {x.full_lock_ref=}" - if not x.backuped: - if self.cache_controller.write_policy == "write_back": - self.write_backup(x, write_back=True) - self.writing_check(write_back=True) - - self.cache_controller.evict_device(x.value) - self.full_evictable_size_ -= len(x.value) - - self.req_to_token_pool.mamba_pool.free(x.mamba_value) - mamba_num_evicted += len(x.mamba_value) - x_next = self.mamba_lru_list.get_prev_no_lock(x) - if self.full_lru_list.in_list(x): - self.full_lru_list.remove_node(x) - self.mamba_lru_list.remove_node(x) - self.mamba_evictable_size_ -= len(x.mamba_value) - - if x.backuped: - self.cache_controller.evict_host(x.host_value) - x.host_value = None - - x.value = None - x.mamba_value = None - - self._discard_from_leaf_sets(x) - - parent = x.parent - child_key = self.get_child_key_fn(x.key) - v = parent.children.pop(child_key, None) - assert v == x, f"parent does not have child key, {x.id=}" - - self._update_leaf_status(parent) - _, _, cascade_mamba = self._iteratively_delete_tombstone_leaf(x) - mamba_num_evicted += cascade_mamba + _, mamba_evicted = self._evict_device_leaf(x) + mamba_num_evicted += mamba_evicted if not self.mamba_lru_list.in_list(x_next): x_next = self.mamba_lru_list.get_lru_no_lock() @@ -630,20 +848,14 @@ class HiMambaRadixCache(MambaRadixCache): return mamba_num_evicted def _unevict_node(self, node: TreeNode, fresh_value: torch.Tensor): - """Restore an evicted node with fresh device KV from the request.""" + assert node.evicted, f"not evicted, {node.id=}" + assert node.mamba_value is None, f"evicted node has device mamba, {node.id=}" n = len(fresh_value) node.value = fresh_value.clone() self.full_lru_list.insert_mru(node) self.full_evictable_size_ += n - if node.mamba_value is not None: - if self.mamba_lru_list.in_list(node): - self.mamba_lru_list.reset_node_mru(node) - else: - self.mamba_lru_list.insert_mru(node) - self.mamba_evictable_size_ += len(node.mamba_value) - self._update_leaf_status(node) if node.parent is not None: self._update_leaf_status(node.parent) @@ -752,7 +964,6 @@ class HiMambaRadixCache(MambaRadixCache): device_indices=torch.empty((0,), dtype=torch.int64, device=self.device), last_device_node=self.root_node, last_host_node=self.root_node, - last_host_backup_node=self.root_node, host_hit_length=0, ) @@ -760,17 +971,13 @@ class HiMambaRadixCache(MambaRadixCache): page_aligned_len = len(key) // self.page_size * self.page_size key = key[:page_aligned_len] - value, best_last_node, best_value_len, deepest_node = self._match_prefix_helper( - key - ) - return self._match_post_processor( - params, value, best_last_node, best_value_len, deepest_node - ) + value, best_last_node, best_value_len = self._match_prefix_helper(key) + return self._match_post_processor(params, value, best_last_node, best_value_len) def _match_prefix_helper( self, key: RadixKey - ) -> Tuple[List[torch.Tensor], TreeNode, int, TreeNode]: - """Walk tree to find best_last_node (mamba boundary) and deepest node.""" + ) -> Tuple[List[torch.Tensor], TreeNode, int]: + """Walk tree to find best_last_node (mamba boundary).""" node = self.root_node child_key = self.get_child_key_fn(key) @@ -784,7 +991,7 @@ class HiMambaRadixCache(MambaRadixCache): if child.evicted and not child.backuped: break - if node.mamba_value is not None: + if node.mamba_value is not None or node.mamba_backuped: best_value_len = len(value) best_last_node = node @@ -803,12 +1010,11 @@ class HiMambaRadixCache(MambaRadixCache): if len(key): child_key = self.get_child_key_fn(key) - if node.mamba_value is not None: + if node.mamba_value is not None or node.mamba_backuped: best_value_len = len(value) best_last_node = node - deepest_node = node - return value, best_last_node, best_value_len, deepest_node + return value, best_last_node, best_value_len def _match_post_processor( self, @@ -816,7 +1022,6 @@ class HiMambaRadixCache(MambaRadixCache): value: List[torch.Tensor], best_last_node: TreeNode, best_value_len: int, - deepest_node: TreeNode, ) -> MatchResult: cow_mamba = params.cow_mamba req = params.req @@ -850,33 +1055,31 @@ class HiMambaRadixCache(MambaRadixCache): else: mamba_branching_seqlen = None - # last_device_node & host_hit_length: from best_last_node (mamba boundary) - host_hit_length = 0 + kv_host_hit_length = 0 last_device_node = best_last_node - while last_device_node.evicted: - host_hit_length += len(last_device_node.host_value) + while last_device_node is not self.root_node and last_device_node.evicted: + kv_host_hit_length += len(last_device_node.host_value) last_device_node = last_device_node.parent last_host_node = best_last_node + while last_host_node is not self.root_node and not last_host_node.backuped: + last_host_node = last_host_node.parent - # last_host_backup_node: from deepest_node, find backuped ancestor - last_host_backup_node = deepest_node - while ( - last_host_backup_node is not self.root_node - and not last_host_backup_node.backuped - ): - last_host_backup_node = last_host_backup_node.parent + mamba_host_hit = ( + 1 if (last_host_node.mamba_evicted and last_host_node.mamba_backuped) else 0 + ) + host_hit_length = max(kv_host_hit_length, mamba_host_hit) mamba_node = best_last_node if cow_mamba and mamba_node.mamba_value is not None: if req.mamba_pool_idx is None: - dst_index = self.req_to_token_pool.mamba_pool.alloc(1) - if dst_index is None: - self.inc_lock_ref(mamba_node) - self.evict_mamba(1) - dst_index = self.req_to_token_pool.mamba_pool.alloc(1) - self.dec_lock_ref(mamba_node) - assert dst_index is not None, "Can not alloc mamba cache" + dst_index = self._alloc_with_evict( + self.req_to_token_pool.mamba_pool, + 1, + self.evict_mamba, + lock_node=mamba_node, + error_message="Can not alloc mamba cache", + ) src_index = mamba_node.mamba_value self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) req.mamba_pool_idx = dst_index[0] @@ -895,7 +1098,6 @@ class HiMambaRadixCache(MambaRadixCache): device_indices=value, last_device_node=last_device_node, last_host_node=last_host_node, - last_host_backup_node=last_host_backup_node, host_hit_length=host_hit_length, mamba_branching_seqlen=mamba_branching_seqlen, ) @@ -1184,6 +1386,7 @@ class HiMambaRadixCache(MambaRadixCache): prefetch_threshold=prefetch_threshold, model_name=served_model_name, storage_backend_extra_config=extra_config, + host_pools=self.host_pool_group.entries, ) except Exception as e: logger.exception( @@ -1220,6 +1423,14 @@ class HiMambaRadixCache(MambaRadixCache): self.storage_metrics_collector = None return True, "Detached HiCache storage backend successfully." + def prefetch_abort(self, pool_transfers: Optional[list[PoolTransfer]]) -> None: + """Free any allocated mamba host slots on prefetch abort/revoke.""" + for transfer in pool_transfers or []: + if transfer.name == PoolName.MAMBA: + if transfer.host_indices is not None: + self.mamba_pool_host.free(transfer.host_indices) + break + def _force_release_pending_storage_ops(self): cc = self.cache_controller @@ -1237,6 +1448,12 @@ class HiMambaRadixCache(MambaRadixCache): logger.exception( "Failed to free host indices for prefetch %s", req_id ) + try: + self.prefetch_abort(getattr(_operation, "pool_transfers", None)) + except Exception: + logger.exception( + "Failed to release mamba host indices for prefetch %s", req_id + ) try: self._release_host_node(last_host_node) except Exception: @@ -1296,7 +1513,8 @@ class HiMambaRadixCache(MambaRadixCache): for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke): info = self.ongoing_prefetch.pop(req_id, None) if info is not None: - last_host_node, token_ids, _, _ = info + last_host_node, token_ids, _, operation = info + self.prefetch_abort(operation.pool_transfers) self._release_host_node(last_host_node) cc.prefetch_tokens_occupied -= len(token_ids) if cc.prefetch_tokens_occupied < 0: @@ -1499,12 +1717,13 @@ class HiMambaRadixCache(MambaRadixCache): if self.hicache_storage_pass_prefix_keys else None ) - + extra_pools = self.mamba_archive_transfers(node) operation_id = self.cache_controller.write_storage( node.host_value, node.key, node.hash_value, prefix_keys, + extra_pools=extra_pools, ) self.ongoing_backup[operation_id] = node self._protect_host_node(node) @@ -1529,15 +1748,31 @@ class HiMambaRadixCache(MambaRadixCache): return self._protect_host_node(last_host_node) - host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) - if host_indices is None: - self.evict_host(prefetch_length) - host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) + + # Allocate host KV memory + host_indices = self._alloc_with_evict( + self.cache_controller.mem_pool_host, + prefetch_length, + self.evict_host, + ) if host_indices is None: self._release_host_node(last_host_node) return + + # Allocate host mamba slot + extra_pools = self.mamba_prefetch_alloc(new_input_tokens, last_hash) + if extra_pools is None: + self.cache_controller.mem_pool_host.free(host_indices) + self._release_host_node(last_host_node) + return + operation = self.cache_controller.prefetch( - req_id, host_indices, new_input_tokens, last_hash, prefix_keys + req_id, + host_indices, + new_input_tokens, + last_hash, + prefix_keys, + extra_pools=extra_pools, ) self.ongoing_prefetch[req_id] = ( last_host_node, @@ -1576,6 +1811,20 @@ class HiMambaRadixCache(MambaRadixCache): group=self.tp_group, ) min_completed_tokens = completed_tokens_tensor.item() + + mamba_host_indices = None + mamba_loaded = False + for transfer in operation.pool_transfers or []: + if transfer.name == PoolName.MAMBA: + mamba_host_indices = transfer.host_indices + mamba_loaded = ( + operation.pool_storage_result.extra_pool_hit_pages.get( + PoolName.MAMBA, 0 + ) + >= 1 + ) + break + fetched_token_ids = token_ids[:min_completed_tokens] written_indices = host_indices[:min_completed_tokens] matched_length = self._insert_helper_host( @@ -1586,12 +1835,22 @@ class HiMambaRadixCache(MambaRadixCache): ), written_indices, hash_value[: min_completed_tokens // self.page_size], + mamba_host_indices, + mamba_loaded, ) + # Free host KV memory: matched portion is already in tree, tail was unused self.cache_controller.mem_pool_host.free(host_indices[:matched_length]) self.cache_controller.append_host_mem_release( host_indices[min_completed_tokens:completed_tokens] ) + + # Free mamba host slot if it wasn't inserted into the tree + if mamba_host_indices is not None: + inserted_new = matched_length < min_completed_tokens + if not inserted_new or not mamba_loaded: + self.mamba_pool_host.free(mamba_host_indices) + self._release_host_node(last_host_node) del self.ongoing_prefetch[req_id] self.cache_controller.prefetch_tokens_occupied -= len(token_ids) @@ -1601,11 +1860,24 @@ class HiMambaRadixCache(MambaRadixCache): if self.enable_storage_metrics: self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage) + if loaded_from_storage > 0 and operation.pool_transfers: + logger.debug( + "HiCache mamba prefetch completed for request %s: prefetched_tokens=%s mamba_states=%s", + req_id, + loaded_from_storage, + int(mamba_loaded), + ) return True def _insert_helper_host( - self, node: TreeNode, key: RadixKey, host_value, hash_value + self, + node: TreeNode, + key: RadixKey, + host_value, + hash_value, + mamba_host_value: Optional[torch.Tensor] = None, + mamba_loaded: bool = False, ): node.last_access_time = get_last_access_time() if len(key) == 0: @@ -1614,7 +1886,6 @@ class HiMambaRadixCache(MambaRadixCache): child_key = self.get_child_key_fn(key) matched_length = 0 - host_value_inserted = False while len(key) > 0 and child_key in node.children.keys(): node = node.children[child_key] node.last_access_time = get_last_access_time() @@ -1622,22 +1893,10 @@ class HiMambaRadixCache(MambaRadixCache): self.mamba_lru_list.reset_node_mru(node) prefix_len = self.key_match_fn(node.key, key) - if node.evicted and not node.backuped: - node.host_value = host_value[:prefix_len].clone() - host_value_inserted = True - self._update_full_host_leaf_status(node) - if node.parent is not None: - self._update_full_host_leaf_status(node.parent) - else: - assert not host_value_inserted, ( - f"matched node after host_value insertion would cause incorrect free, " - f"{node.id=} {matched_length=} {prefix_len=}" - ) - matched_length += prefix_len - key = key[prefix_len:] host_value = host_value[prefix_len:] hash_value = hash_value[prefix_len // self.page_size :] + matched_length += prefix_len if prefix_len < len(node.key): new_node = self._split_node(node.key, node, prefix_len) @@ -1646,6 +1905,7 @@ class HiMambaRadixCache(MambaRadixCache): if len(key): child_key = self.get_child_key_fn(key) + leaf_node: Optional[TreeNode] = None if len(key): new_node = TreeNode() new_node.parent = node @@ -1655,8 +1915,15 @@ class HiMambaRadixCache(MambaRadixCache): new_node.host_value = host_value.clone() new_node.hash_value = hash_value node.children[child_key] = new_node + leaf_node = new_node self._update_full_host_leaf_status(new_node) self._update_full_host_leaf_status(node) + + # Attach mamba state to the new leaf + if leaf_node is not None and mamba_host_value is not None and mamba_loaded: + leaf_node.mamba_host_value = mamba_host_value.clone() + if not self.mamba_host_lru_list.in_list(leaf_node): + self.mamba_host_lru_list.insert_mru(leaf_node) return matched_length def release_aborted_request(self, rid: str): @@ -1675,4 +1942,177 @@ class HiMambaRadixCache(MambaRadixCache): self._release_host_node(last_host_node) del self.ongoing_prefetch[rid] self.cache_controller.append_host_mem_release(host_indices[:completed_tokens]) + self.prefetch_abort(operation.pool_transfers) self.cache_controller.prefetch_tokens_occupied -= len(token_ids) + + def _flush_pending_storage_backups_before_reset(self) -> None: + if not self.enable_storage: + return + + self.writing_check(write_back=True) + deadline = time.monotonic() + 30.0 + while time.monotonic() < deadline: + self.drain_storage_control_queues() + backup_qsize = self.cache_controller.backup_queue.qsize() + ack_backup_qsize = self.cache_controller.ack_backup_queue.qsize() + ongoing_backup = len(self.ongoing_backup) + ongoing_write = len(self.ongoing_write_through) + if ( + backup_qsize == 0 + and ack_backup_qsize == 0 + and ongoing_backup == 0 + and ongoing_write == 0 + ): + return + time.sleep(0.05) + + logger.warning( + "Timed out waiting for HiCache storage backups to drain before reset: " + "ongoing_write=%s ongoing_backup=%s backup_queue=%s ack_backup_queue=%s", + len(self.ongoing_write_through), + len(self.ongoing_backup), + self.cache_controller.backup_queue.qsize(), + self.cache_controller.ack_backup_queue.qsize(), + ) + + def _alloc_with_evict( + self, + pool, + size: int, + evict_fn, + lock_node: Optional[TreeNode] = None, + error_message: Optional[str] = None, + ) -> Optional[torch.Tensor]: + indices = pool.alloc(size) + if indices is None: + if lock_node is not None: + self.inc_lock_ref(lock_node) + evict_fn(size) + indices = pool.alloc(size) + if lock_node is not None: + self.dec_lock_ref(lock_node) + if indices is None and error_message is not None: + raise RuntimeError(error_message) + return indices + + # -- mamba PoolTransfer builders (D↔H↔S) ---------------------------------- + + def mamba_backup_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: + # build D→H transfer descriptor for mamba state + if node.mamba_value is None: + return None + return [ + PoolTransfer( + name=PoolName.MAMBA, + host_indices=node.mamba_host_value, + device_indices=node.mamba_value, + ) + ] + + def mamba_backup_commit( + self, node: TreeNode, transfers: list[PoolTransfer] + ) -> None: + # store auto-allocated mamba host indices into the node after D→H backup + if not transfers: + return + host_indices = transfers[0].host_indices + if node.mamba_host_value is None and host_indices is not None: + node.mamba_host_value = host_indices + self.mamba_host_lru_list.insert_mru(node) + + def mamba_archive_transfers(self, node: TreeNode) -> Optional[list[PoolTransfer]]: + # build H→Storage transfer descriptor for mamba state + if node.mamba_host_value is None or not node.hash_value: + return None + return [ + PoolTransfer( + name=PoolName.MAMBA, + host_indices=node.mamba_host_value, + keys=[node.hash_value[-1]], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ] + + def mamba_prefetch_alloc( + self, + token_ids: List[int], + last_hash: Optional[str], + ) -> Optional[list[PoolTransfer]]: + # allocate a mamba host slot and build Storage→H transfer descriptor + if not token_ids: + return None + host_indices = self._alloc_with_evict( + self.mamba_pool_host, 1, self.evict_mamba_host + ) + if host_indices is None: + return None + # placeholder key; I/O thread replaces with correct hash after hit query + return [ + PoolTransfer( + name=PoolName.MAMBA, + host_indices=host_indices, + keys=["__placeholder__"], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ] + + def mamba_restore_transfers( + self, + last_hit_node: TreeNode, + nodes_to_restore: list[TreeNode], + req, + ) -> Optional[list[PoolTransfer]]: + # build H→D transfer descriptors for mamba state + backed_up_host_indices: list[torch.Tensor] = [] + for node in nodes_to_restore: + if not node.mamba_backuped: + continue + backed_up_host_indices.append(node.mamba_host_value) + + transfers: list[PoolTransfer] = [] + if backed_up_host_indices: + transfers.append( + PoolTransfer( + name=PoolName.MAMBA, + host_indices=torch.cat(backed_up_host_indices), + device_indices=None, + ) + ) + + if ( + req is not None + and last_hit_node in nodes_to_restore + and last_hit_node.mamba_host_value is not None + ): + if req.mamba_pool_idx is None: + req.mamba_pool_idx = self._alloc_with_evict( + self.req_to_token_pool.mamba_pool, + len(last_hit_node.mamba_host_value), + self.evict_mamba, + lock_node=last_hit_node, + error_message="Cannot alloc request mamba cache for host load back", + )[0] + transfers.append( + PoolTransfer( + name=PoolName.MAMBA, + host_indices=last_hit_node.mamba_host_value, + device_indices=req.mamba_pool_idx.unsqueeze(0), + ) + ) + + return transfers if transfers else None + + def mamba_restore_commit( + self, + restored_nodes: list[TreeNode], + transfers: Optional[list[PoolTransfer]], + ) -> None: + # write back controller-allocated device indices after H→D restore + if not restored_nodes or not transfers or transfers[0].device_indices is None: + return + device_indices = transfers[0].device_indices + offset = 0 + for node in restored_nodes: + count = len(node.mamba_host_value) + node.mamba_value = device_indices[offset : offset + count].clone() + offset += count diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index e76cfebd4..8e39fd0ac 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -3,7 +3,8 @@ import logging import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any, List, Optional +from enum import Enum +from typing import Any, List, Optional, Set import torch @@ -65,6 +66,61 @@ class HiCacheStorageExtraInfo: extra_info: Optional[dict] = None +class PoolName(str, Enum): + """Well-known pool names used as PoolTransfer/PoolEntry identifiers.""" + + KV = "kv" + MAMBA = "mamba" + + +class PoolHitPolicy(str, Enum): + """Hit policy for batch_exists_v2 per-pool prefix matching. + + ALL_PAGES : every page in [0, kv_hit) must exist (default). + TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states). + """ + + ALL_PAGES = "all_pages" + TRAILING_PAGES = "trailing_pages" + + +@dataclass +class PoolTransfer: + """Unified per-pool transfer descriptor for batch v2 interface. + + device<->host path : host_indices + device_indices + host<->storage path: host_indices + keys + """ + + name: PoolName + host_indices: Optional[torch.Tensor] = None + device_indices: Optional[torch.Tensor] = None + keys: Optional[List[str]] = None + hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES + + +@dataclass +class PoolTransferResult: + """Tracks how many pages were successfully processed per pool.""" + + kv_hit_pages: int + extra_pool_hit_pages: dict[str, int] + + @classmethod + def empty(cls) -> "PoolTransferResult": + return cls(0, {}) + + def update_kv_hit_pages(self, kv_hit_pages: int) -> None: + """Accumulate kv_hit_pages across batches (max = last successful batch).""" + self.kv_hit_pages = max(self.kv_hit_pages, kv_hit_pages) + + def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None: + """Record actual load/write success counts per extra pool.""" + self.extra_pool_hit_pages.update( + {name: sum(rs) for name, rs in results.items()} + ) + + class HiCacheStorage(ABC): """ HiCacheStorage is a class that provides a generic key-value interface for storing and retrieving KV cache. @@ -72,10 +128,69 @@ class HiCacheStorage(ABC): """ # todo, the page size of storage backend does not have to be the same as the same as host memory pool - def register_mem_pool_host(self, mem_pool_host: HostKVCache): self.mem_pool_host = mem_pool_host + def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name): + if not hasattr(self, "registered_pools"): + self.registered_pools = {} + self.registered_pools[host_pool_name] = host_pool + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + """Check which cache pages exist in storage, respecting per-pool hit policies. + + Longest-prefix semantics + Extra-pool hit policies (``PoolTransfer.hit_policy``) + ------------------------------------------------------ + Each ``PoolTransfer`` in ``pool_transfers`` describes a secondary + cache pool (e.g. Mamba SSM states) that must be co-present with the + KV pages. The final ``final_pages`` is the minimum across all pools, + so a missing auxiliary page shrinks the usable prefix. + + - ``"all_pages"`` (default): every page in [0, kv_hit) must exist + for this pool. Used for pools that are required for every token + in the prefix (e.g. DeepSeek DSA pool). + + - ``"trailing_pages"``: only the *last* ``len(transfer.keys)`` pages + of the KV prefix need to exist. Used for pools whose data covers + only the tail of a prefix (e.g. Mamba/SWA Pool). + + Returns + ------- + PoolTransferResult + ``kv_hit_pages`` = length of the usable KV prefix. + ``extra_pool_hit_pages`` maps each pool name to the number of pages + that were found. + """ + raise NotImplementedError() + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + """Read data from storage into host memory for each PoolTransfer. + + Returns a dict mapping pool name to a per-entry success list. + """ + raise NotImplementedError() + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + """Write data from host memory to storage for each PoolTransfer. + + Returns a dict mapping pool name to a per-entry success list. + """ + raise NotImplementedError() + def batch_get_v1( self, keys: List[str], @@ -203,7 +318,6 @@ class HiCacheFile(HiCacheStorage): self.config_suffix = f"_{model_name}" else: self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}" - if not os.path.exists(self.file_path) and tp_rank == 0: os.makedirs(self.file_path) logger.info(f"Created HiCacheFile storage directory at {self.file_path}") @@ -211,6 +325,18 @@ class HiCacheFile(HiCacheStorage): def _get_suffixed_key(self, key: str) -> str: return key + self.config_suffix + def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str: + if component_name is None or component_name in ("__default__", PoolName.KV): + return self._get_suffixed_key(key) + return self._get_suffixed_key(f"{key}.{component_name}") + + def _get_component_path( + self, key: str, component_name: Optional[str] = None + ) -> str: + return os.path.join( + self.file_path, f"{self._get_component_key(key, component_name)}.bin" + ) + def get( self, key: str, @@ -280,6 +406,133 @@ class HiCacheFile(HiCacheStorage): tensor_path = os.path.join(self.file_path, f"{key}.bin") return os.path.exists(tensor_path) + def _collect_existing_component_keys( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + ) -> Set[str]: + target_files = {f"{self._get_component_key(key)}.bin" for key in keys} + for transfer in pool_transfers or []: + for key in keys: + target_files.add(f"{self._get_component_key(key, transfer.name)}.bin") + + existing_files = set() + with os.scandir(self.file_path) as entries: + for entry in entries: + if entry.is_file() and entry.name in target_files: + existing_files.add(entry.name) + return existing_files + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + existing_files = self._collect_existing_component_keys(keys, pool_transfers) + + def has_component(page_idx: int, name: str) -> bool: + return ( + f"{self._get_component_key(keys[page_idx], name)}.bin" in existing_files + ) + + # Longest contiguous KV prefix present in storage. + kv_pages = next( + ( + i + for i in range(len(keys)) + if f"{self._get_component_key(keys[i])}.bin" not in existing_files + ), + len(keys), + ) + + hit_count: dict[str, int] = {PoolName.KV: kv_pages} if kv_pages else {} + final_pages = kv_pages + + for transfer in pool_transfers or []: + if final_pages == 0: + break + name = transfer.name + if transfer.hit_policy == PoolHitPolicy.ALL_PAGES: + boundary = next( + (i for i in range(kv_pages) if not has_component(i, name)), kv_pages + ) + else: # trailing_pages + trailing = max(1, len(transfer.keys) if transfer.keys else 1) + boundary = 0 + for prefix_len in range(kv_pages, 0, -1): + if all( + has_component(i, name) + for i in range(max(0, prefix_len - trailing), prefix_len) + ): + boundary = prefix_len + break + if boundary: + hit_count[name] = boundary + final_pages = min(final_pages, boundary) + + return PoolTransferResult(final_pages, hit_count) + + def _log_key(self, pool_name: str, key: str) -> str: + return key if pool_name == PoolName.KV else f"{key}.{pool_name}" + + def _read_page(self, pool_name: str, key: str, host_pool, page_offset: int) -> bool: + """Read one page from storage into host_pool at page_offset.""" + storage_key = self._log_key(pool_name, key) + data_page = self.get(storage_key, host_pool.get_dummy_flat_data_page()) + if data_page is None: + return False + host_pool.set_from_flat_data_page(page_offset, data_page) + return True + + def _write_page( + self, pool_name: str, key: str, host_pool, page_offset: int + ) -> bool: + """Write one page from host_pool at page_offset to storage as raw bytes.""" + storage_key = self._log_key(pool_name, key) + data_page = host_pool.get_data_page(page_offset, flat=True) + return self.set(storage_key, data_page) + + def _batch_io_v2(self, transfers: List[PoolTransfer], op_fn): + results: dict[str, List[bool]] = {} + for transfer in transfers: + host_pool = self.registered_pools[transfer.name] + keys = transfer.keys or [] + page_size = getattr(host_pool, "page_size", 1) or 1 + expected = len(keys) * page_size + host_indices = transfer.host_indices + + if host_indices is None or host_indices.numel() != expected: + logger.error( + "%s indices length mismatch for %s: expected %s, got %s", + op_fn.__name__, + transfer.name, + expected, + host_indices.numel() if host_indices is not None else 0, + ) + results[transfer.name] = [False] * len(keys) + continue + + results[transfer.name] = [ + op_fn(transfer.name, key, host_pool, host_indices[i * page_size].item()) + for i, key in enumerate(keys) + ] + return results + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + return self._batch_io_v2(transfers, self._read_page) + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional["HiCacheStorageExtraInfo"] = None, + ) -> dict[str, List[bool]]: + return self._batch_io_v2(transfers, self._write_page) + def clear(self) -> bool: try: for filename in os.listdir(self.file_path): diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py new file mode 100644 index 000000000..7d6abdab4 --- /dev/null +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -0,0 +1,500 @@ +from __future__ import annotations + +import logging +import threading +import time +from typing import TYPE_CHECKING, Any, List, Optional + +import torch + +from sglang.srt.managers.cache_controller import CacheOperation as BaseCacheOperation +from sglang.srt.managers.cache_controller import ( + HiCacheAck, +) +from sglang.srt.managers.cache_controller import ( + HiCacheController as BaseHiCacheController, +) +from sglang.srt.managers.cache_controller import ( + LayerDoneCounter, +) +from sglang.srt.managers.cache_controller import ( + StorageOperation as BaseStorageOperation, +) +from sglang.srt.mem_cache.hicache_storage import ( + HiCacheStorageExtraInfo, + PoolHitPolicy, + PoolTransfer, + PoolTransferResult, +) +from sglang.srt.mem_cache.memory_pool_host import PoolEntry +from sglang.srt.utils import get_device_module + +if TYPE_CHECKING: + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + +logger = logging.getLogger(__name__) +device_module = get_device_module() + + +class CacheOperation(BaseCacheOperation): + def __init__( + self, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + node_id: int, + priority: Optional[int] = None, + pool_transfers: Optional[list[PoolTransfer]] = None, + ): + super().__init__(host_indices, device_indices, node_id, priority) + self.pool_transfers = pool_transfers + + @staticmethod + def merge_pool_transfers( + ops: List["CacheOperation"], + ) -> Optional[list[PoolTransfer]]: + grouped: dict[str, list[PoolTransfer]] = {} + for op in ops: + for t in op.pool_transfers or []: + grouped.setdefault(t.name, []).append(t) + if not grouped: + return None + + def cat_or_none(tensors): + parts = [x for x in tensors if x is not None] + return torch.cat(parts) if parts else None + + return [ + PoolTransfer( + name=name, + host_indices=cat_or_none(t.host_indices for t in ts), + device_indices=cat_or_none(t.device_indices for t in ts), + keys=[k for t in ts if t.keys for k in t.keys] or None, + ) + for name, ts in grouped.items() + ] + + @staticmethod + def merge_ops(ops: List["CacheOperation"]) -> "CacheOperation": + if len(ops) == 1: + return ops[0] + host_indices = torch.cat([op.host_indices for op in ops]) + device_indices = torch.cat([op.device_indices for op in ops]) + node_ids = [] + priority = min(op.priority for op in ops) + for op in ops: + node_ids.extend(op.node_ids) + merged = CacheOperation( + host_indices, + device_indices, + -1, + priority, + pool_transfers=CacheOperation.merge_pool_transfers(ops), + ) + merged.node_ids = node_ids + return merged + + +class StorageOperation(BaseStorageOperation): + def __init__( + self, + host_indices: torch.Tensor, + token_ids: List[int], + last_hash: Optional[str] = None, + hash_value: Optional[List[str]] = None, + prefix_keys: Optional[List[str]] = None, + pool_transfers: Optional[list[PoolTransfer]] = None, + ): + super().__init__(host_indices, token_ids, last_hash, hash_value, prefix_keys) + self.pool_transfers = pool_transfers + self.pool_storage_result = PoolTransferResult.empty() + + +class PrefetchOperation(StorageOperation): + def __init__( + self, + request_id: str, + host_indices: torch.Tensor, + token_ids: List[int], + last_hash: Optional[str] = None, + prefix_keys: Optional[List[str]] = None, + pool_transfers: Optional[list[PoolTransfer]] = None, + ): + self.request_id = request_id + self._lock = threading.Lock() + self._terminated_flag = False + self.start_time = time.monotonic() + super().__init__( + host_indices, + token_ids, + last_hash, + prefix_keys=prefix_keys, + pool_transfers=pool_transfers, + ) + + def increment(self, num_tokens: int): + with self._lock: + if self._terminated_flag: + return False + self.completed_tokens += num_tokens + return True + + def mark_terminate(self): + with self._lock: + self._terminated_flag = True + + def is_terminated(self) -> bool: + return self._terminated_flag + + +class HybridCacheController(BaseHiCacheController): + def __init__( + self, + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, + mem_pool_host: Any, + page_size: int, + tp_group: torch.distributed.ProcessGroup, + load_cache_event: threading.Event, + write_policy: str = "write_through_selective", + io_backend: str = "", + storage_backend: Optional[str] = None, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[dict] = None, + pp_rank: int = 0, + pp_size: int = 1, + transfer_layer_num: Optional[int] = None, + ): + startup_storage_backend = storage_backend + super().__init__( + token_to_kv_pool_allocator=token_to_kv_pool_allocator, + mem_pool_host=mem_pool_host, + page_size=page_size, + tp_group=tp_group, + load_cache_event=load_cache_event, + write_policy=write_policy, + io_backend=io_backend, + storage_backend=None, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + pp_rank=pp_rank, + pp_size=pp_size, + ) + # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)), + # not just the full attention layers reported by full_kv_pool. + if transfer_layer_num is not None and transfer_layer_num != self.layer_num: + self.layer_num = transfer_layer_num + self.layer_done_counter = LayerDoneCounter(self.layer_num) + + if startup_storage_backend is not None: + self.attach_storage_backend( + storage_backend=startup_storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + host_pools=getattr(mem_pool_host, "entries", None), + ) + + def attach_storage_backend( + self, + storage_backend: str, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[dict] = None, + host_pools: Optional[list[PoolEntry]] = None, + ): + super().attach_storage_backend( + storage_backend=storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + ) + + for entry in host_pools or []: + self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name) + + def reset(self): + super().reset() + if self.enable_storage: + self.host_mem_release_queue.queue.clear() + self.prefetch_tokens_occupied = 0 + + def write( + self, + device_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> Optional[torch.Tensor]: + host_indices = self.mem_pool_host.alloc(len(device_indices)) + if host_indices is None: + return None + pool_transfers = self._resolve_pool_transfers_allocation( + extra_pools, alloc_host=True + ) + if pool_transfers is None and extra_pools: + self.mem_pool_host.free(host_indices) + return None + + self.write_queue.append( + CacheOperation( + host_indices, + device_indices, + node_id, + priority, + pool_transfers=pool_transfers or None, + ) + ) + self.start_writing() + return host_indices + + def start_writing(self) -> None: + if not self.write_queue: + return + op = CacheOperation.merge_ops(self.write_queue) + host_indices, device_indices = self.move_indices(op) + self.write_queue.clear() + start_event = device_module.Event() + finish_event = device_module.Event() + start_event.record() + with device_module.stream(self.write_stream): + start_event.wait(self.write_stream) + self.mem_pool_host.backup_from_device_all_layer( + self.mem_pool_device, + host_indices, + device_indices, + self.io_backend, + pool_transfers=op.pool_transfers, + ) + finish_event.record() + if host_indices.is_cuda: + host_indices.record_stream(self.write_stream) + if device_indices.is_cuda: + device_indices.record_stream(self.write_stream) + self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + + def load( + self, + host_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> Optional[torch.Tensor]: + need_load_kv = host_indices.numel() > 0 + if need_load_kv: + device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) + if device_indices is None: + return None + else: + device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) + + pool_transfers = self._resolve_pool_transfers_allocation( + extra_pools, alloc_host=False + ) + if pool_transfers is None and extra_pools: + if need_load_kv: + self.mem_pool_device_allocator.free(device_indices) + return None + + self.load_queue.append( + CacheOperation( + host_indices, + device_indices, + node_id, + priority, + pool_transfers=pool_transfers or None, + ) + ) + return device_indices + + def start_loading(self) -> int: + if not self.load_queue: + return -1 + producer_id = self.layer_done_counter.update_producer() + op = CacheOperation.merge_ops(self.load_queue) + host_indices, device_indices = self.move_indices(op) + self.load_queue.clear() + producer_event = self.layer_done_counter.events[producer_id] + producer_event.start_event.record() + with device_module.stream(self.load_stream): + producer_event.start_event.wait(self.load_stream) + for i in range(self.layer_num): + self.mem_pool_host.load_to_device_per_layer( + self.mem_pool_device, + host_indices, + device_indices, + i, + self.io_backend, + pool_transfers=op.pool_transfers, + ) + producer_event.complete(i) + if host_indices.is_cuda: + host_indices.record_stream(self.load_stream) + if device_indices.is_cuda: + device_indices.record_stream(self.load_stream) + self.ack_load_queue.append( + HiCacheAck( + producer_event.start_event, + producer_event.finish_event, + op.node_ids, + ) + ) + return producer_id + + def prefetch( + self, + request_id: str, + host_indices: torch.Tensor, + new_input_tokens: List[int], + last_hash: Optional[str] = None, + prefix_keys: Optional[List[str]] = None, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> PrefetchOperation: + operation = PrefetchOperation( + request_id, + host_indices, + new_input_tokens, + last_hash, + prefix_keys=prefix_keys, + pool_transfers=extra_pools, + ) + self.prefetch_queue.put(operation) + return operation + + def write_storage( + self, + host_indices: torch.Tensor, + token_ids: List[int], + hash_value: Optional[List[str]] = None, + prefix_keys: Optional[List[str]] = None, + extra_pools: Optional[list[PoolTransfer]] = None, + ) -> int: + operation = StorageOperation( + host_indices, + token_ids, + hash_value=hash_value, + prefix_keys=prefix_keys, + pool_transfers=extra_pools, + ) + self.backup_queue.put(operation) + return operation.id + + def _storage_hit_query(self, operation) -> tuple[list[str], int]: + last_hash = operation.last_hash + hash_value = [] + for start in range(0, len(operation.token_ids), self.page_size): + last_hash = self.get_hash_str( + operation.token_ids[start : start + self.page_size], last_hash + ) + hash_value.append(last_hash) + + extra_info = HiCacheStorageExtraInfo( + prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None + ) + if operation.pool_transfers: + hit_result = self.storage_backend.batch_exists_v2( + hash_value, operation.pool_transfers, extra_info + ) + else: + kv_hit_count = self.storage_backend.batch_exists(hash_value, extra_info) + hit_result = PoolTransferResult( + kv_hit_pages=kv_hit_count, extra_pool_hit_pages={} + ) + + kv_hit_pages = hit_result.kv_hit_pages + operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages) + + if kv_hit_pages > 0 and operation.pool_transfers: + self._sync_trailing_keys(operation.pool_transfers, hash_value, kv_hit_pages) + + return ( + hash_value[:kv_hit_pages], + kv_hit_pages * self.page_size, + ) + + def _page_transfer(self, operation): + # Transfer extra pools + if operation.pool_transfers and not operation.is_terminated(): + results = self.storage_backend.batch_get_v2(operation.pool_transfers) + operation.pool_storage_result.update_extra_pool_hit_pages(results) + + # Transfer kv pools + super()._page_transfer(operation) + + def _page_backup(self, operation): + # Backup extra pools + if operation.pool_transfers: + results = self.storage_backend.batch_set_v2(operation.pool_transfers) + operation.pool_storage_result.update_extra_pool_hit_pages(results) + + # Backup kv pools + super()._page_backup(operation) + + def _sync_trailing_keys( + self, + pool_transfers: list[PoolTransfer], + all_hashes: list[str], + kv_hit_pages: int, + ) -> None: + """Re-align trailing-page sidecar keys after KV hit truncation. + + When the storage hit is shorter than the original target prefix, each + pool transfer's keys must be updated to the last N hashes of the actual + hit range instead of the last N hashes of the original target range. + For mamba (N=1) this is just the last hit page hash; for SWA (N>1) it + is a sliding window of the last N hit pages. + """ + for transfer in pool_transfers: + if transfer.hit_policy != PoolHitPolicy.TRAILING_PAGES: + continue + trailing_n = len(transfer.keys) if transfer.keys else 1 + transfer.keys = all_hashes[max(0, kv_hit_pages - trailing_n) : kv_hit_pages] + + def _resolve_pool_transfers_allocation( + self, + extra_pools: Optional[list[PoolTransfer]], + alloc_host: bool, + ) -> Optional[list[PoolTransfer]]: + """Auto-alloc host or device indices for PoolTransfers where they are None.""" + if not extra_pools: + return None + newly_allocated: list[tuple[PoolTransfer, Any, torch.Tensor]] = [] + for pool in extra_pools: + entry = self.mem_pool_host.entry_map.get(pool.name) + if entry is None: + continue + if alloc_host: + if pool.host_indices is not None or pool.device_indices is None: + continue + entry_pool, evict_fn, size = ( + entry.host_pool, + entry.host_evict_fn, + len(pool.device_indices), + ) + else: + if pool.device_indices is not None or pool.host_indices is None: + continue + entry_pool, evict_fn, size = ( + entry.device_pool, + entry.device_evict_fn, + len(pool.host_indices), + ) + indices = entry_pool.alloc(size) + if indices is None and evict_fn: + evict_fn(size) + indices = entry_pool.alloc(size) + if indices is None: + # Roll back all previous allocations using each pool's own entry_pool. + for prev_pool, prev_entry_pool, prev_indices in newly_allocated: + prev_entry_pool.free(prev_indices) + if alloc_host: + prev_pool.host_indices = None + else: + prev_pool.device_indices = None + return None + if alloc_host: + pool.host_indices = indices + else: + pool.device_indices = indices + newly_allocated.append((pool, entry_pool, indices)) + return extra_pools diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 147738e55..d02eb7d9b 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -74,6 +74,7 @@ class TreeNode: self.key: RadixKey = None self.value: Optional[torch.Tensor] = None self.mamba_value: Optional[torch.Tensor] = None + self.mamba_host_value: Optional[torch.Tensor] = None # invariant: for any node, if mamba_lock_ref is locked, full_lock_ref must be locked; # if full_lock_ref is locked, mamba_lock_ref doesn't need to be locked. So, # full_lock_ref is always >= mamba_lock_ref. @@ -98,6 +99,8 @@ class TreeNode: self.next = None self.mamba_prev = None self.mamba_next = None + self.host_mamba_prev = None + self.host_mamba_next = None self.id = TreeNode.counter if id is None else id TreeNode.counter += 1 @@ -106,10 +109,18 @@ class TreeNode: def evicted(self): return self.value is None + @property + def mamba_evicted(self): + return self.mamba_value is None + @property def backuped(self): return self.host_value is not None + @property + def mamba_backuped(self): + return self.mamba_host_value is not None + def protect_host(self): """Protect the host value from eviction.""" self.host_ref_counter += 1 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 51530e865..0e2f63cb4 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -475,6 +475,9 @@ class HybridReqToTokenPool(ReqToTokenPool): self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1 self.enable_mamba_extra_buffer = enable_mamba_extra_buffer self.enable_memory_saver = enable_memory_saver + # TODO: Support PP + self.start_layer = 0 + self.layer_transfer_counter = None self._init_mamba_pool( size=mamba_size, mamba_spec_state_size=mamba_spec_state_size, @@ -516,6 +519,11 @@ class HybridReqToTokenPool(ReqToTokenPool): ) ) + def register_layer_transfer_counter( + self, layer_transfer_counter: "LayerDoneCounter" + ): + self.layer_transfer_counter = layer_transfer_counter + # For chunk prefill req, we do not need to allocate mamba cache, # We could use allocated mamba cache instead. def alloc(self, reqs: List["Req"]) -> Optional[List[int]]: @@ -570,6 +578,8 @@ class HybridReqToTokenPool(ReqToTokenPool): def mamba2_layer_cache(self, layer_id: int): assert layer_id in self.mamba_map + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) return self.mamba_pool.mamba2_layer_cache(self.mamba_map[layer_id]) def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState: @@ -1238,8 +1248,8 @@ class HybridLinearKVPool(KVCache): self.device = device self.full_layer_nums = len(full_attention_layer_ids) self.page_size = page_size - # TODO support pp? - self.start_layer = 0 + self.start_layer = 0 # TODO: Support PP + self.layer_transfer_counter = None self.head_num = head_num self.head_dim = head_dim self.mamba_pool = mamba_pool @@ -1323,15 +1333,30 @@ class HybridLinearKVPool(KVCache): ) return self.full_attention_layer_id_mapping[layer_id] + def register_layer_transfer_counter( + self, layer_transfer_counter: "LayerDoneCounter" + ): + self.layer_transfer_counter = layer_transfer_counter + # The layer-wise wait logic is executed at the Hybrid LinearPool level; + # no additional wait is needed in the full_kv_pool + self.full_kv_pool.register_layer_transfer_counter(None) + + def _wait_for_layer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + def get_key_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_key_buffer(layer_id) def get_value_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_value_buffer(layer_id) def get_kv_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_kv_buffer(layer_id) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index ea220bef0..10fc35239 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1,10 +1,17 @@ +from __future__ import annotations + import abc import logging import threading from collections import defaultdict +from dataclasses import dataclass from functools import wraps -from typing import Optional +from typing import TYPE_CHECKING, Any, Callable, Optional +if TYPE_CHECKING: + from sglang.srt.mem_cache.hicache_storage import PoolName + +import numpy as np import psutil import torch @@ -19,6 +26,7 @@ from sglang.jit_kernel.hicache import ( ) from sglang.srt.mem_cache.memory_pool import ( KVCache, + MambaPool, MHATokenToKVPool, MLATokenToKVPool, NSATokenToKVPool, @@ -1074,6 +1082,526 @@ class MLATokenToKVPoolHost(HostKVCache): return ptr_list, element_size_list +class MambaPoolHost(HostKVCache): + + def __init__( + self, + device_pool: MambaPool, + host_to_device_ratio: float, + host_size: int, + pin_memory: bool = True, + device: str = "cpu", + allocator_type: str = "default", + layout: str = "layer_first", + ): + self.device_pool = device_pool + self.page_size = 1 + assert layout in [ + "page_first", + "page_first_direct", + "layer_first", + ], "Unsupported layout: {layout}" + + self.layout = layout + self.pin_memory = pin_memory + self.device = device + self.allocator = get_allocator_from_storage(allocator_type) + self.num_mamba_layers = device_pool.num_mamba_layers + + self.conv_state_shapes = [ + conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv + ] + self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:] + self.conv_dtype = device_pool.mamba_cache.conv[0].dtype + self.temporal_dtype = device_pool.mamba_cache.temporal.dtype + self.dtype = self.conv_dtype + self.size_per_token = self.get_size_per_token() + + if host_size > 0: + self.size = int(host_size * 1e9 // self.size_per_token) + else: + self.size = int(device_pool.size * host_to_device_ratio) + + self.page_num = self.size // self.page_size + 1 + self.size = self.page_num * self.page_size + + assert ( + self.size > device_pool.size + ), "The host memory should be larger than the device memory with the current protocol" + + host_mem = psutil.virtual_memory() + requested_bytes = self.size * self.size_per_token + ten_gb = 10 * (1024**3) + available_bytes = host_mem.available - ten_gb + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory available. Requesting " + f"{requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free. Please reduce the " + f"size of the hierarchical cache." + ) + logger.info( + "Allocating %.2f GB host memory for hierarchical Mamba cache (layout=%s).", + requested_bytes / 1e9, + self.layout, + ) + + self.init_kv_buffer() + self.lock = threading.RLock() + self.clear() + + def init_kv_buffer(self): + alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] + + if self.layout in ["page_first", "page_first_direct"]: + # page-first: (page_num, num_layers, 1, *shape) — per-page data is contiguous + temporal_dims = ( + self.size, + self.num_mamba_layers, + 1, + ) + self.temporal_state_shape + self.temporal_buffer = alloc_func( + temporal_dims, + dtype=self.temporal_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + self.conv_buffer = [] + for conv_shape in self.conv_state_shapes: + conv_dims = (self.size, self.num_mamba_layers, 1) + conv_shape + self.conv_buffer.append( + alloc_func( + conv_dims, + dtype=self.conv_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + ) + else: + # layer-first: (num_layers, size, *shape) + temporal_dims = ( + self.num_mamba_layers, + self.size, + ) + self.temporal_state_shape + self.temporal_buffer = alloc_func( + temporal_dims, + dtype=self.temporal_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + self.conv_buffer = [] + for conv_shape in self.conv_state_shapes: + conv_dims = (self.num_mamba_layers, self.size) + conv_shape + self.conv_buffer.append( + alloc_func( + conv_dims, + dtype=self.conv_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + ) + + def _iter_page_tensors(self, index: int): + if self.layout in ["page_first", "page_first_direct"]: + yield self.temporal_buffer[index] + for conv_buf in self.conv_buffer: + yield conv_buf[index] + else: + yield self.temporal_buffer[:, index : index + self.page_size] + for conv_buf in self.conv_buffer: + yield conv_buf[:, index : index + self.page_size] + + @staticmethod + def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor: + return tensor.contiguous().view(torch.uint8).reshape(-1) + + @synchronized + def clear(self): + self.mem_state = torch.zeros( + (self.size,), dtype=torch.uint8, device=self.device + ) + self.free_slots = torch.arange(self.size, dtype=torch.int64) + + def available_size(self): + return len(self.free_slots) + + @synchronized + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + assert ( + need_size % self.page_size == 0 + ), "The requested size should be a multiple of the page size." + if need_size > self.available_size(): + return None + select_index = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + return select_index + + @synchronized + def free(self, indices: torch.Tensor) -> int: + self.free_slots = torch.cat([self.free_slots, indices]) + return len(indices) + + def get_size_per_token(self): + conv_total_size = 0 + for conv_shape in self.conv_state_shapes: + conv_total_size += int(np.prod(conv_shape)) * self.conv_dtype.itemsize + temporal_size = ( + int(np.prod(self.temporal_state_shape)) * self.temporal_dtype.itemsize + ) + return (conv_total_size + temporal_size) * self.num_mamba_layers + + def get_ksize_per_token(self): + return self.get_size_per_token() + + @staticmethod + def _item_size_per_index(tensor: torch.Tensor) -> int: + if tensor.shape[0] == 0: + return 0 + return int(tensor[0].numel() * tensor.element_size()) + + @staticmethod + def _copy_tensor( + src: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + io_backend: str, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + # TODO: Rename the interface for clarity. + # Here, transfer_kv_per_layer_mla is reused to transfer the Mamba state. + # This has nothing to do with MLA; it's only reused because this interface happens to transfer a single Pool. + transfer_kv_per_layer_mla( + src=src, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=MambaPoolHost._item_size_per_index(src), + ) + elif io_backend == "direct": + transfer_kv_direct( + src_layers=[src], + dst_layers=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + + @staticmethod + def _copy_tensor_pf_lf( + src: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + layer_id: int, + num_layers: int, + io_backend: str, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + item_size = MambaPoolHost._item_size_per_index(dst) + transfer_kv_per_layer_mla_pf_lf( + src=src, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + layer_id=layer_id, + item_size=item_size, + src_layout_dim=item_size * num_layers, + ) + elif io_backend == "direct": + transfer_kv_per_layer_direct_pf_lf( + src_ptrs=[src], + dst_ptrs=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + layer_id=layer_id, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + + @staticmethod + def _copy_tensor_all_layers_lf_pf( + src_layers: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + num_layers: int, + device: str, + io_backend: str, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + item_size = MambaPoolHost._item_size_per_index(src_layers[0]) + src_ptrs = torch.tensor( + [src_layers[i].data_ptr() for i in range(num_layers)], + dtype=torch.uint64, + device=device, + ) + transfer_kv_all_layer_mla_lf_pf( + src_layers=src_ptrs, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=item_size, + dst_layout_dim=item_size * num_layers, + num_layers=num_layers, + ) + elif io_backend == "direct": + src_ptrs = [src_layers[i] for i in range(num_layers)] + transfer_kv_all_layer_direct_lf_pf( + src_ptrs=src_ptrs, + dst_ptrs=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + + def load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend="kernel", + ): + if self.layout in ["page_first", "page_first_direct"]: + self._copy_tensor_pf_lf( + src=self.temporal_buffer, + dst=device_pool.mamba_cache.temporal[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor_pf_lf( + src=self.conv_buffer[conv_idx], + dst=device_pool.mamba_cache.conv[conv_idx][layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) + else: + self._copy_tensor( + self.temporal_buffer[layer_id], + device_pool.mamba_cache.temporal[layer_id], + host_indices, + device_indices, + io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + self.conv_buffer[conv_idx][layer_id], + device_pool.mamba_cache.conv[conv_idx][layer_id], + host_indices, + device_indices, + io_backend, + ) + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend="kernel" + ): + if self.layout in ["page_first", "page_first_direct"]: + self._copy_tensor_all_layers_lf_pf( + src_layers=device_pool.mamba_cache.temporal, + dst=self.temporal_buffer, + src_indices=device_indices, + dst_indices=host_indices, + num_layers=self.num_mamba_layers, + device=self.device_pool.device, + io_backend=io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor_all_layers_lf_pf( + src_layers=device_pool.mamba_cache.conv[conv_idx], + dst=self.conv_buffer[conv_idx], + src_indices=device_indices, + dst_indices=host_indices, + num_layers=self.num_mamba_layers, + device=self.device_pool.device, + io_backend=io_backend, + ) + else: + for layer_id in range(self.num_mamba_layers): + self._copy_tensor( + device_pool.mamba_cache.temporal[layer_id], + self.temporal_buffer[layer_id], + device_indices, + host_indices, + io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + device_pool.mamba_cache.conv[conv_idx][layer_id], + self.conv_buffer[conv_idx][layer_id], + device_indices, + host_indices, + io_backend, + ) + + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + data_page = torch.cat( + [ + self._flatten_tensor_bytes(tensor) + for tensor in self._iter_page_tensors(index) + ] + ) + return data_page.flatten() if flat else data_page + + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros( + self.page_size * self.size_per_token, + dtype=torch.uint8, + device=self.device, + pin_memory=self.pin_memory, + ) + + def set_from_flat_data_page( + self, + index: int, + data_page: torch.Tensor, + ) -> None: + flat_bytes = data_page.contiguous().view(torch.uint8).reshape(-1) + start = 0 + for tensor in self._iter_page_tensors(index): + num_bytes = tensor.numel() * tensor.element_size() + tensor_bytes = flat_bytes[start : start + num_bytes] + start += num_bytes + restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape) + tensor.copy_(restored) + + +@dataclass +class PoolEntry: + name: PoolName + host_pool: Any + device_pool: Any + layer_mapper: Callable[[int], Optional[int]] + is_primary_index_anchor: bool = False + # Optional eviction callbacks for auto-alloc in HybridCacheController. + # host_evict_fn(n): evict n slots from the host pool (used by write()). + # device_evict_fn(n): evict n slots from the device pool (used by load()). + host_evict_fn: Optional[Callable] = None + device_evict_fn: Optional[Callable] = None + + +class HostPoolGroup: + def __init__(self, entries: list[PoolEntry]): + if not entries: + raise ValueError("HostPoolGroup requires at least one pool entry.") + self.entries = entries + self.entry_map = {entry.name: entry for entry in entries} + self.anchor_entry = next( + (entry for entry in entries if entry.is_primary_index_anchor), + entries[0], + ) + + self.layout = self.anchor_entry.host_pool.layout + self.page_size = self.anchor_entry.host_pool.page_size + self.device = self.anchor_entry.host_pool.device + self.size = self.anchor_entry.host_pool.size + + def clear(self) -> None: + for entry in self.entries: + entry.host_pool.clear() + + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + return self.anchor_entry.host_pool.alloc(need_size) + + def free(self, indices: torch.Tensor) -> int: + return self.anchor_entry.host_pool.free(indices) + + def get_data_page(self, index, flat: bool = True): + return self.anchor_entry.host_pool.get_data_page(index, flat) + + def get_dummy_flat_data_page(self): + return self.anchor_entry.host_pool.get_dummy_flat_data_page() + + def set_from_flat_data_page(self, index: int, data_page) -> None: + return self.anchor_entry.host_pool.set_from_flat_data_page(index, data_page) + + def load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + pool_transfers: Optional[list] = None, + ) -> None: + # 1. Anchor (KV) transfer + anchor = self.anchor_entry + local_layer_id = anchor.layer_mapper(layer_id) + if local_layer_id is not None and host_indices.numel() > 0: + anchor.host_pool.load_to_device_per_layer( + anchor.device_pool, + host_indices, + device_indices, + local_layer_id, + io_backend, + ) + + # 2. Extra pool transfers + for transfer in pool_transfers or []: + entry = self.entry_map.get(transfer.name) + if entry is None or transfer.host_indices is None: + continue + local_layer_id = entry.layer_mapper(layer_id) + if local_layer_id is None: + continue + entry.host_pool.load_to_device_per_layer( + entry.device_pool, + transfer.host_indices, + transfer.device_indices, + local_layer_id, + io_backend, + ) + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + pool_transfers: Optional[list] = None, + ) -> None: + # 1. Anchor (KV) backup + self.anchor_entry.host_pool.backup_from_device_all_layer( + self.anchor_entry.device_pool, + host_indices, + device_indices, + io_backend, + ) + # 2. Extra pool backup + for transfer in pool_transfers or []: + entry = self.entry_map.get(transfer.name) + if entry is None or transfer.host_indices is None: + continue + entry.host_pool.backup_from_device_all_layer( + entry.device_pool, + transfer.host_indices, + transfer.device_indices, + io_backend, + ) + + class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device_pool: NSATokenToKVPool diff --git a/test/registered/4-gpu-models/test_qwen35_models.py b/test/registered/4-gpu-models/test_qwen35_models.py index f23a111fa..cfd1626f0 100644 --- a/test/registered/4-gpu-models/test_qwen35_models.py +++ b/test/registered/4-gpu-models/test_qwen35_models.py @@ -1,5 +1,6 @@ import shutil import tempfile +import time import unittest from types import SimpleNamespace @@ -254,6 +255,10 @@ class TestQwen35WithHiCache(CustomTestCase): other_args=[ "--tp-size", "4", + "--max-mamba-cache-size", + "500", + "--max-total-tokens", + "120000", "--chunked-prefill-size", "2048", "--mamba-scheduler-strategy", @@ -289,14 +294,14 @@ class TestQwen35WithHiCache(CustomTestCase): kill_process_tree(cls.process.pid) shutil.rmtree(cls.storage_dir, ignore_errors=True) - def test_gsm8k(self): + def _run_gsm8k(self): args = SimpleNamespace( model=self.model, eval_name="gsm8k", num_shots=5, - num_examples=200, - max_tokens=16000, - num_threads=128, + num_examples=100, + max_tokens=4000, + num_threads=50, repeat=1, temperature=0.6, top_p=0.95, @@ -305,9 +310,30 @@ class TestQwen35WithHiCache(CustomTestCase): host="http://127.0.0.1", port=int(self.base_url.split(":")[-1]), ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]) + return run_eval(args) + + def test_gsm8k(self): + first_metrics = self._run_gsm8k() + print(f"first_metrics={first_metrics}") + self.assertGreaterEqual( + first_metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + + print(f"flush cache") + time.sleep(2) + requests.post(f"{self.base_url}/flush_cache", timeout=10) + + second_metrics = self._run_gsm8k() + print(f"second_metrics={second_metrics}") + self.assertGreaterEqual( + second_metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + self.assertLessEqual( + abs(second_metrics["score"] - first_metrics["score"]), + 0.05, + f"HiCache prefetch accuracy drift too large: " + f"first={first_metrics['score']}, second={second_metrics['score']}", + ) if __name__ == "__main__":