From c0f59506364f8f16527b574c06a9d4c5f58ce99a Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Sun, 3 May 2026 22:13:22 +0800 Subject: [PATCH] [UnifiedRadixTree]: Support HiCache Framework for UnifiedRadixTree (#23316) Co-authored-by: JINZ <1023553676@qq.com> Co-authored-by: diemchai --- python/sglang/srt/managers/scheduler.py | 35 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 2 +- .../sglang/srt/mem_cache/hicache_storage.py | 2 + .../hybrid_cache/hybrid_pool_assembler.py | 50 +- .../unified_cache_components/__init__.py | 4 + .../full_component.py | 206 ++- .../mamba_component.py | 215 +++- .../unified_cache_components/swa_component.py | 43 +- .../tree_component.py | 44 +- .../srt/mem_cache/unified_radix_cache.py | 1137 ++++++++++++++--- .../test_unified_radix_cache_kl.py | 50 + .../test_unified_radix_cache_unittest.py | 659 +++++++++- 12 files changed, 2193 insertions(+), 254 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c1f79bb0d..dc44902e4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -892,6 +892,26 @@ class Scheduler( logger.info("Using experimental C++ radix tree implementation.") self.tree_cache = RadixCacheCpp(params=params, server_args=server_args) + elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get(): + from sglang.srt.mem_cache.unified_cache_components import ( + ComponentType, + ) + from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, + ) + + tree_components = [ComponentType.FULL] + if self.is_hybrid_swa or self.is_hybrid_ssm: + tree_components.append( + ComponentType.SWA if self.is_hybrid_swa else ComponentType.MAMBA + ) + params.tree_components = tuple(tree_components) + self.tree_cache = UnifiedRadixCache(params) + if self.enable_hierarchical_cache: + self.tree_cache.init_hicache(server_args, params) + self.tp_worker.register_hicache_layer_transfer_counter( + self.tree_cache.cache_controller.layer_done_counter + ) elif self.enable_hierarchical_cache: if self.is_hybrid_ssm: from sglang.srt.mem_cache.hi_mamba_radix_cache import ( @@ -910,21 +930,6 @@ class Scheduler( self.tp_worker.register_hicache_layer_transfer_counter( self.tree_cache.cache_controller.layer_done_counter ) - elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get(): - from sglang.srt.mem_cache.unified_cache_components import ( - ComponentType, - ) - from sglang.srt.mem_cache.unified_radix_cache import ( - UnifiedRadixCache, - ) - - tree_components = [ComponentType.FULL] - if self.is_hybrid_swa or self.is_hybrid_ssm: - tree_components.append( - ComponentType.SWA if self.is_hybrid_swa else ComponentType.MAMBA - ) - params.tree_components = tuple(tree_components) - self.tree_cache = UnifiedRadixCache(params) elif self.is_hybrid_swa: from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 403dc6ed4..999f46f87 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -74,7 +74,7 @@ class InsertResult: class EvictParams: """Unified parameters for evict across different cache types""" - num_tokens: int + num_tokens: int = 0 swa_num_tokens: int = 0 mamba_num: int = 0 diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 29b72ae3d..21bb4a9ec 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -64,6 +64,7 @@ class PoolTransfer: device<->host path : host_indices + device_indices host<->storage path: host_indices + keys + nodes_to_load : evicted nodes this transfer covers """ name: PoolName @@ -71,6 +72,7 @@ class PoolTransfer: device_indices: Optional[torch.Tensor] = None keys: Optional[List[str]] = None hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES + nodes_to_load: Optional[List[Any]] = None @dataclass diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 2b1ec8aac..d88ba3da2 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any, Callable, Optional -from sglang.srt.mem_cache.hicache_storage import PoolName +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, ) @@ -325,7 +325,11 @@ def attach_hybrid_pool_to_unified_cache( ) -> None: """Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache.""" from sglang.srt.mem_cache.base_prefix_cache import EvictParams - from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MLATokenToKVPool + from sglang.srt.mem_cache.memory_pool import ( + HybridLinearKVPool, + MLATokenToKVPool, + NSATokenToKVPool, + ) from sglang.srt.mem_cache.unified_cache_components import ComponentType try: @@ -345,6 +349,7 @@ def attach_hybrid_pool_to_unified_cache( }, "Non-hybrid KV pool currently only supports FULL-only UnifiedRadixCache." mamba_stack = isinstance(kvcache, HybridLinearKVPool) + nsa_stack = isinstance(kvcache, NSATokenToKVPool) if mamba_stack: full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping) mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) @@ -381,6 +386,45 @@ def attach_hybrid_pool_to_unified_cache( cache_controller.layer_done_counter ) transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping) + elif nsa_stack: + full_layer_mapping = { + layer_id: layer_id for layer_id in range(full_kv_pool.layer_num) + } + host_pool_group, cache_controller = build_shared_anchor_stack( + params=params, + server_args=server_args, + kv_pool=full_kv_pool, + shared_pool_name=PoolName.INDEXER, + full_layer_mapping=full_layer_mapping, + page_size=cache.page_size, + tp_group=params.tp_cache_group, + load_cache_event=load_cache_event, + storage_backend=None, + use_mla=use_mla, + shared_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( + full_kv_pool, + kv_host_pool, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ), + pp_rank=params.pp_rank, + pp_size=params.pp_size, + attn_cp_rank=params.attn_cp_rank, + attn_cp_size=params.attn_cp_size, + ) + cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) + cache.host_pool_group = host_pool_group + cache.cache_controller = cache_controller + # Register the NSA indexer pool as sharing anchor-KV indices so + # HiCache backup/load emits its PoolTransfer together with KV. + cache.register_hicache_anchor_kv_shared_indices_pool( + PoolName.INDEXER, + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + cache.components[ComponentType.FULL]._full_kv_pool_host = ( + cache.full_kv_pool_host + ) + transfer_layer_num = len(full_layer_mapping) else: full_layer_mapping = { layer_id: layer_id for layer_id in range(full_kv_pool.layer_num) @@ -414,7 +458,7 @@ def attach_hybrid_pool_to_unified_cache( logger.info( "Attached hybrid pool stack to UnifiedRadixCache: pools=%s, transfer_layer_num=%s", - "KV + MAMBA" if mamba_stack else "KV", + "KV + MAMBA" if mamba_stack else "KV + INDEXER" if nsa_stack else "KV", transfer_layer_num, ) except Exception: diff --git a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py b/python/sglang/srt/mem_cache/unified_cache_components/__init__.py index d0fde2786..2d848b3ec 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/__init__.py @@ -4,8 +4,10 @@ from sglang.srt.mem_cache.unified_cache_components.swa_component import SWACompo from sglang.srt.mem_cache.unified_cache_components.tree_component import ( _NUM_COMPONENT_TYPES, BASE_COMPONENT_TYPE, + CacheTransferPhase, ComponentData, ComponentType, + EvictLayer, TreeComponent, get_and_increase_time_counter, next_component_uuid, @@ -15,7 +17,9 @@ __all__ = [ "BASE_COMPONENT_TYPE", "ComponentData", "ComponentType", + "EvictLayer", "FullComponent", + "CacheTransferPhase", "MambaComponent", "SWAComponent", "TreeComponent", diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index d963775d1..c9470f66f 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py @@ -1,14 +1,22 @@ from __future__ import annotations +import heapq from typing import TYPE_CHECKING, Callable, Optional +import torch + from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, EvictParams, IncLockRefResult, + MatchPrefixParams, + MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + CacheTransferPhase, ComponentType, + EvictLayer, TreeComponent, ) @@ -30,27 +38,80 @@ class FullComponent(TreeComponent): self._free_full = allocator.full_attn_allocator.free else: self._free_full = allocator.free - - def node_has_component_data(self, node: UnifiedTreeNode) -> bool: - # Override so _for_each_component_lru includes Full in LRU operations - return node.component_data[self.component_type].value is not None + # HiCache state: set to host KV pool when HiCache enabled + self._full_kv_pool_host = None def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]: - return lambda node: True + # HiCache: evicted + backuped nodes are valid match boundaries + return lambda node: ( + node.component_data[self.component_type].value is not None or node.backuped + ) + + def finalize_match_result( + self, + result: MatchResult, + params: MatchPrefixParams, + value_chunks: list[torch.Tensor], + best_value_len: int, + ) -> MatchResult: + # Compute Full KV host hit length: walk from last_host_node up to + # last_device_node, summing host_value lengths of evicted nodes. + ct = self.component_type + kv_host_hit = 0 + node = result.last_host_node + root_node = self.cache.root_node + while node is not result.last_device_node and node is not root_node: + full_host = node.component_data[ct].host_value + if full_host is not None: + kv_host_hit += len(full_host) + node = node.parent + if kv_host_hit > 0: + return result._replace( + host_hit_length=max(result.host_hit_length, kv_host_hit) + ) + return result def redistribute_on_node_split( self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode ): - new_parent.component_data[self.component_type].lock_ref = child.component_data[ - self.component_type - ].lock_ref + ct = self.component_type + new_parent.component_data[ct].lock_ref = child.component_data[ct].lock_ref + child_cd = child.component_data[ct] + split_len = len(new_parent.key) + if child_cd.value is not None: + new_parent.component_data[ct].value = child_cd.value[:split_len].clone() + child_cd.value = child_cd.value[split_len:].clone() + if child_cd.host_value is not None: + new_parent.component_data[ct].host_value = child_cd.host_value[ + :split_len + ].clone() + child_cd.host_value = child_cd.host_value[split_len:].clone() - def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: + def evict_component( + self, + node: UnifiedTreeNode, + target: EvictLayer = EvictLayer.DEVICE, + ) -> tuple[int, int]: cd = node.component_data[self.component_type] - self._free_full(cd.value) - freed = len(cd.value) - self.cache.component_evictable_size_[self.component_type] -= freed - return freed + freed = 0 + host_freed = 0 + + # Device layer + if EvictLayer.DEVICE in target and cd.value is not None: + self._free_full(cd.value) + freed = len(cd.value) + self.cache.component_evictable_size_[self.component_type] -= freed + # NOTE: cd.value = None is deferred to _cascade_evict (Full as trigger) + # because SWA's free_swa still needs to read Full.value. + # cd.value = None + + # Host layer + if EvictLayer.HOST in target and cd.host_value is not None: + host_freed = len(cd.host_value) + if self._full_kv_pool_host is not None: + self._full_kv_pool_host.free(cd.host_value) + cd.host_value = None + return freed, host_freed def eviction_priority(self, is_leaf: bool) -> int: return 0 if is_leaf else 2 @@ -59,30 +120,55 @@ class FullComponent(TreeComponent): self, params: EvictParams, tracker: dict[ComponentType, int] ) -> None: request = params.num_tokens - lru = self.cache.lru_lists[self.component_type] - while tracker[self.component_type] < request: - x = lru.get_leaf_lru_no_lock() - if x is None: - break - self.cache._evict_component_and_detach_lru( - x, self, is_leaf=True, tracker=tracker - ) - self.cache._cascade_evict(x, self, tracker) + # Heap-based eviction from evictable_device_leaves, ordered by LRU. + heap = [(n.last_access_time, n) for n in self.cache.evictable_device_leaves] + heapq.heapify(heap) + ct = self.component_type + while tracker[ct] < request and heap: + _, x = heapq.heappop(heap) + if x not in self.cache.evictable_device_leaves: + continue + self.cache._evict_device_leaf(x, tracker) + if x.parent is not None and x.parent in self.cache.evictable_device_leaves: + heapq.heappush(heap, (x.parent.last_access_time, x.parent)) + + def drive_host_eviction( + self, num_tokens: int, tracker: dict[ComponentType, int] + ) -> None: + """Evict host leaves to free KV host pool space.""" + heap = [(n.last_access_time, n) for n in self.cache.evictable_host_leaves] + heapq.heapify(heap) + ct = self.component_type + while tracker[ct] < num_tokens and heap: + _, x = heapq.heappop(heap) + if x not in self.cache.evictable_host_leaves: + continue + self.cache._evict_host_leaf(x, tracker) + if x.parent is not None and x.parent in self.cache.evictable_host_leaves: + heapq.heappush(heap, (x.parent.last_access_time, x.parent)) def acquire_component_lock( self, node: UnifiedTreeNode, result: IncLockRefResult ) -> IncLockRefResult: ct = self.component_type root = self.cache.root_node + delta = 0 cur = node while cur != root: cd = cur.component_data[ct] + assert cd.value is not None + if cd.lock_ref == 0: key_len = len(cd.value) self.cache.component_evictable_size_[ct] -= key_len self.cache.component_protected_size_[ct] += key_len + delta += key_len cd.lock_ref += 1 + self.cache.evictable_device_leaves.discard(cur) cur = cur.parent + result = IncLockRefResult( + delta=delta, swa_uuid_for_lock=result.swa_uuid_for_lock + ) return result def release_component_lock( @@ -93,10 +179,86 @@ class FullComponent(TreeComponent): cur = node while cur != root: cd = cur.component_data[ct] + assert cd.value is not None assert cd.lock_ref > 0 + if cd.lock_ref == 1: key_len = len(cd.value) self.cache.component_evictable_size_[ct] += key_len self.cache.component_protected_size_[ct] -= key_len cd.lock_ref -= 1 + if cd.lock_ref == 0: + self.cache._update_evictable_leaf_sets(cur) cur = cur.parent + + # ---- HiCache Hooks ---- + + def build_hicache_transfers( + self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw + ) -> Optional[list[PoolTransfer]]: + ct = self.component_type + + if phase == CacheTransferPhase.BACKUP_HOST: + # Full KV backup is handled by the main flow + # (write_backup → cache_controller.write on host_value directly). + # No extra PoolTransfer needed. + return None + + if phase == CacheTransferPhase.LOAD_BACK: + # Walk evicted chain, collect host_values and nodes + backed_up: list[torch.Tensor] = [] + nodes: list = [] + cur = node + while cur.evicted: + cd = cur.component_data[ct] + if cd.host_value is not None: + backed_up.append(cd.host_value) + nodes.append(cur) + cur = cur.parent + backed_up.reverse() + nodes.reverse() + return [ + PoolTransfer( + name=PoolName.KV, + host_indices=( + torch.cat(backed_up) + if backed_up + else torch.empty((0,), dtype=torch.int64, device="cpu") + ), + device_indices=None, + nodes_to_load=nodes, + ) + ] + + return None + + def commit_hicache_transfer( + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + transfers: list[PoolTransfer] = (), + ) -> None: + ct = self.component_type + + if phase == CacheTransferPhase.BACKUP_HOST: + if transfers and transfers[0].host_indices is not None: + node.component_data[ct].host_value = transfers[0].host_indices.clone() + + elif phase == CacheTransferPhase.LOAD_BACK: + if not transfers or transfers[0].device_indices is None: + self.cache._update_evictable_leaf_sets(node) + return + + xfer = transfers[0] + device_indices = xfer.device_indices + offset = 0 + for n in xfer.nodes_to_load or []: + cd = n.component_data[ct] + n_len = len(cd.host_value) + cd.value = device_indices[offset : offset + n_len].clone() + offset += n_len + # Full uses leaf sets, not LRU + self.cache.component_evictable_size_[ct] += n_len + self.cache._update_evictable_leaf_sets(n) + + self.cache._update_evictable_leaf_sets(node) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index 032e905dd..9f75352e5 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -13,8 +13,11 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + CacheTransferPhase, ComponentType, + EvictLayer, TreeComponent, get_and_increase_time_counter, ) @@ -44,10 +47,16 @@ class MambaComponent(TreeComponent): ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {cache.page_size}" super().__init__(cache, params) self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer + # HiCache state + self._mamba_pool_host = None # set to host mamba pool when HiCache enabled def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]: ct = self.component_type - return lambda node: node.component_data[ct].value is not None + # HiCache: evicted + backuped (host_value present) is also a valid match + return lambda node: ( + node.component_data[ct].value is not None + or node.component_data[ct].host_value is not None + ) def finalize_match_result( self, @@ -90,6 +99,13 @@ class MambaComponent(TreeComponent): mamba_value, dst_index ) + # HiCache: if mamba was evicted from device but has host backup, + # ensure host_hit_length >= 1 so load_back is triggered. + host_node = result.last_host_node + cd = host_node.component_data[self.component_type] + if cd.value is None and cd.host_value is not None: + result = result._replace(host_hit_length=max(result.host_hit_length, 1)) + return result._replace(mamba_branching_seqlen=branching_seqlen) def commit_insert_component_data( @@ -109,6 +125,10 @@ class MambaComponent(TreeComponent): return if node.component_data[self.component_type].value is None: node.component_data[self.component_type].value = params.mamba_value + # move from host LRU to device LRU + host_lru = self.cache.host_lru_lists[self.component_type] + if host_lru.in_list(node): + host_lru.remove_node(node) self.cache.lru_lists[self.component_type].insert_mru(node) self.cache.component_evictable_size_[self.component_type] += len( params.mamba_value @@ -122,41 +142,74 @@ class MambaComponent(TreeComponent): def redistribute_on_node_split( self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode ): - new_parent.component_data[self.component_type].value = None - new_parent.component_data[self.component_type].lock_ref = 0 + ct = self.component_type + new_parent.component_data[ct].value = None + new_parent.component_data[ct].lock_ref = 0 + # HiCache: mamba host_value stays on child (mamba = leaf-only data) + new_parent.component_data[ct].host_value = None + new_parent.component_data[ct].host_lock_ref = 0 - def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: - value = node.component_data[self.component_type].value - self.cache.req_to_token_pool.mamba_pool.free(value) - freed = len(value) - self.cache.component_evictable_size_[self.component_type] -= freed - if not is_leaf: - node.component_data[self.component_type].value = None - return freed + def evict_component( + self, + node: UnifiedTreeNode, + target: EvictLayer = EvictLayer.DEVICE, + ) -> tuple[int, int]: + cd = node.component_data[self.component_type] + freed = 0 + host_freed = 0 + + # Device layer + if EvictLayer.DEVICE in target and cd.value is not None: + self.cache.req_to_token_pool.mamba_pool.free(cd.value) + freed = len(cd.value) + self.cache.component_evictable_size_[self.component_type] -= freed + cd.value = None + + # Host layer + host_lru = self.cache.host_lru_lists[self.component_type] + if EvictLayer.HOST in target and cd.host_value is not None: + host_freed = len(cd.host_value) + if self._mamba_pool_host is not None: + self._mamba_pool_host.free(cd.host_value) + cd.host_value = None + if host_lru.in_list(node): + host_lru.remove_node(node) + + # After device tombstone: if only host_value remains, insert into host LRU + if ( + target is EvictLayer.DEVICE + and cd.value is None + and cd.host_value is not None + ): + if not host_lru.in_list(node): + host_lru.insert_mru(node) + + return freed, host_freed def drive_eviction( self, params: EvictParams, tracker: dict[ComponentType, int] ) -> None: request = params.mamba_num - lru = self.cache.lru_lists[self.component_type] + ct = self.component_type + lru = self.cache.lru_lists[ct] x = lru.get_lru_no_lock() - while ( - tracker[self.component_type] < request and x is not None and lru.in_list(x) - ): - assert x.component_data[self.component_type].value is not None - if len(x.children) > 0: + while tracker[ct] < request and x is not None and lru.in_list(x): + assert x.component_data[ct].value is not None + if x in self.cache.evictable_device_leaves: + # D-leaf: atomic eviction of all components + x_next = lru.get_prev_no_lock(x) + self.cache._evict_device_leaf(x, tracker) + if not lru.in_list(x_next): + x_next = lru.get_lru_no_lock() + x = x_next + else: + # Internal: tombstone Mamba + cascade x_next = lru.get_prev_no_lock(x) self.cache._evict_component_and_detach_lru( - x, self, is_leaf=False, tracker=tracker + x, self, target=EvictLayer.DEVICE, tracker=tracker ) self.cache._cascade_evict(x, self, tracker) x = x_next - else: - self.cache._evict_component_and_detach_lru( - x, self, is_leaf=True, tracker=tracker - ) - self.cache._cascade_evict(x, self, tracker) - x = lru.get_lru_no_lock() def acquire_component_lock( self, node: UnifiedTreeNode, result: IncLockRefResult @@ -178,8 +231,7 @@ class MambaComponent(TreeComponent): ct = self.component_type cd = node.component_data[ct] value = cd.value - if value is not None: - assert cd.lock_ref > 0 + if value is not None and cd.lock_ref > 0: if cd.lock_ref == 1: vlen = len(value) self.cache.component_evictable_size_[ct] += vlen @@ -268,3 +320,114 @@ class MambaComponent(TreeComponent): ): self.cache.req_to_token_pool.mamba_pool.free(insert_params.mamba_value) req.mamba_last_track_seqlen = None + + # ---- HiCache Hooks ---- + + def build_hicache_transfers( + self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw + ) -> Optional[list[PoolTransfer]]: + ct = self.component_type + + if phase == CacheTransferPhase.BACKUP_HOST: + cd = node.component_data[ct] + if cd.value is None: + return None + return [ + PoolTransfer( + name=PoolName.MAMBA, + device_indices=cd.value, + ) + ] + + if phase == CacheTransferPhase.LOAD_BACK: + req = kw.get("req") + transfers: list[PoolTransfer] = [] + + cd = node.component_data[ct] + if cd.value is not None: + return None + + # restore single node if host_value exists and + if cd.host_value is not None and cd.value is None: + transfers.append( + PoolTransfer( + name=PoolName.MAMBA, + host_indices=cd.host_value, + nodes_to_load=[node], + ) + ) + + # Per-request mamba CoW (H→D copy into request's device slot) + cd = node.component_data[ct] + if req is not None and cd.host_value is not None: + if req.mamba_pool_idx is None: + dst = self.cache.req_to_token_pool.mamba_pool.alloc(1) + if dst is None: + self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) + dst = self.cache.req_to_token_pool.mamba_pool.alloc(1) + assert dst is not None, "Cannot alloc mamba for load_back" + req.mamba_pool_idx = dst[0] + transfers.append( + PoolTransfer( + name=PoolName.MAMBA, + host_indices=cd.host_value, + device_indices=req.mamba_pool_idx.unsqueeze(0), + ) + ) + + return transfers if transfers else None + + return None + + def commit_hicache_transfer( + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + transfers: list[PoolTransfer] = (), + ) -> None: + ct = self.component_type + + if phase == CacheTransferPhase.BACKUP_HOST: + if transfers and transfers[0].host_indices is not None: + cd = node.component_data[ct] + if cd.host_value is None: + cd.host_value = transfers[0].host_indices.clone() + + elif phase == CacheTransferPhase.LOAD_BACK: + if not transfers: + return + transfer = transfers[0] + if transfer.device_indices is not None: + cd = node.component_data[ct] + cd.value = transfer.device_indices.clone() + count = len(cd.value) + # Move from host LRU to device LRU + host_lru = self.cache.host_lru_lists[ct] + if host_lru.in_list(node): + host_lru.remove_node(node) + self.cache.lru_lists[ct].insert_mru(node) + self.cache.component_evictable_size_[ct] += count + + def drive_host_eviction( + self, num_tokens: int, tracker: dict[ComponentType, int] + ) -> None: + """Evict mamba host resources. + Internal nodes: private tombstone (free host mamba only). + Host leaves: atomic eviction via _evict_host_leaf.""" + ct = self.component_type + host_lru = self.cache.host_lru_lists[ct] + x = host_lru.get_lru_no_lock() + while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x): + x_next = host_lru.get_prev_no_lock(x) + cd = x.component_data[ct] + if x in self.cache.evictable_host_leaves: + # Host leaf: atomic eviction (all components host + delete) + self.cache._evict_host_leaf(x, tracker) + else: + # Internal: tombstone Mamba + cascade + assert cd.host_value is not None + self.cache._evict_component_and_detach_lru( + x, self, target=EvictLayer.HOST, tracker=tracker + ) + self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST) + x = x_next diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index b38b8842f..4eccfecfc 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -14,6 +14,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.unified_cache_components.tree_component import ( BASE_COMPONENT_TYPE, ComponentType, + EvictLayer, TreeComponent, next_component_uuid, ) @@ -186,10 +187,17 @@ class SWAComponent(TreeComponent): ) child.component_data[self.component_type].metadata.pop("uuid", None) - def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: + def evict_component( + self, + node: UnifiedTreeNode, + target: EvictLayer = EvictLayer.DEVICE, + ) -> tuple[int, int]: + if target is EvictLayer.HOST: + return 0, 0 # TODO:SWA has no host layer currently + swa_value = node.component_data[self.component_type].value if swa_value is None: - return 0 + return 0, 0 # Direct swa_attn_allocator.free(swa_value) would double-free # free_swa(full_value) has the mapping guard to avoid double-free # TODO: decoupling full and swa free, need further discussion on mapping necessity @@ -198,9 +206,9 @@ class SWAComponent(TreeComponent): ) freed = len(swa_value) self.cache.component_evictable_size_[self.component_type] -= freed - if not is_leaf: + if target is EvictLayer.DEVICE: node.component_data[self.component_type].value = None - return freed + return freed, 0 def eviction_priority(self, is_leaf: bool) -> int: return 0 if is_leaf else 1 @@ -209,25 +217,26 @@ class SWAComponent(TreeComponent): self, params: EvictParams, tracker: dict[ComponentType, int] ) -> None: request = params.swa_num_tokens - lru = self.cache.lru_lists[self.component_type] + ct = self.component_type + lru = self.cache.lru_lists[ct] x = lru.get_lru_no_lock() - while ( - tracker[self.component_type] < request and x is not None and lru.in_list(x) - ): - assert x.component_data[self.component_type].value is not None - if len(x.children) > 0: + while tracker[ct] < request and x is not None and lru.in_list(x): + assert x.component_data[ct].value is not None + if x in self.cache.evictable_device_leaves: + # D-leaf: atomic eviction of all components + x_next = lru.get_prev_no_lock(x) + self.cache._evict_device_leaf(x, tracker) + if not lru.in_list(x_next): + x_next = lru.get_lru_no_lock() + x = x_next + else: + # Internal: tombstone SWA + cascade x_next = lru.get_prev_no_lock(x) self.cache._evict_component_and_detach_lru( - x, self, is_leaf=False, tracker=tracker + x, self, target=EvictLayer.DEVICE, tracker=tracker ) self.cache._cascade_evict(x, self, tracker) x = x_next - else: - self.cache._evict_component_and_detach_lru( - x, self, is_leaf=True, tracker=tracker - ) - self.cache._cascade_evict(x, self, tracker) - x = lru.get_lru_no_lock() def acquire_component_lock( self, node: UnifiedTreeNode, result: IncLockRefResult diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index 31a2bad5d..75e3123ed 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py @@ -2,7 +2,7 @@ from __future__ import annotations import dataclasses from abc import ABC, abstractmethod -from enum import Enum +from enum import Enum, IntFlag from typing import TYPE_CHECKING, Any, Callable, Optional import torch @@ -67,6 +67,14 @@ class ComponentData: host_lock_ref: int = 0 +class EvictLayer(IntFlag): + """Which storage layer(s) to evict. Combinable via bitwise OR.""" + + DEVICE = 1 + HOST = 2 + ALL = DEVICE | HOST + + class CacheTransferPhase(str, Enum): BACKUP_HOST = "backup_host" # D→H @@ -95,8 +103,13 @@ class TreeComponent(ABC): # Subclasses MUST set this as a class attribute (not @property) component_type: ComponentType - def node_has_component_data(self, node: UnifiedTreeNode) -> bool: - return node.component_data[self.component_type].value is not None + def node_has_component_data( + self, node: UnifiedTreeNode, target: EvictLayer = EvictLayer.DEVICE + ) -> bool: + cd = node.component_data[self.component_type] + if target is EvictLayer.DEVICE: + return cd.value is not None + return cd.host_value is not None def value_len(self, node: UnifiedTreeNode) -> int: value = node.component_data[self.component_type].value @@ -186,17 +199,22 @@ class TreeComponent(ABC): ... @abstractmethod - def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: + def evict_component( + self, + node: UnifiedTreeNode, + target: EvictLayer = EvictLayer.DEVICE, + ) -> tuple[int, int]: """Free this component's KV resources on a node being evicted. - For internal (non-leaf) nodes: free memory and tombstone the value - (set to None); the node structure is kept. - For leaf nodes: free memory; the node will be deleted by caller. - Returns the number of tokens/slots freed. - - Full: frees full_value via token_to_kv_pool_allocator. - - SWA: frees swa value via swa_token_to_kv_pool_allocator; - only tombstones on internal nodes. - - Mamba: frees mamba value via mamba_token_to_kv_pool_allocator; - only tombstones on internal nodes.""" + + *target* controls which layer(s) to evict: + - DEVICE: free device memory and tombstone (value = None). + Host data is untouched. + - HOST: free host memory (host_value = None). + Device data is untouched. + - ALL: free both device and host memory. + No tombstone — caller will delete the node. + + Returns (device_freed, host_freed) token counts.""" ... def eviction_priority(self, is_leaf: bool) -> int: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 38dcb5652..5b0c56a84 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import threading import time from collections import defaultdict from functools import partial @@ -15,17 +16,21 @@ from sglang.srt.mem_cache.base_prefix_cache import ( EvictParams, EvictResult, IncLockRefResult, + InitLoadBackParams, InsertParams, InsertResult, MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache_components import ( _NUM_COMPONENT_TYPES, BASE_COMPONENT_TYPE, + CacheTransferPhase, ComponentData, ComponentType, + EvictLayer, FullComponent, MambaComponent, SWAComponent, @@ -38,6 +43,7 @@ from sglang.srt.session.streaming_session import StreamingSession if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.server_args import ServerArgs class UnifiedTreeNode: @@ -53,45 +59,68 @@ class UnifiedTreeNode: ComponentData() for _ in range(_NUM_COMPONENT_TYPES) ] self.last_access_time = get_and_increase_time_counter() - self.host_value = None + self.hash_value = None self.hit_count = 0 - self.lru_prev: list[UnifiedTreeNode | None] = [None] * _NUM_COMPONENT_TYPES - self.lru_next: list[UnifiedTreeNode | None] = [None] * _NUM_COMPONENT_TYPES + self.lru_prev: list[UnifiedTreeNode | None] = [None] * ( + _NUM_COMPONENT_TYPES * 2 + ) + self.lru_next: list[UnifiedTreeNode | None] = [None] * ( + _NUM_COMPONENT_TYPES * 2 + ) self.id = UnifiedTreeNode.counter UnifiedTreeNode.counter += 1 def component(self, component_type: ComponentType) -> ComponentData: return self.component_data[component_type] + @property + def backuped(self) -> bool: + """Tree-level: Full KV present on host.""" + return self.component_data[ComponentType.FULL].host_value is not None + + @property + def evicted(self) -> bool: + """Tree-level: Full KV not on device (non-root with value=None).""" + return ( + self.parent is not None + and self.component_data[ComponentType.FULL].value is None + ) + def __lt__(self, other: UnifiedTreeNode): return self.last_access_time < other.last_access_time class UnifiedLRUList: def __init__( - self, component_type: ComponentType, tree_components: tuple[ComponentType, ...] + self, + component_type: ComponentType, + tree_components: tuple[ComponentType, ...], + use_host_ptr: bool = False, ): self.component_type = component_type + # Pointer slot: host LRU uses offset slots so device/host pointers + # never collide on the same node. + self._pt: int = component_type + (_NUM_COMPONENT_TYPES if use_host_ptr else 0) self.head = UnifiedTreeNode(tree_components) self.tail = UnifiedTreeNode(tree_components) - self.head.lru_next[component_type] = self.tail - self.tail.lru_prev[component_type] = self.head + self.head.lru_next[self._pt] = self.tail + self.tail.lru_prev[self._pt] = self.head self.cache: dict[int, UnifiedTreeNode] = {} def _add_node_after(self, prev_node: UnifiedTreeNode, new_node: UnifiedTreeNode): - ct = self.component_type - new_node.lru_prev[ct] = prev_node - new_node.lru_next[ct] = prev_node.lru_next[ct] - prev_node.lru_next[ct].lru_prev[ct] = new_node - prev_node.lru_next[ct] = new_node + pt = self._pt + new_node.lru_prev[pt] = prev_node + new_node.lru_next[pt] = prev_node.lru_next[pt] + prev_node.lru_next[pt].lru_prev[pt] = new_node + prev_node.lru_next[pt] = new_node def _add_node(self, node: UnifiedTreeNode): self._add_node_after(self.head, node) def _remove_node(self, node: UnifiedTreeNode): - ct = self.component_type - node.lru_prev[ct].lru_next[ct] = node.lru_next[ct] - node.lru_next[ct].lru_prev[ct] = node.lru_prev[ct] + pt = self._pt + node.lru_prev[pt].lru_next[pt] = node.lru_next[pt] + node.lru_next[pt].lru_prev[pt] = node.lru_prev[pt] def insert_mru(self, node: UnifiedTreeNode): assert node.id not in self.cache @@ -129,10 +158,11 @@ class UnifiedLRUList: def get_prev_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): if check_id: assert node.id in self.cache + pt = self._pt ct = self.component_type - x = node.lru_prev[ct] + x = node.lru_prev[pt] while x.component_data[ct].lock_ref > 0: - x = x.lru_prev[ct] + x = x.lru_prev[pt] if x == self.head: return None return x @@ -140,10 +170,11 @@ class UnifiedLRUList: def get_prev_leaf_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): if check_id: assert node.id in self.cache + pt = self._pt ct = self.component_type - x = node.lru_prev[ct] + x = node.lru_prev[pt] while x.component_data[ct].lock_ref > 0 or len(x.children) > 0: - x = x.lru_prev[ct] + x = x.lru_prev[pt] if x == self.head: return None return x @@ -191,6 +222,9 @@ class UnifiedRadixCache(BasePrefixCache): self._components_tuple: tuple[TreeComponent, ...] = tuple( self.components.values() ) + self.hicache_anchor_kv_shared_indices_pools: list[ + tuple[PoolName, PoolHitPolicy] + ] = [] if self.is_eagle: self.key_convert_fn = convert_to_bigram_key else: @@ -203,10 +237,25 @@ class UnifiedRadixCache(BasePrefixCache): # internal fall-through to self.inner.xxx never fires -- no recursion. self.session = StreamingSession(inner=self) + self.tp_group = params.tp_cache_group + self.tp_world_size = ( + 1 + if self.tp_group is None + else torch.distributed.get_world_size(group=self.tp_group) + ) + + # HiCache D↔H defaults (overridden by init_hicache) + self.cache_controller = None + self.write_through_threshold = 256 + self.reset() logger.info(f"Init Unified RadixTree with components {self.tree_components}") def reset(self) -> None: + self._reset_full() + + def _reset_full(self) -> None: + """Full reset: destroy entire tree and all state.""" self.root_node = UnifiedTreeNode(self.tree_components) self.root_node.key = RadixKey([], None) self.root_node.component_data[BASE_COMPONENT_TYPE].value = [] @@ -214,11 +263,73 @@ class UnifiedRadixCache(BasePrefixCache): self.root_node.component_data[ct].lock_ref = 1 self.component_evictable_size_ = {ct: 0 for ct in self.tree_components} self.component_protected_size_ = {ct: 0 for ct in self.tree_components} + self.lru_lists = { ct: UnifiedLRUList(ct, self.tree_components) for ct in self.tree_components } self.session.slots.clear() + self.evictable_device_leaves: set[UnifiedTreeNode] = set() + self.evictable_host_leaves: set[UnifiedTreeNode] = set() + self.host_lru_lists = { + ct: UnifiedLRUList(ct, self.tree_components, use_host_ptr=True) + for ct in self.tree_components + } + self.ongoing_write_through: dict[int, UnifiedTreeNode] = {} + self.ongoing_load_back: dict[int, UnifiedTreeNode] = {} + self.enable_storage = False + self.ongoing_prefetch: dict = {} + self.ongoing_backup: dict = {} + + if self.cache_controller is not None: + self.cache_controller.reset() + self.cache_controller.mem_pool_host.clear() + + def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None: + """Initialize HiCache infrastructure.""" + from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + attach_hybrid_pool_to_unified_cache, + ) + + # Direct IO layout fixup (must happen before pool creation) + if server_args.hicache_io_backend == "direct": + if server_args.hicache_mem_layout == "page_first": + server_args.hicache_mem_layout = "page_first_direct" + logger.warning( + "Page first layout is not supported with direct IO backend, " + "switching to page first direct layout" + ) + + self.load_cache_event = threading.Event() + self.hicache_anchor_kv_shared_indices_pools.clear() + attach_hybrid_pool_to_unified_cache( + self, + params, + server_args, + load_cache_event=self.load_cache_event, + ) + + # State initialization + self.write_through_threshold = ( + 1 if server_args.hicache_write_policy == "write_through" else 2 + ) + self.load_back_threshold = 256 + + logger.info( + f"HiCache D\u2194H initialized: " + f"host_pool_size={self.host_pool_group.size}, " + f"write_policy={server_args.hicache_write_policy}, " + f"tp_world_size={self.tp_world_size}, " + f"transfer_layer_num={self.cache_controller.layer_num}" + ) + + def register_hicache_anchor_kv_shared_indices_pool( + self, + pool_name: PoolName, + hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES, + ) -> None: + self.hicache_anchor_kv_shared_indices_pools.append((pool_name, hit_policy)) + def match_prefix(self, params: MatchPrefixParams) -> MatchResult: result = self.session.try_match_prefix(params) if result is not None: @@ -282,6 +393,8 @@ class UnifiedRadixCache(BasePrefixCache): result = IncLockRefResult() for component in self._components_tuple: result = component.acquire_component_lock(node=node, result=result) + + self._update_evictable_leaf_sets(node) return result def dec_lock_ref( @@ -294,6 +407,8 @@ class UnifiedRadixCache(BasePrefixCache): return DecLockRefResult() for component in self._components_tuple: component.release_component_lock(node=node, params=params) + + self._update_evictable_leaf_sets(node) # TODO: delta is not aggregated from components; no caller uses it yet. return DecLockRefResult() @@ -386,7 +501,9 @@ class UnifiedRadixCache(BasePrefixCache): ] # components prepare insert data + return effective cache_len - insert_params = InsertParams(prev_prefix_len=req.cache_protected_len) + insert_params = InsertParams( + prev_prefix_len=req.cache_protected_len, chunked=chunked + ) effective_cache_len = len(token_ids) for comp in self._components_tuple: cl = comp.prepare_for_caching_req( @@ -488,11 +605,18 @@ class UnifiedRadixCache(BasePrefixCache): while len(key) > 0 and child_key in node.children: child = node.children[child_key] + + # HiCache: dead node (evicted + not backuped) — stop traversal + if child.evicted and not child.backuped: + break + prefix_len = child.key.match(key, page_size=self.page_size) if prefix_len < len(child.key): # Read-only: do not split, ignore partial match and stop break - value.append(child.component_data[BASE_COMPONENT_TYPE].value) + + if not child.evicted: + value.append(child.component_data[BASE_COMPONENT_TYPE].value) node = child _update_best_if_valid(node) key = key[prefix_len:] @@ -520,13 +644,22 @@ class UnifiedRadixCache(BasePrefixCache): while len(key) > 0 and child_key in node.children: child = node.children[child_key] + + # HiCache: dead node (evicted + not backuped) — stop traversal + if child.evicted and not child.backuped: + break + prefix_len = child.key.match(key, page_size=self.page_size) if prefix_len < len(child.key): + if child.evicted: + break node = self._split_node(child.key, child, prefix_len) value.append(node.component_data[BASE_COMPONENT_TYPE].value) _update_best_if_valid(node) break - value.append(child.component_data[BASE_COMPONENT_TYPE].value) + + if not child.evicted: + value.append(child.component_data[BASE_COMPONENT_TYPE].value) node = child _update_best_if_valid(node) key = key[prefix_len:] @@ -543,23 +676,37 @@ class UnifiedRadixCache(BasePrefixCache): ) -> MatchResult: node_update = last_node for comp in self._components_tuple: + if comp.component_type == BASE_COMPONENT_TYPE: + continue # Full uses last_access_time, not LRU self.lru_lists[comp.component_type].reset_node_and_parents_mru( node_update, self.root_node, comp.node_has_component_data ) + cur_time = get_and_increase_time_counter() while node_update: node_update.last_access_time = cur_time cur_time -= 0.00001 node_update = node_update.parent + # Walk up to find last_device_node + last_device_node = last_node + while last_device_node is not self.root_node and last_device_node.evicted: + last_device_node = last_device_node.parent + + # Walk up to find last_host_node + last_host_node = last_node + while last_host_node is not self.root_node and not last_host_node.backuped: + last_host_node = last_host_node.parent + if best_value_len > 0: device_indices = torch.cat(value[:best_value_len]) else: device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) result = MatchResult( device_indices=device_indices, - last_device_node=last_node, - last_host_node=last_node, + last_device_node=last_device_node, + last_host_node=last_host_node, + host_hit_length=0, ) for component in self._components_tuple: @@ -578,25 +725,26 @@ class UnifiedRadixCache(BasePrefixCache): new_node.children = {key[split_len:].child_key(self.page_size): child} new_node.parent = child.parent new_node.key = child.key[:split_len] - new_node.component_data[BASE_COMPONENT_TYPE].value = ( - child.component_data[BASE_COMPONENT_TYPE].value[:split_len].clone() - ) self._for_each_component_lru(child, UnifiedLRUList.remove_node) child.parent = new_node child.key = child.key[split_len:] - child.component_data[BASE_COMPONENT_TYPE].value = ( - child.component_data[BASE_COMPONENT_TYPE].value[split_len:].clone() - ) for component in self._components_tuple: component.redistribute_on_node_split(new_parent=new_node, child=child) new_node.parent.children[key.child_key(self.page_size)] = new_node - self._for_each_component_lru(new_node, UnifiedLRUList.insert_mru) - self._for_each_component_lru(child, UnifiedLRUList.insert_mru) + self._for_each_component_lru( + new_node, UnifiedLRUList.insert_mru, skip_existing=True + ) + self._for_each_component_lru( + child, UnifiedLRUList.insert_mru, skip_existing=True + ) child.last_access_time = get_and_increase_time_counter() + + self._update_evictable_leaf_sets(new_node) + self._update_evictable_leaf_sets(child) return new_node def _touch_node(self, node: UnifiedTreeNode): @@ -615,10 +763,27 @@ class UnifiedRadixCache(BasePrefixCache): new_node.key = key new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone() parent.children[key.child_key(self.page_size)] = new_node - self.lru_lists[BASE_COMPONENT_TYPE].insert_mru(new_node) self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value) + + self._update_evictable_leaf_sets(new_node) + self._update_evictable_leaf_sets(parent) return new_node + def _unevict_node_on_insert( + self, node: UnifiedTreeNode, fresh_value: torch.Tensor + ) -> None: + """Restore an evicted node's Full device value from fresh KV indices + during insert.""" + ct = BASE_COMPONENT_TYPE + cd = node.component_data[ct] + assert cd.value is None + n = len(fresh_value) + cd.value = fresh_value.clone() + self.component_evictable_size_[ct] += n + self._update_evictable_leaf_sets(node) + if node.parent is not None: + self._update_evictable_leaf_sets(node.parent) + def _insert_helper( self, node: UnifiedTreeNode, @@ -639,25 +804,29 @@ class UnifiedRadixCache(BasePrefixCache): if prefix_len < len(node.key): node = self._split_node(node.key, node, prefix_len) - value_slice = value[:prefix_len] - consumed_from = prefix_len - # Let each component claim ownership of overlapping KV slots - for component in self._components_tuple: - comp_consumed_from = component.update_component_on_insert_overlap( - node=node, - prefix_len=prefix_len, - total_prefix_len=total_prefix_length, - value_slice=value_slice, - params=params, - ) - consumed_from = min(consumed_from, comp_consumed_from) + if node.evicted: + self._unevict_node_on_insert(node, value[:prefix_len]) + else: + value_slice = value[:prefix_len] + consumed_from = prefix_len + # Let each component claim ownership of overlapping KV slots + for component in self._components_tuple: + comp_consumed_from = component.update_component_on_insert_overlap( + node=node, + prefix_len=prefix_len, + total_prefix_len=total_prefix_length, + value_slice=value_slice, + params=params, + ) + consumed_from = min(consumed_from, comp_consumed_from) - dup_start = max(0, params.prev_prefix_len - total_prefix_length) - if dup_start < consumed_from: - self.token_to_kv_pool_allocator.free( - value_slice[dup_start:consumed_from] - ) + dup_start = max(0, params.prev_prefix_len - total_prefix_length) + if dup_start < consumed_from: + self.token_to_kv_pool_allocator.free( + value_slice[dup_start:consumed_from] + ) + self._inc_hit_count(node, params.chunked) total_prefix_length += prefix_len key = key[prefix_len:] value = value[prefix_len:] @@ -695,34 +864,45 @@ class UnifiedRadixCache(BasePrefixCache): params=params, result=result, ) + if is_new_leaf: + self._inc_hit_count(target_node, params.chunked) return result + # ---- Evict Helpers ---- + def _cascade_evict( self, node: UnifiedTreeNode, trigger: TreeComponent, tracker: dict[ComponentType, int], + target: EvictLayer = EvictLayer.DEVICE, ): - """Cascade eviction from trigger to lower-or-equal priority components. - - When a component evicts a node, all other components with equal or - lower eviction_priority on the same node are also evicted. - If the node is a leaf, it is removed from the tree and any - resulting tombstone ancestors are cleaned up recursively.""" + """Cascade eviction from trigger to lower-or-equal priority components.""" is_leaf = len(node.children) == 0 trigger_priority = trigger.eviction_priority(is_leaf) for comp in self._components_tuple: if comp.eviction_priority(is_leaf) <= trigger_priority: - if comp is not trigger and comp.node_has_component_data(node): - assert node.component_data[comp.component_type].lock_ref == 0 + if comp is not trigger and comp.node_has_component_data(node, target): + cd = node.component_data[comp.component_type] + if EvictLayer.DEVICE in target: + assert cd.lock_ref == 0 + if EvictLayer.HOST in target: + assert cd.host_lock_ref == 0 self._evict_component_and_detach_lru( - node, comp, is_leaf=is_leaf, tracker=tracker + node, comp, target=target, tracker=tracker ) - if is_leaf: - self._remove_leaf_from_parent(node) - self._iteratively_delete_tombstone_leaf(node, tracker) + # Now that all components (including SWA which depends on Full.value) + # have been freed, we can safely tombstone Full.value. + # This is deferred from evict_component because free_swa needs it. + if ( + target is EvictLayer.DEVICE + and trigger.component_type == BASE_COMPONENT_TYPE + ): + node.component_data[trigger.component_type].value = None + + self._update_evictable_leaf_sets(node) def _remove_leaf_from_parent(self, node: UnifiedTreeNode): key = node.key.child_key(self.page_size) @@ -733,49 +913,497 @@ class UnifiedRadixCache(BasePrefixCache): self, node: UnifiedTreeNode, comp: TreeComponent, - is_leaf: bool, - tracker: dict[ComponentType, int], - ) -> int: - freed = comp.evict_component(node, is_leaf=is_leaf) - tracker[comp.component_type] += freed - lru = self.lru_lists[comp.component_type] - if lru.in_list(node): - lru.remove_node(node) - return freed + target: EvictLayer = EvictLayer.DEVICE, + tracker: dict[ComponentType, int] = None, + ) -> tuple[int, int]: + device_freed, host_freed = comp.evict_component(node, target=target) + if tracker is not None: + if EvictLayer.DEVICE in target: + tracker[comp.component_type] += device_freed + elif EvictLayer.HOST in target: + tracker[comp.component_type] += host_freed + + # Detach from the appropriate LRU list(s) + ct = comp.component_type + for layer, lru_lists in ( + (EvictLayer.DEVICE, self.lru_lists), + (EvictLayer.HOST, self.host_lru_lists), + ): + if layer in target: + lru = lru_lists[ct] + if lru.in_list(node): + lru.remove_node(node) + return device_freed, host_freed def _iteratively_delete_tombstone_leaf( self, deleted_node: UnifiedTreeNode, tracker: dict[ComponentType, int] ): - """After a leaf is removed, walk up the parent chain and delete - any ancestor that is leaf node and has lost any component data (tombstoned).""" + """Walk up from *deleted_node* and cascade-delete childless ancestors. + + Only the Full (base) component decides whether a node survives: + - Full device present → keep as D-leaf + - Full host present → keep as H-leaf + - neither → evict all remaining data, delete, continue up + """ + ct = BASE_COMPONENT_TYPE cur = deleted_node.parent while cur != self.root_node and len(cur.children) == 0: - has_tombstone = any( - not comp.node_has_component_data(cur) - for comp in self.components.values() - ) - if not has_tombstone: - break - if any( - cur.component_data[comp.component_type].lock_ref > 0 - for comp in self.components.values() - if comp.node_has_component_data(cur) + cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data ): break + has_device = cur.component_data[ct].value is not None + has_host = cur.component_data[ct].host_value is not None + + if has_device: + self._update_evictable_leaf_sets(cur) + break + + # Full device absent — clean up orphaned aux device data. for comp in self.components.values(): if comp.node_has_component_data(cur): self._evict_component_and_detach_lru( - cur, comp, is_leaf=True, tracker=tracker + cur, comp, target=EvictLayer.DEVICE, tracker=tracker ) - self._remove_leaf_from_parent(cur) - cur = cur.parent - def _for_each_component_lru(self, node: UnifiedTreeNode, lru_op): + if has_host: + self._update_evictable_leaf_sets(cur) + break + + # Full absent on both layers — evict remaining host data, delete. + for comp in self.components.values(): + if comp.node_has_component_data(cur, target=EvictLayer.HOST): + self._evict_component_and_detach_lru( + cur, comp, target=EvictLayer.HOST, tracker=tracker + ) + + self.evictable_host_leaves.discard(cur) + self._remove_leaf_from_parent(cur) + parent = cur.parent + self._update_evictable_leaf_sets(parent) + cur = parent + + def _for_each_component_lru( + self, + node: UnifiedTreeNode, + lru_op, + target: EvictLayer = EvictLayer.DEVICE, + skip_existing: bool = False, + ): + """Apply lru_op to each aux component's LRU that has data on this node. + If skip_existing=True, skip components already in the target LRU list.""" + lru_dict = self.host_lru_lists if target is EvictLayer.HOST else self.lru_lists for ct in self.tree_components: - if node.component_data[ct].value is not None: - lru_op(self.lru_lists[ct], node) + if ct == BASE_COMPONENT_TYPE: + continue # Full uses leaf sets, not LRU + cd = node.component_data[ct] + if (cd.host_value if target is EvictLayer.HOST else cd.value) is not None: + lru = lru_dict[ct] + if skip_existing and lru.in_list(node): + continue + lru_op(lru, node) + + def evict_host( + self, num_tokens: int, component_type: ComponentType = BASE_COMPONENT_TYPE + ) -> int: + """Evict host resources for a specific component to free host pool space.""" + tracker: dict[ComponentType, int] = {ct: 0 for ct in self.tree_components} + comp = self.components.get(component_type) + if comp is not None: + comp.drive_host_eviction(num_tokens, tracker) + return tracker[component_type] + + def _is_device_leaf(self, node: UnifiedTreeNode) -> bool: + """D-leaf: Full device value present, no child with Full KV on device, + unlocked, not root. + + Only the Full (base) component is required; auxiliary components + (Mamba, SWA) are not mandatory for D-leaf membership.""" + ct = BASE_COMPONENT_TYPE + if node is self.root_node or node.evicted: + return False + if any(cd.lock_ref > 0 for cd in node.component_data): + return False + if any( + child.component_data[ct].value is not None + for child in node.children.values() + ): + return False + return True + + def _is_host_leaf(self, node: UnifiedTreeNode) -> bool: + """H-leaf: evicted, Full host value present, no children, unlocked, not root. + + Only the Full (base) component host_value is required; auxiliary + components are not mandatory for H-leaf membership.""" + if node is self.root_node or not node.evicted: + return False + if not node.backuped: + return False + if any(cd.host_lock_ref > 0 for cd in node.component_data): + return False + if len(node.children) > 0: + return False + return True + + def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None: + """Update both device and host leaf sets for a node.""" + if self._is_device_leaf(node): + self.evictable_device_leaves.add(node) + else: + self.evictable_device_leaves.discard(node) + + if self._is_host_leaf(node): + self.evictable_host_leaves.add(node) + else: + self.evictable_host_leaves.discard(node) + + def _evict_to_host( + self, node: UnifiedTreeNode, tracker: dict[ComponentType, int] = None + ) -> None: + """GPU→CPU demotion: release all device resources, node stays in tree.""" + assert not node.evicted and node.backuped + trigger = self.components[BASE_COMPONENT_TYPE] + self._evict_component_and_detach_lru( + node, trigger, target=EvictLayer.DEVICE, tracker=tracker + ) + self._cascade_evict(node, trigger, tracker) + + # after device eviction, insert aux components into host LRU. + self._for_each_component_lru( + node, UnifiedLRUList.insert_mru, target=EvictLayer.HOST, skip_existing=True + ) + self._update_evictable_leaf_sets(node.parent) + + def _evict_device_leaf( + self, node: UnifiedTreeNode, tracker: dict[ComponentType, int] + ) -> None: + """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: Cascade evict all components + + All freed device tokens are accumulated into *tracker*. + """ + assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf" + if not node.backuped: + if ( + self.cache_controller is not None + and self.cache_controller.write_policy == "write_back" + ): + self.write_backup(node, write_back=True) + self._evict_to_host(node, tracker) + return + else: + # Write-through: node has no backup, delete entirely. + for comp in self._components_tuple: + self._evict_component_and_detach_lru( + node, comp, target=EvictLayer.ALL, tracker=tracker + ) + self.evictable_device_leaves.discard(node) + parent = node.parent + self._remove_leaf_from_parent(node) + self._update_evictable_leaf_sets(parent) + self._iteratively_delete_tombstone_leaf(node, tracker) + return + self._evict_to_host(node, tracker) + + def _evict_host_leaf( + self, node: UnifiedTreeNode, tracker: dict[ComponentType, int] + ) -> None: + """Atomically evict all components on a host leaf. + + All freed tokens are accumulated into *tracker*.""" + assert self._is_host_leaf(node), f"node {node.id} is not an H-leaf" + + for comp in self._components_tuple: + _, hf = self._evict_component_and_detach_lru( + node, comp, target=EvictLayer.ALL, tracker=None + ) + tracker[comp.component_type] += hf + self.evictable_host_leaves.discard(node) + self._remove_leaf_from_parent(node) + self._iteratively_delete_tombstone_leaf(node, tracker) + + # ---- HiCache: Backup / LoadBack ---- + + def write_backup(self, node: UnifiedTreeNode, write_back: bool = False) -> int: + """Backup a node's data from device to host (D->H).""" + if self.cache_controller is None: + return 0 + + # Backup invariant (write-through): parent must be backuped first + if not write_back and ( + node.parent is not self.root_node and not node.parent.backuped + ): + return 0 + + # Build aux transfers, keyed per component + comp_xfers: dict[ComponentType, list] = {} + for comp in self._components_tuple: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST) + if t: + comp_xfers[comp.component_type] = t + anchor_kv_shared_indices_xfers = [ + PoolTransfer(name=pool_name, hit_policy=hit_policy) + for pool_name, hit_policy in self.hicache_anchor_kv_shared_indices_pools + ] + + # Pre-evict host if insufficient + device_value = node.component_data[BASE_COMPONENT_TYPE].value + kv_tokens = len(device_value) + host_avail = self.cache_controller.mem_pool_host.available_size() + if host_avail < kv_tokens: + needed = kv_tokens - host_avail + evicted = self.evict_host(needed) + if evicted < needed: + return 0 + + aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] + aux_xfers.extend(anchor_kv_shared_indices_xfers) + host_indices = self.cache_controller.write( + device_value, node_id=node.id, extra_pools=aux_xfers or None + ) + if host_indices is None: + return 0 + + # Commit + kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices) + self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer( + node, + CacheTransferPhase.BACKUP_HOST, + transfers=[kv_xfer], + ) + for ct, xfers in comp_xfers.items(): + self.components[ct].commit_hicache_transfer( + node, + CacheTransferPhase.BACKUP_HOST, + transfers=xfers, + ) + + self.ongoing_write_through[node.id] = node + if not write_back: + self.inc_lock_ref(node) + return len(host_indices) + + def load_back( + self, + node: UnifiedTreeNode, + mem_quota: Optional[int] = None, + req=None, + ) -> Optional[torch.Tensor]: + """Load evicted KV data from host back to device (H→D).""" + if self.cache_controller is None: + return None + + # Build KV transfer + last_hit_node = node + kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers( + last_hit_node, CacheTransferPhase.LOAD_BACK + )[0] + + # Lock path & pre-evict if device pool is insufficient + nodes_to_load = kv_xfer.nodes_to_load + ancestor_node = nodes_to_load[0].parent if nodes_to_load else last_hit_node + result = self.inc_lock_ref(ancestor_node) + kv_tokens = len(kv_xfer.host_indices) + + # Build aux transfers, keyed per component. + comp_xfers: dict[ComponentType, list] = {} + for comp in self._components_tuple: + if comp.component_type == BASE_COMPONENT_TYPE: + continue + t = comp.build_hicache_transfers( + last_hit_node, CacheTransferPhase.LOAD_BACK, req=req + ) + if t: + comp_xfers[comp.component_type] = t + anchor_kv_shared_indices_xfers = [ + PoolTransfer(name=pool_name, hit_policy=hit_policy) + for pool_name, hit_policy in self.hicache_anchor_kv_shared_indices_pools + ] + + # Skip if there is nothing to load, or if the Full-KV transfer is too + # small / exceeds memory quota. Aux transfers should still run even + # when the Full-KV load is skipped by thresholding. + if (kv_tokens < self.load_back_threshold and not comp_xfers) or ( + mem_quota is not None and kv_tokens > mem_quota + result.delta + ): + self.dec_lock_ref(ancestor_node) + return None + + avail = self.token_to_kv_pool_allocator.available_size() + if avail < kv_tokens: + needed = kv_tokens - avail + result = self.evict(EvictParams(num_tokens=needed)) + if result.num_tokens_evicted < needed: + self.dec_lock_ref(ancestor_node) + return None + + logger.info( + "load_back: kv_tokens=%d, node_id=%d", + kv_tokens, + last_hit_node.id, + ) + + # Load H→D + aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] + aux_xfers.extend(anchor_kv_shared_indices_xfers) + device_indices = self.cache_controller.load( + host_indices=kv_xfer.host_indices, + node_id=last_hit_node.id, + extra_pools=aux_xfers or None, + ) + + self.dec_lock_ref(ancestor_node) + if device_indices is None: + return None + + # Commit: each component gets only its own transfers + kv_xfer.device_indices = device_indices + self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer( + last_hit_node, + CacheTransferPhase.LOAD_BACK, + [kv_xfer], + ) + for ct, xfers in comp_xfers.items(): + self.components[ct].commit_hicache_transfer( + last_hit_node, + CacheTransferPhase.LOAD_BACK, + xfers, + ) + + self._update_evictable_leaf_sets(ancestor_node) + self.inc_lock_ref(last_hit_node) + self.ongoing_load_back[last_hit_node.id] = last_hit_node + return device_indices + + def _inc_hit_count(self, node: UnifiedTreeNode, chunked: bool = False) -> None: + """Increment hit count; trigger write_backup when threshold reached.""" + if self.cache_controller is None: + return + if node.evicted or chunked: + return + if self.cache_controller.write_policy == "write_back": + return + node.hit_count += 1 + if not node.backuped and node.hit_count >= self.write_through_threshold: + self.write_backup(node) + + # ---- HiCache: Async Event Management ---- + + def writing_check(self, write_back: bool = False) -> None: + """Poll write-through completions.""" + cc = self.cache_controller + if cc is None: + return + + if write_back: + # Blocking: wait for all pending write-backs + while self.ongoing_write_through: + for _, finish_event, ack_list in cc.ack_write_queue: + finish_event.synchronize() + for ack_id in ack_list: + self.ongoing_write_through.pop(ack_id, None) + cc.ack_write_queue.clear() + assert len(self.ongoing_write_through) == 0 + return + + if len(self.ongoing_write_through) == 0: + return + + finish_count = 0 + for _, finish_event, ack_list in cc.ack_write_queue: + if not finish_event.query(): + break + finish_count += 1 + + # TP sync: MIN across all ranks for consistent tree updates + queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") + if self.tp_world_size > 1: + torch.distributed.all_reduce( + queue_size, op=torch.distributed.ReduceOp.MIN, group=self.tp_group + ) + finish_count = int(queue_size.item()) + + # Process completed acks + while finish_count > 0: + _, finish_event, ack_list = cc.ack_write_queue.pop(0) + finish_event.synchronize() + for ack_id in ack_list: + node = self.ongoing_write_through.pop(ack_id) + self.dec_lock_ref(node) + finish_count -= 1 + + def loading_check(self) -> None: + """Poll load-back completions.""" + cc = self.cache_controller + if cc is None or not self.ongoing_load_back: + return + finish_count = 0 + for _, finish_event, ack_list in cc.ack_load_queue: + if not finish_event.query(): + break + finish_count += 1 + for ack_id in ack_list: + node = self.ongoing_load_back.pop(ack_id) + self.dec_lock_ref(node) + del cc.ack_load_queue[:finish_count] + + # ---- HiCache: Scheduler Entry Points ---- + + def init_load_back( + self, + params: InitLoadBackParams, + ) -> tuple[torch.Tensor, UnifiedTreeNode]: + """Prepare KV cache loading from host to device. + Returns (device_indices, last_node) tuple.""" + last_node = params.last_host_node + mem_quota = params.mem_quota + req = params.req + + if last_node.evicted or params.host_hit_length > 0: + logger.info( + "init_load_back triggered: node_id=%d, host_hit_length=%d", + last_node.id, + params.host_hit_length, + ) + loading_values = self.load_back(last_node, mem_quota, req=req) + if loading_values is not None: + logger.info( + "init_load_back success: loaded %d tokens for node %d", + len(loading_values), + last_node.id, + ) + return loading_values, last_node + + # Fallback: walk up to non-evicted ancestor + while last_node is not self.root_node and last_node.evicted: + last_node = last_node.parent + + return ( + torch.empty((0,), dtype=torch.int64, device=self.device), + last_node, + ) + + def check_hicache_events(self) -> None: + """Called per scheduler step to poll async HiCache events.""" + self.writing_check() + self.loading_check() + + def flush_write_through_acks(self) -> None: + """Flush pending write-through acknowledgements.""" + self.writing_check() + + def ready_to_load_host_cache(self) -> int: + """Notify the cache controller to start the KV cache loading.""" + if self.cache_controller is not None: + return self.cache_controller.start_loading() + return 0 # ---- Query / Inspection APIs ---- # These APIs exist for compatibility with other RadixTree implementations. @@ -845,7 +1473,9 @@ class UnifiedRadixCache(BasePrefixCache): stack = [self.root_node] while stack: node = stack.pop() - total_size += len(node.component_data[BASE_COMPONENT_TYPE].value) + full_value = node.component_data[BASE_COMPONENT_TYPE].value + if full_value is not None: + total_size += len(full_value) for ct in self.tree_components: if ct == BASE_COMPONENT_TYPE: continue @@ -861,7 +1491,9 @@ class UnifiedRadixCache(BasePrefixCache): def _dfs(node: UnifiedTreeNode): for child in node.children.values(): - values.append(child.component_data[BASE_COMPONENT_TYPE].value) + v = child.component_data[BASE_COMPONENT_TYPE].value + if v is not None: + values.append(v) _dfs(child) _dfs(self.root_node) @@ -931,82 +1563,268 @@ class UnifiedRadixCache(BasePrefixCache): return nodes def sanity_check(self): - """Thorough sanity check: verify LRU membership, lock state, linked-list - integrity, and evictable sizes for every component. - Expensive — use only in tests or idle checks.""" + """Verify tree invariants. + + TODO(hzh): This method has relatively high latency; simplify the + check logic once the tree implementation stabilizes. + """ # Skip when streaming sessions hold tree locks: the check asserts # all nodes are unlocked during idle, which streaming sessions break # by design (they hold a first-turn lock across turns). if self.session.any_holding_kv(): return - try: - # 1. Collect all nodes from tree - all_nodes = self._collect_all_nodes() + errors: list[str] = [] + E = errors.append + all_nodes = self._collect_all_nodes() + all_node_set = set(all_nodes) + FCT = BASE_COMPONENT_TYPE + + # ── PART 1: Tree Structure ── + # Root state + if self.root_node.component_data[FCT].value is None: + E("[Root] root missing Full device value") + if self.root_node.component_data[FCT].lock_ref <= 0: + E( + f"[Root] root Full lock_ref={self.root_node.component_data[FCT].lock_ref}" + ) + if self.root_node.parent is not None: + E("[Root] root has a parent pointer") + # Parent ↔ child bidirectional consistency + for node in all_nodes: + for child in node.children.values(): + if child.parent is not node: + pid = child.parent.id if child.parent else None + E(f"[Tree] child {child.id} parent={pid}, expected {node.id}") + if child.key is None: + E(f"[Tree] node {child.id} has no key") + + # ── PART 2: Per-node state machine and leaf qualification ── + expected_dev_leaves: set[UnifiedTreeNode] = set() + expected_hst_leaves: set[UnifiedTreeNode] = set() + + for node in all_nodes: + if node is self.root_node: + continue + nid = node.id + full_dev = node.component_data[FCT].value is not None + full_hst = node.component_data[FCT].host_value is not None + + # Full is the tree backbone, so aux data requires Full data. for ct in self.tree_components: - # 2. Basic size invariants - assert ( - self.component_evictable_size_[ct] >= 0 - ), f"component_evictable_size_[{ct}] = {self.component_evictable_size_[ct]} < 0" - assert ( - self.component_protected_size_[ct] >= 0 - ), f"component_protected_size_[{ct}] = {self.component_protected_size_[ct]} < 0" + if ct == FCT: + continue + cd = node.component_data[ct] + if cd.value is not None and not full_dev: + E(f"node {nid} {ct} device present but Full.value=None") + if cd.host_value is not None and not full_hst: + E(f"node {nid} {ct} host present but Full.host_value=None") - # 3. Verify LRU membership: tree nodes with data == LRU cache entries - lru = self.lru_lists[ct] + # Every node must keep Full data on at least one layer. + if not full_dev and not full_hst: + E(f"node {nid} dead: no Full device and no Full host") + + # Parent prefixes must keep data whenever the child does. + if node.parent is not None and node.parent is not self.root_node: + p_dev = node.parent.component_data[FCT].value is not None + p_hst = node.parent.component_data[FCT].host_value is not None + if full_dev and not p_dev: + E(f"node {nid} device present but parent {node.parent.id} evicted") + if full_hst and not p_hst: + E(f"node {nid} backed up but parent {node.parent.id} not backed up") + + # Lock hierarchy and counters must stay sane. + fl = node.component_data[FCT].lock_ref + for ct in self.tree_components: + cd = node.component_data[ct] + if cd.lock_ref < 0: + E(f"node {nid} {ct} lock_ref={cd.lock_ref}") + if cd.host_lock_ref < 0: + E(f"node {nid} {ct} host_lock_ref={cd.host_lock_ref}") + if ct != FCT and fl < cd.lock_ref: + E(f"node {nid} full_lock={fl} < {ct}_lock={cd.lock_ref}") + if cd.value is None and cd.lock_ref > 0: + E(f"node {nid} {ct} evicted but lock_ref={cd.lock_ref}") + + # Collect expected leaf qualification (single pass) + if self._is_device_leaf(node): + expected_dev_leaves.add(node) + if self._is_host_leaf(node): + expected_hst_leaves.add(node) + + # ── PART 3: Tracking structures ── + + # Device leaf set must match the expected leaves. + if self.evictable_device_leaves != expected_dev_leaves: + extra = self.evictable_device_leaves - expected_dev_leaves + missing = expected_dev_leaves - self.evictable_device_leaves + if extra: + E(f"D-leaf extra: {[n.id for n in list(extra)[:5]]}") + if missing: + E(f"D-leaf missing: {[n.id for n in list(missing)[:5]]}") + + # Host leaf set must match the expected leaves. + if self.evictable_host_leaves != expected_hst_leaves: + extra = self.evictable_host_leaves - expected_hst_leaves + missing = expected_hst_leaves - self.evictable_host_leaves + if extra: + E(f"H-leaf extra: {[n.id for n in list(extra)[:5]]}") + if missing: + E(f"H-leaf missing: {[n.id for n in list(missing)[:5]]}") + + # D-leaf ∩ H-leaf = ∅ + overlap = self.evictable_device_leaves & self.evictable_host_leaves + if overlap: + E( + f"[Leaf] {len(overlap)} in both sets: {[n.id for n in list(overlap)[:5]]}" + ) + + # Stale nodes: leaf sets must only contain tree-reachable nodes + stale = self.evictable_device_leaves - all_node_set + if stale: + E( + f"{len(stale)} stale nodes in device_leaves: {[n.id for n in list(stale)[:5]]}" + ) + stale = self.evictable_host_leaves - all_node_set + if stale: + E( + f"{len(stale)} stale nodes in host_leaves: {[n.id for n in list(stale)[:5]]}" + ) + + # Per-component LRU tracking + for ct in self.tree_components: + lru = self.lru_lists[ct] + if ct == FCT: + # Full uses leaf sets, not LRU + if len(lru.cache) > 0: + E(f"Full device LRU not empty: {len(lru.cache)}") + if len(self.host_lru_lists[ct].cache) > 0: + E(f"Full host LRU not empty: {len(self.host_lru_lists[ct].cache)}") + else: + # Aux device values must match the device LRU. tree_ids = { n.id for n in all_nodes - if n != self.root_node and n.component_data[ct].value is not None + if n is not self.root_node + and n.component_data[ct].value is not None } lru_ids = set(lru.cache.keys()) - assert tree_ids == lru_ids, ( - f"[{ct}] LRU membership mismatch: " - f"in_tree_not_lru={tree_ids - lru_ids}, " - f"in_lru_not_tree={lru_ids - tree_ids}" - ) - - # 4. Walk LRU doubly-linked list: verify structural integrity - # and that all nodes are unlocked (idle check) - visited = set() - x = lru.head.lru_next[ct] - prev = lru.head - while x != lru.tail: - assert ( - x.lru_prev[ct] == prev - ), f"[{ct}] broken prev link at node {x.id}" - assert ( - x.id in lru.cache - ), f"[{ct}] node {x.id} in linked list but not in cache dict" - assert x.id not in visited, f"[{ct}] cycle detected at node {x.id}" - assert x.component_data[ct].lock_ref == 0, ( - f"[{ct}] node {x.id} should not be locked when idle, " - f"lock_ref={x.component_data[ct].lock_ref}" + if tree_ids != lru_ids: + E( + f"{ct} device LRU: " + f"+tree={tree_ids - lru_ids}, +lru={lru_ids - tree_ids}" ) - visited.add(x.id) - prev = x - x = x.lru_next[ct] - assert len(visited) == len(lru.cache), ( - f"[{ct}] linked list has {len(visited)} nodes, " - f"cache dict has {len(lru.cache)}" + # Aux host-only states must match the host LRU. + host_lru = self.host_lru_lists[ct] + s3_ids = { + n.id + for n in all_nodes + if n is not self.root_node + and n.component_data[ct].value is None + and n.component_data[ct].host_value is not None + } + host_lru_ids = set(host_lru.cache.keys()) + if s3_ids != host_lru_ids: + E( + f"{ct} host LRU: " + f"+S3={s3_ids - host_lru_ids}, +lru={host_lru_ids - s3_ids}" + ) + # The same aux node must not appear in both device and host LRU. + inv5_overlap = lru_ids & host_lru_ids + if inv5_overlap: + E(f"{ct} in both device and host LRU: {inv5_overlap}") + # Linked-list integrity + self._check_lru_linked_list(lru, ct, "device", errors) + self._check_lru_linked_list(host_lru, ct, "host", errors) + + # ── PART 4: Size Accounting ── + for ct in self.tree_components: + evictable = 0 + protected = 0 + for n in all_nodes: + if n is self.root_node: + continue + cd = n.component_data[ct] + if cd.value is not None: + toks = len(cd.value) + if cd.lock_ref > 0: + protected += toks + else: + evictable += toks + if self.component_evictable_size_[ct] != evictable: + E( + f"[Size] {ct} evictable={self.component_evictable_size_[ct]} " + f"!= recomputed={evictable}" + ) + if self.component_protected_size_[ct] != protected: + E( + f"[Size] {ct} protected={self.component_protected_size_[ct]} " + f"!= recomputed={protected}" ) - # 5. Verify evictable size by walking unlocked LRU nodes - recomputed = 0 - x = lru.get_lru_no_lock() - while lru.in_list(x): - v = x.component_data[ct].value - recomputed += len(v) if v is not None else 0 - x = lru.get_prev_no_lock(x) - assert self.component_evictable_size_[ct] == recomputed, ( - f"[{ct}] evictable_size_={self.component_evictable_size_[ct]} " - f"!= recomputed={recomputed}" + # ── PART 5: Ongoing Operations ── + for nid, n in self.ongoing_write_through.items(): + if n not in all_node_set: + E(f"[Ongoing] write_through node {nid} not in tree") + elif n.component_data[FCT].lock_ref <= 0: + E( + f"[Ongoing] write_through node {nid} lock_ref={n.component_data[FCT].lock_ref}" + ) + for nid, n in self.ongoing_load_back.items(): + if n not in all_node_set: + E(f"[Ongoing] load_back node {nid} not in tree") + elif n.component_data[FCT].lock_ref <= 0: + E( + f"[Ongoing] load_back node {nid} lock_ref={n.component_data[FCT].lock_ref}" ) - except Exception as e: - logger.error(f"Unified RadixTree sanity check failed: {e}") + # ── Result ── + if errors: + msg = ( + f"Sanity check FAILED ({len(errors)} violations " + f"across {len(all_nodes)} nodes):\n" + + "\n".join(f" {e}" for e in errors) + ) + logger.error(msg) self.pretty_print() - raise + raise AssertionError(msg) + logger.debug( + f"Sanity check PASSED: {len(all_nodes)} nodes, " + f"{len(self.tree_components)} components" + ) + + def _check_lru_linked_list( + self, + lru: "UnifiedLRUList", + ct: ComponentType, + label: str, + errors: list[str], + ) -> None: + """Walk a LRU doubly-linked list, collect integrity errors.""" + pt = lru._pt # use LRU's own pointer slot + visited: set[int] = set() + x = lru.head.lru_next[pt] + prev = lru.head + while x is not None and x != lru.tail: + if x.lru_prev[pt] != prev: + errors.append(f"[{label}][{ct}] broken prev at node {x.id}") + if x.id not in lru.cache: + errors.append(f"[{label}][{ct}] node {x.id} in list not cache") + if x.id in visited: + errors.append(f"[{label}][{ct}] cycle at node {x.id}") + break + visited.add(x.id) + prev = x + x = x.lru_next[pt] + if x is None: + errors.append( + f"[{label}][{ct}] broken chain: lru_next is None " + f"after node {prev.id if hasattr(prev, 'id') else 'head'}" + ) + if len(visited) != len(lru.cache): + errors.append( + f"[{label}][{ct}] list={len(visited)} != cache={len(lru.cache)}" + ) def pretty_print(self) -> None: stack = [(self.root_node, 0)] @@ -1025,3 +1843,28 @@ class UnifiedRadixCache(BasePrefixCache): ) for child in node.children.values(): stack.append((child, indent + 2)) + + def _rebuild_host_leaf_sets(self) -> None: + """Rebuild evictable_host_leaves after L1-only reset.""" + stack = [self.root_node] + while stack: + node = stack.pop() + if node is not self.root_node: + self._update_evictable_leaf_sets(node) + stack.extend(node.children.values()) + + def _rebuild_host_lru_lists(self) -> None: + """Rebuild host_lru_lists for extra components after L1-only reset. + Walks the tree and adds nodes with host component data to the + appropriate host LRU list.""" + stack = [self.root_node] + while stack: + node = stack.pop() + if node is not self.root_node: + for ct in self.tree_components: + if ct == BASE_COMPONENT_TYPE: + continue # Full uses evictable_host_leaves, not host LRU + cd = node.component_data[ct] + if cd.host_value is not None: + self.host_lru_lists[ct].insert_mru(node) + stack.extend(node.children.values()) diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl.py b/test/registered/radix_cache/test_unified_radix_cache_kl.py index 6ac338518..0543bd5c7 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl.py @@ -262,5 +262,55 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): kill_process_tree(cls.process.pid) +# TODO(hzh): Currently, HiCache is not fully compatible with the CI CUDA13 environment; we need to wait for the fix before re-enabling the tests below. +# class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): +# """Mamba hybrid + UnifiedRadixCache.""" +# +# kl_threshold = 0.003 +# prefill_cache_assert = staticmethod( +# make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE) +# ) +# decode_cache_assert = staticmethod( +# make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL) +# ) +# +# @classmethod +# def setUpClass(cls): +# cls.model = MAMBA_MODEL +# cls.base_url = DEFAULT_URL_FOR_TEST +# cls.process = popen_launch_server( +# cls.model, +# cls.base_url, +# timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, +# other_args=[ +# "--tp-size", +# "4", +# "--chunked-prefill-size", +# "2048", +# "--mem-fraction-static", +# "0.85", +# "--mamba-scheduler-strategy", +# "extra_buffer", +# "--mamba-track-interval", +# str(MAMBA_TRACK_INTERVAL), +# "--enable-hierarchical-cache", +# "--hicache-ratio", +# "1.5", +# "--hicache-write-policy", +# "write_through", +# "--hicache-io-backend", +# "direct", +# "--hicache-mem-layout", +# "page_first_direct", +# ], +# env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, +# ) +# cls.input_ids = get_input_ids(cls.model, num_samples=18) +# +# @classmethod +# def tearDownClass(cls): +# kill_process_tree(cls.process.pid) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 069205311..66bf55994 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -3,6 +3,7 @@ import unittest from dataclasses import dataclass from typing import Optional +from unittest import mock import torch @@ -207,17 +208,17 @@ def build_fixture(cfg: CacheConfig): need_sort=False, ) - tree = UnifiedRadixCache( - params=CacheInitParams( - req_to_token_pool=req_to_token_pool, - token_to_kv_pool_allocator=allocator, - page_size=cfg.page_size, - disable=False, - sliding_window_size=cfg.sliding_window_size, - tree_components=cfg.components, - enable_mamba_extra_buffer=cfg.enable_mamba_extra_buffer, - ), + cache_init_params = CacheInitParams( + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=allocator, + page_size=cfg.page_size, + disable=False, + sliding_window_size=cfg.sliding_window_size, + tree_components=cfg.components, + enable_mamba_extra_buffer=cfg.enable_mamba_extra_buffer, ) + tree = UnifiedRadixCache(params=cache_init_params) + tree.cache_init_params = cache_init_params return tree, allocator, req_to_token_pool @@ -934,6 +935,644 @@ class UnifiedRadixCacheSuite: tree.sanity_check() + # ================================================================ + # Evict chain tests covering demotion, cascade, and tombstone cleanup. + # ================================================================ + + def test_evict_leaf_frees_all_components(self): + """Evicting a device leaf frees Full and all aux components atomically.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + seq = self._make_seq(1, 3) + self._insert(tree, allocator, req_to_token_pool, seq) + + full_before = tree.full_evictable_size() + mamba_before = tree.mamba_evictable_size() if self.cfg.has_mamba else 0 + swa_before = tree.swa_evictable_size() if self.cfg.has_swa else 0 + self.assertGreater(full_before, 0) + + result = tree.evict(EvictParams(num_tokens=full_before * 2)) + self.assertGreaterEqual(result.num_tokens_evicted, full_before) + self.assertEqual(tree.full_evictable_size(), 0) + if self.cfg.has_mamba: + self.assertEqual(tree.mamba_evictable_size(), 0) + if self.cfg.has_swa: + self.assertEqual(tree.swa_evictable_size(), 0) + tree.sanity_check() + + def test_evict_cascade_parent_becomes_d_leaf(self): + """After evicting a D-leaf child, parent may become a new D-leaf.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + base = self._make_seq(1, 2) + leaf = base + self._make_seq(500, 2) + self._insert(tree, allocator, req_to_token_pool, base) + self._insert(tree, allocator, req_to_token_pool, leaf) + + # Lock the base node to prevent it from being evicted + m_base = tree.match_prefix(MatchPrefixParams(key=RadixKey(base))) + lock_result = tree.inc_lock_ref(m_base.last_device_node) + + # Evict the leaf — parent (base) should become D-leaf after unlock + result = tree.evict(EvictParams(num_tokens=len(leaf))) + tree.sanity_check() + + tree.dec_lock_ref( + m_base.last_device_node, + DecLockRefParams( + swa_uuid_for_lock=getattr(lock_result, "swa_uuid_for_lock", None) + ), + ) + # After unlock, base should be in evictable_device_leaves + self.assertIn(m_base.last_device_node, tree.evictable_device_leaves) + tree.sanity_check() + + def test_evict_iterative_tombstone_cleanup(self): + """Tombstone cascade: evicting a leaf triggers cleanup up the tree.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + # Create a chain: root -> A -> B -> C (3 levels) + ps = self.cfg.page_size + chain = self._make_seq(1, 6) + self._insert(tree, allocator, req_to_token_pool, chain[: 2 * ps]) + self._insert(tree, allocator, req_to_token_pool, chain[: 4 * ps]) + self._insert(tree, allocator, req_to_token_pool, chain) + + initial_evictable = tree.full_evictable_size() + self.assertGreater(initial_evictable, 0) + + # Evict everything — tombstone cascade should clean up all + result = tree.evict(EvictParams(num_tokens=initial_evictable * 2)) + self.assertGreaterEqual(result.num_tokens_evicted, initial_evictable) + self.assertEqual(tree.full_evictable_size(), 0) + # Only root should remain + self.assertEqual(len(tree.root_node.children), 0) + tree.sanity_check() + + def test_evict_respects_lru_order(self): + """Older (less recently accessed) nodes are evicted first.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + ps = self.cfg.page_size + seq_old = self._make_seq(1, 2) + seq_new = self._make_seq(500, 2) + + self._insert(tree, allocator, req_to_token_pool, seq_old) + self._insert(tree, allocator, req_to_token_pool, seq_new) + + # Touch seq_new to make it MRU + tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new))) + + # Evict just enough for one sequence + tree.evict(EvictParams(num_tokens=len(seq_old))) + + # seq_old should be gone (LRU), seq_new should remain + m_old = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_old))) + m_new = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new))) + self.assertEqual(len(m_old.device_indices), 0) + self.assertEqual(len(m_new.device_indices), len(seq_new)) + tree.sanity_check() + + def test_evict_multiple_independent_leaves(self): + """Evicting multiple independent leaves works correctly.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + seqs = [self._make_seq(i * 100, 2) for i in range(4)] + for s in seqs: + self._insert(tree, allocator, req_to_token_pool, s) + + total = sum(len(s) for s in seqs) + self.assertEqual(tree.full_evictable_size(), total) + + # Evict half + half = total // 2 + result = tree.evict(EvictParams(num_tokens=half)) + self.assertGreaterEqual(result.num_tokens_evicted, half) + self.assertLessEqual(tree.full_evictable_size(), total - half) + tree.sanity_check() + + # Evict remainder + remaining = tree.full_evictable_size() + result = tree.evict(EvictParams(num_tokens=remaining * 2)) + self.assertGreaterEqual(result.num_tokens_evicted, remaining) + self.assertEqual(tree.full_evictable_size(), 0) + tree.sanity_check() + + def test_evict_shared_prefix_keeps_common_path(self): + """Evicting one branch preserves the shared prefix for other branch.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + base = self._make_seq(1, 2) + branch_a = base + self._make_seq(100, 2) + branch_b = base + self._make_seq(200, 2) + + self._insert(tree, allocator, req_to_token_pool, branch_a) + self._insert(tree, allocator, req_to_token_pool, branch_b) + + # Lock branch_b + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b))) + lr = tree.inc_lock_ref(m.last_device_node) + + # Evict — branch_a should go, base + branch_b stay + tree.evict(EvictParams(num_tokens=len(branch_a))) + + m_b = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b))) + self.assertEqual(len(m_b.device_indices), len(branch_b)) + + tree.dec_lock_ref( + m.last_device_node, + DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), + ) + tree.sanity_check() + + def test_evict_result_accounting_matches_actual(self): + """EvictResult.num_tokens_evicted matches actual size change.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + seqs = [self._make_seq(i * 100, 2) for i in range(5)] + for s in seqs: + self._insert(tree, allocator, req_to_token_pool, s) + + before = tree.full_evictable_size() + result = tree.evict(EvictParams(num_tokens=before)) + after = tree.full_evictable_size() + self.assertEqual(result.num_tokens_evicted, before - after) + tree.sanity_check() + + def test_evict_locked_subtree_skipped(self): + """All nodes in a locked path are skipped during eviction.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + seq_a = self._make_seq(1, 3) + seq_b = self._make_seq(500, 2) + self._insert(tree, allocator, req_to_token_pool, seq_a) + self._insert(tree, allocator, req_to_token_pool, seq_b) + + # Lock seq_a + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) + lr = tree.inc_lock_ref(m.last_device_node) + + # Try to evict everything + total = tree.full_evictable_size() + tree.full_protected_size() + result = tree.evict(EvictParams(num_tokens=total)) + + # seq_a should still be matchable (protected) + m2 = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) + self.assertEqual(len(m2.device_indices), len(seq_a)) + + tree.dec_lock_ref( + m.last_device_node, + DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), + ) + tree.sanity_check() + + def test_mamba_internal_tombstone_evict(self): + """Mamba eviction on internal node tombstones mamba only, keeps Full.""" + if not self.cfg.has_mamba: + self.skipTest("requires Mamba component") + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + # Create internal node with mamba and leaf extending it + seq_short = self._make_seq(1, 2) + seq_long = seq_short + self._make_seq(500, 2) + self._insert(tree, allocator, req_to_token_pool, seq_short) + self._insert(tree, allocator, req_to_token_pool, seq_long) + + # Evict only mamba + result = tree.evict(EvictParams(num_tokens=0, mamba_num=10)) + self.assertEqual(tree.mamba_evictable_size(), 0) + + # Full should still be accessible for at least the long seq base + # (mamba gone breaks match, but full data might still be in tree) + tree.sanity_check() + + def test_evict_reinsert_after_full_eviction(self): + """After evicting everything, new inserts work correctly.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + seq_a = self._make_seq(1, 2) + self._insert(tree, allocator, req_to_token_pool, seq_a) + tree.evict(EvictParams(num_tokens=len(seq_a) * 2)) + self.assertEqual(tree.full_evictable_size(), 0) + + # Re-insert + seq_b = self._make_seq(500, 2) + self._insert(tree, allocator, req_to_token_pool, seq_b) + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_b))) + self.assertEqual(len(m.device_indices), len(seq_b)) + tree.sanity_check() + + def test_swa_evict_internal_tombstone(self): + """SWA eviction on internal node cascades to lower-priority components.""" + if not self.cfg.has_swa: + self.skipTest("requires SWA component") + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + base = self._make_seq(1, 3) + leaf = base + self._make_seq(500, 3) + self._insert(tree, allocator, req_to_token_pool, base) + self._insert(tree, allocator, req_to_token_pool, leaf) + + swa_before = tree.swa_evictable_size() + result = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=swa_before * 2)) + self.assertEqual(tree.swa_evictable_size(), 0) + tree.sanity_check() + + def test_evict_d_leaf_set_consistency(self): + """evictable_device_leaves is consistent after mixed operations.""" + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + seqs = [self._make_seq(i * 100, 2) for i in range(6)] + for s in seqs: + self._insert(tree, allocator, req_to_token_pool, s) + + # Lock some, evict some, unlock + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[0]))) + lr = tree.inc_lock_ref(m.last_device_node) + + tree.evict(EvictParams(num_tokens=len(seqs[1]))) + tree.sanity_check() + + tree.dec_lock_ref( + m.last_device_node, + DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), + ) + tree.sanity_check() + + # Insert more + extra = self._make_seq(9000, 2) + self._insert(tree, allocator, req_to_token_pool, extra) + tree.sanity_check() + + # ================================================================ + # HiCache Unit Tests (real cache_controller D<->H backup/load) + # ================================================================ + + def _skip_unsupported_hicache_test(self): + if self.cfg.has_swa: + self.skipTest("HiCache tests do not run on SWA stacks") + return False + + def _init_hicache(self, tree): + import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler + + orig_kv_host_pool = assembler.MHATokenToKVPoolHost + orig_mamba_host_pool = assembler.MambaPoolHost + + def kv_host_pool_wrapper(*args, **kwargs): + kwargs["pin_memory"] = False + return orig_kv_host_pool(*args, **kwargs) + + def mamba_host_pool_wrapper(*args, **kwargs): + kwargs["pin_memory"] = False + return orig_mamba_host_pool(*args, **kwargs) + + patchers = [ + mock.patch.object( + assembler, + "MHATokenToKVPoolHost", + side_effect=kv_host_pool_wrapper, + ), + mock.patch.object( + assembler, + "MambaPoolHost", + side_effect=mamba_host_pool_wrapper, + ), + ] + for patcher in patchers: + patcher.start() + self.addCleanup(patcher.stop) + + server_args = ServerArgs( + model_path="dummy", + page_size=self.cfg.page_size, + hicache_io_backend="direct", + hicache_write_policy="write_through", + ) + set_global_server_args_for_scheduler(server_args) + tree.init_hicache(server_args, tree.cache_init_params) + tree.write_through_threshold = 1 << 30 + tree.load_back_threshold = 0 + + def _build_hicache_fixture(self): + fixture = build_fixture(self.cfg) + tree, _, _ = fixture + self._init_hicache(tree) + return fixture + + def _backup_node(self, tree, node): + backed_up = tree.write_backup(node, write_back=True) + self.assertGreater(backed_up, 0) + tree.writing_check(write_back=True) + return backed_up + + def _backup_tree(self, tree): + stack = [tree.root_node] + while stack: + node = stack.pop() + children = list(node.children.values()) + stack.extend(reversed(children)) + if node is not tree.root_node: + self._backup_node(tree, node) + + def _load_back_node(self, tree, node): + device_indices = tree.load_back(node) + self.assertIsNotNone(device_indices) + producer_id = tree.ready_to_load_host_cache() + self.assertNotEqual(producer_id, -1) + for _, finish_event, _ in list(tree.cache_controller.ack_load_queue): + finish_event.synchronize() + tree.loading_check() + return device_indices + + def _get_full_kv_pool(self, allocator): + kv_pool = allocator.get_kvcache() + return getattr(kv_pool, "full_kv_pool", kv_pool) + + def _fill_full_kv(self, allocator, indices, marker): + kv_pool = self._get_full_kv_pool(allocator) + layer_id = kv_pool.start_layer + k_buf = kv_pool.get_key_buffer(layer_id) + v_buf = kv_pool.get_value_buffer(layer_id) + k_buf[indices].fill_(marker) + v_buf[indices].fill_(marker + 1) + + def _snapshot_full_kv(self, allocator, indices): + kv_pool = self._get_full_kv_pool(allocator) + layer_id = kv_pool.start_layer + return ( + kv_pool.get_key_buffer(layer_id)[indices].float().cpu().clone(), + kv_pool.get_value_buffer(layer_id)[indices].float().cpu().clone(), + ) + + def _fill_mamba_state(self, req_to_token_pool, indices, marker): + if not self.cfg.has_mamba: + return + mamba_indices = indices.reshape(-1) + mamba_cache = req_to_token_pool.mamba_pool.mamba_cache + mamba_cache.temporal[:, mamba_indices].fill_(marker) + for offset, conv_buf in enumerate(mamba_cache.conv, start=1): + conv_buf[:, mamba_indices].fill_(marker + offset) + + def _snapshot_mamba_state(self, req_to_token_pool, indices): + mamba_indices = indices.reshape(-1) + mamba_cache = req_to_token_pool.mamba_pool.mamba_cache + return ( + mamba_cache.temporal[:, mamba_indices].float().cpu().clone(), + [conv[:, mamba_indices].float().cpu().clone() for conv in mamba_cache.conv], + ) + + def test_hicache_node_states(self): + """Verify device-only to device+host transition after real backup.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + seq = self._make_seq(1, 2) + self._insert(tree, allocator, req_to_token_pool, seq) + + # Find the leaf node + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) + node = m.last_device_node + self.assertIsNot(node, tree.root_node) + + ct = ComponentType.FULL + # S1: device only + self.assertIsNotNone(node.component_data[ct].value) + self.assertIsNone(node.component_data[ct].host_value) + self.assertFalse(node.backuped) + self.assertFalse(node.evicted) + + self._backup_node(tree, node) + self.assertIsNotNone(node.component_data[ct].value) + self.assertIsNotNone(node.component_data[ct].host_value) + self.assertTrue(node.backuped) + self.assertFalse(node.evicted) + tree.sanity_check() + + def test_hicache_evict_to_host(self): + """Evicting a backed-up device leaf demotes it to host-only state.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + seq = self._make_seq(1, 2) + self._insert(tree, allocator, req_to_token_pool, seq) + + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) + node = m.last_device_node + + self._backup_node(tree, node) + self.assertTrue(node.backuped) + + # Evict -> should demote to host (S3) + result = tree.evict(EvictParams(num_tokens=len(seq))) + self.assertGreaterEqual(result.num_tokens_evicted, len(seq)) + + # Node should now be evicted (S3) + self.assertTrue(node.evicted) + self.assertTrue(node.backuped) + self.assertIsNone(node.component_data[ComponentType.FULL].value) + self.assertIsNotNone(node.component_data[ComponentType.FULL].host_value) + + # Should be in host_leaves, not device_leaves + self.assertNotIn(node, tree.evictable_device_leaves) + self.assertIn(node, tree.evictable_host_leaves) + tree.sanity_check() + + def test_hicache_match_through_evicted_node(self): + """Match can traverse evicted (S3) nodes using host_value.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + base = self._make_seq(1, 2) + leaf = base + self._make_seq(500, 2) + self._insert(tree, allocator, req_to_token_pool, base) + self._insert(tree, allocator, req_to_token_pool, leaf) + + self._backup_tree(tree) + + # Lock leaf so only base can be evicted + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf))) + lr = tree.inc_lock_ref(m.last_device_node) + + # Evict base (inner node won't be evicted while child is locked) + tree.evict(EvictParams(num_tokens=len(base))) + + tree.dec_lock_ref( + m.last_device_node, + DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), + ) + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf))) + self.assertGreaterEqual(len(m.device_indices), len(base)) + tree.sanity_check() + + def test_hicache_d_leaf_h_leaf_mutual_exclusion(self): + """D-leaf and H-leaf sets are always disjoint.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + seqs = [self._make_seq(i * 100, 2) for i in range(4)] + for s in seqs: + self._insert(tree, allocator, req_to_token_pool, s) + + for i in range(2): + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i]))) + self._backup_node(tree, m.last_device_node) + + # Evict one backed-up node + tree.evict(EvictParams(num_tokens=len(seqs[0]))) + + # Check mutual exclusion + overlap = tree.evictable_device_leaves & tree.evictable_host_leaves + self.assertEqual(len(overlap), 0) + tree.sanity_check() + + def test_hicache_host_leaf_eviction(self): + """Evicting a host leaf removes the node from the tree entirely.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + seq = self._make_seq(1, 2) + self._insert(tree, allocator, req_to_token_pool, seq) + + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) + node = m.last_device_node + + self._backup_node(tree, node) + tree.evict(EvictParams(num_tokens=len(seq))) + + self.assertTrue(node.evicted) + self.assertIn(node, tree.evictable_host_leaves) + + # Now evict host + tree.evict_host(len(seq)) + + # Node should be removed from tree + self.assertNotIn(node, tree.evictable_host_leaves) + self.assertEqual(len(tree.root_node.children), 0) + tree.sanity_check() + + def test_hicache_load_back_restores_data(self): + """Loading back an evicted node restores the backed-up cache data.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + base = self._make_seq(1, 2) + self._insert(tree, allocator, req_to_token_pool, base) + + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(base))) + node = m.last_device_node + original_device_indices = m.device_indices.clone() + self._fill_full_kv(allocator, original_device_indices, marker=3) + expected_k, expected_v = self._snapshot_full_kv( + allocator, original_device_indices + ) + original_mamba_indices = None + expected_temporal = None + expected_conv = None + if self.cfg.has_mamba: + original_mamba_indices = node.component_data[ + ComponentType.MAMBA + ].value.clone() + self._fill_mamba_state(req_to_token_pool, original_mamba_indices, marker=11) + expected_temporal, expected_conv = self._snapshot_mamba_state( + req_to_token_pool, original_mamba_indices + ) + + self._backup_node(tree, node) + tree.evict(EvictParams(num_tokens=len(base))) + self.assertTrue(node.evicted) + self._fill_full_kv(allocator, original_device_indices, marker=9) + if original_mamba_indices is not None: + self._fill_mamba_state(req_to_token_pool, original_mamba_indices, marker=21) + + loaded_indices = self._load_back_node(tree, node) + self.assertFalse(node.evicted) + self.assertIsNotNone(node.component_data[ComponentType.FULL].value) + loaded_k, loaded_v = self._snapshot_full_kv(allocator, loaded_indices) + self.assertTrue(torch.equal(loaded_k, expected_k)) + self.assertTrue(torch.equal(loaded_v, expected_v)) + if self.cfg.has_mamba: + loaded_mamba_indices = node.component_data[ComponentType.MAMBA].value + loaded_temporal, loaded_conv = self._snapshot_mamba_state( + req_to_token_pool, loaded_mamba_indices + ) + self.assertTrue(torch.equal(loaded_temporal, expected_temporal)) + self.assertEqual(len(loaded_conv), len(expected_conv)) + for actual_conv, expected_conv_buf in zip(loaded_conv, expected_conv): + self.assertTrue(torch.equal(actual_conv, expected_conv_buf)) + tree.sanity_check() + + def test_hicache_backup_continuity(self): + """Backed-up nodes form a continuous prefix from the root.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._make_seq(1, 4) + ps = self.cfg.page_size + self._insert(tree, allocator, req_to_token_pool, chain[: 2 * ps]) + self._insert(tree, allocator, req_to_token_pool, chain) + + self._backup_tree(tree) + + # Verify: every backed-up node's parent is also backed-up (or root) + all_nodes = tree._collect_all_nodes() + for node in all_nodes: + if node is tree.root_node: + continue + if node.backuped: + parent = node.parent + self.assertTrue( + parent is tree.root_node or parent.backuped, + f"Backup continuity violated: node {node.id} backed up but parent {parent.id} not", + ) + tree.sanity_check() + + def test_hicache_evict_to_host_updates_aux_lru(self): + """Aux components move from device LRU to host LRU on device-to-host eviction.""" + if self._skip_unsupported_hicache_test(): + return + if not self.cfg.has_mamba: + self.skipTest("requires Mamba component") + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + seq = self._make_seq(1, 2) + self._insert(tree, allocator, req_to_token_pool, seq) + + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) + node = m.last_device_node + + # Check mamba is in device LRU + mamba_lru = tree.lru_lists[ComponentType.MAMBA] + host_mamba_lru = tree.host_lru_lists[ComponentType.MAMBA] + self.assertTrue(mamba_lru.in_list(node)) + self.assertFalse(host_mamba_lru.in_list(node)) + + self._backup_node(tree, node) + tree.evict(EvictParams(num_tokens=len(seq))) + + # Mamba should move to host LRU + self.assertFalse(mamba_lru.in_list(node)) + if node.component_data[ComponentType.MAMBA].host_value is not None: + self.assertTrue(host_mamba_lru.in_list(node)) + tree.sanity_check() + + def test_hicache_mixed_backup_evict_insert(self): + """Complex scenario: backup some, evict, insert new, verify invariants.""" + if self._skip_unsupported_hicache_test(): + return + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + seqs = [self._make_seq(i * 100, 2) for i in range(5)] + + # Insert all + for s in seqs: + self._insert(tree, allocator, req_to_token_pool, s) + tree.sanity_check() + + for i in range(3): + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i]))) + self._backup_node(tree, m.last_device_node) + + # Evict to free some tokens + tree.evict(EvictParams(num_tokens=len(seqs[0]) * 2)) + tree.sanity_check() + + # Insert new sequences + new_seqs = [self._make_seq(i * 1000, 2) for i in range(3)] + for s in new_seqs: + self._insert(tree, allocator, req_to_token_pool, s) + tree.sanity_check() + + # Verify D-leaf / H-leaf mutual exclusion + overlap = tree.evictable_device_leaves & tree.evictable_host_leaves + self.assertEqual(len(overlap), 0) + _CONFIGS: list[CacheConfig] = [ CacheConfig(page_size=1, components=(ComponentType.FULL,)),