[UnifiedRadixTree]: Support HiCache Framework for UnifiedRadixTree (#23316)

Co-authored-by: JINZ <1023553676@qq.com>
Co-authored-by: diemchai <diemchai@tencent.com>

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
 

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
 

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

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:

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",

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
-
+        # HiCache state: set to host KV pool when HiCache enabled
-    def node_has_component_data(self, node: UnifiedTreeNode) -> bool:
+        self._full_kv_pool_host = None
-        # Override so _for_each_component_lru includes Full in LRU operations
-        return node.component_data[self.component_type].value is not 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[
+        ct = self.component_type
-            self.component_type
+        new_parent.component_data[ct].lock_ref = child.component_data[ct].lock_ref
-        ].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]
+        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
-        return 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]
+        # Heap-based eviction from evictable_device_leaves, ordered by LRU.
-        while tracker[self.component_type] < request:
+        heap = [(n.last_access_time, n) for n in self.cache.evictable_device_leaves]
-            x = lru.get_leaf_lru_no_lock()
+        heapq.heapify(heap)
-            if x is None:
+        ct = self.component_type
-                break
+        while tracker[ct] < request and heap:
-            self.cache._evict_component_and_detach_lru(
+            _, x = heapq.heappop(heap)
-                x, self, is_leaf=True, tracker=tracker
+            if x not in self.cache.evictable_device_leaves:
-            )
+                continue
-            self.cache._cascade_evict(x, self, tracker)
+            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)

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
+        ct = self.component_type
-        new_parent.component_data[self.component_type].lock_ref = 0
+        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:
+    def evict_component(
-        value = node.component_data[self.component_type].value
+        self,
-        self.cache.req_to_token_pool.mamba_pool.free(value)
+        node: UnifiedTreeNode,
-        freed = len(value)
+        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
-        if not is_leaf:
+            cd.value = None
-            node.component_data[self.component_type].value = None
+
-        return freed
+        # 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 (
+        while tracker[ct] < request and x is not None and lru.in_list(x):
-            tracker[self.component_type] < 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:
-            assert x.component_data[self.component_type].value is not None
+                # D-leaf: atomic eviction of all components
-            if len(x.children) > 0:
+                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:
+        if value is not None and cd.lock_ref > 0:
-            assert 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

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 (
+        while tracker[ct] < request and x is not None and lru.in_list(x):
-            tracker[self.component_type] < 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:
-            assert x.component_data[self.component_type].value is not None
+                # D-leaf: atomic eviction of all components
-            if len(x.children) > 0:
+                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

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:
+    def node_has_component_data(
-        return node.component_data[self.component_type].value is not None
+        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.
+        *target* controls which layer(s) to evict:
-        For leaf nodes: free memory; the node will be deleted by caller.
+          - DEVICE: free device memory and tombstone (value = None).
-        Returns the number of tokens/slots freed.
+                    Host data is untouched.
-        - Full: frees full_value via token_to_kv_pool_allocator.
+          - HOST:   free host memory (host_value = None).
-        - SWA: frees swa value via swa_token_to_kv_pool_allocator;
+                    Device data is untouched.
-          only tombstones on internal nodes.
+          - ALL:    free both device and host memory.
-        - Mamba: frees mamba value via mamba_token_to_kv_pool_allocator;
+                    No tombstone — caller will delete the node.
-          only tombstones on internal nodes."""
+
+        Returns (device_freed, host_freed) token counts."""
         ...
 
     def eviction_priority(self, is_leaf: bool) -> int:

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()

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,8 +208,7 @@ def build_fixture(cfg: CacheConfig):
             need_sort=False,
         )
 
-    tree = UnifiedRadixCache(
+    cache_init_params = CacheInitParams(
-        params=CacheInitParams(
         req_to_token_pool=req_to_token_pool,
         token_to_kv_pool_allocator=allocator,
         page_size=cfg.page_size,
@@ -216,8 +216,9 @@ def build_fixture(cfg: CacheConfig):
         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,)),