[UnifiedRadixTree]: Support L3 HiStorage framework (#26062)

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -75,6 +75,7 @@ class InsertResult:
 
     prefix_len: int
     mamba_exist: bool = False
+    inserted_host_node: Any = None
 
 
 @dataclasses.dataclass
@@ -101,6 +102,7 @@ class IncLockRefResult:
 
     delta: Optional[int] = None
     swa_uuid_for_lock: Optional[int] = None
+    swa_uuid_for_host_lock: Optional[int] = None
     # Component nodes that were tombstones at acquire time. Replaying this set
     # at release prevents a short-lived lock from consuming a later load-back or
     # request lock after that tombstone becomes a valid device value.
@@ -112,6 +114,7 @@ class IncLockRefResult:
         """Convert to the corresponding DecLockRefParams for dec_lock_ref."""
         return DecLockRefParams(
             swa_uuid_for_lock=self.swa_uuid_for_lock,
+            swa_uuid_for_host_lock=self.swa_uuid_for_host_lock,
             skip_lock_node_ids={
                 component_type: set(node_ids)
                 for component_type, node_ids in self.skip_lock_node_ids.items()
@@ -124,6 +127,7 @@ class DecLockRefParams:
     """Parameters for dec_lock_ref operation."""
 
     swa_uuid_for_lock: Optional[int] = None
+    swa_uuid_for_host_lock: Optional[int] = None
     skip_lock_node_ids: dict[ComponentType, set[int]] = dataclasses.field(
         default_factory=dict
     )

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import json
 import logging
+import os
 import threading
 import time
+from queue import Queue
 from typing import TYPE_CHECKING, Any, Callable, List, Optional
 
 import torch
@@ -171,6 +174,7 @@ class HybridCacheController(BaseHiCacheController):
         enable_storage_metrics: bool = False,
     ):
         startup_storage_backend = storage_backend
+        self.extra_host_mem_release_queues: dict[PoolName, Queue] = {}
         super().__init__(
             token_to_kv_pool_allocator=token_to_kv_pool_allocator,
             mem_pool_host=mem_pool_host,
@@ -204,6 +208,10 @@ class HybridCacheController(BaseHiCacheController):
                 host_pools=getattr(mem_pool_host, "entries", None),
             )
 
+    def _start_storage_threads(self):
+        super()._start_storage_threads()
+        self._init_extra_host_mem_release_queues()
+
     def attach_storage_backend(
         self,
         storage_backend: str,
@@ -222,10 +230,133 @@ class HybridCacheController(BaseHiCacheController):
         for entry in host_pools or []:
             self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name)
 
+    @staticmethod
+    def parse_storage_backend_extra_config(
+        storage_backend_extra_config: Optional[str],
+    ) -> tuple[dict, int, float, float, bool]:
+        extra_config = {}
+        if storage_backend_extra_config:
+            if storage_backend_extra_config.startswith("@"):
+                path = storage_backend_extra_config[1:]
+                ext = os.path.splitext(path)[1].lower()
+                with open(path, "rb" if ext == ".toml" else "r") as f:
+                    if ext == ".json":
+                        extra_config = json.load(f)
+                    elif ext == ".toml":
+                        import tomllib
+
+                        extra_config = tomllib.load(f)
+                    elif ext in (".yaml", ".yml"):
+                        import yaml
+
+                        extra_config = yaml.safe_load(f)
+                    else:
+                        raise ValueError(
+                            f"Unsupported config file {path} (config format: {ext})"
+                        )
+            else:
+                extra_config = json.loads(storage_backend_extra_config)
+
+        prefetch_threshold = extra_config.pop("prefetch_threshold", 256)
+        prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", 1)
+        prefetch_timeout_per_ki_token = extra_config.pop(
+            "prefetch_timeout_per_ki_token", 0.25
+        )
+        hicache_storage_pass_prefix_keys = extra_config.pop(
+            "hicache_storage_pass_prefix_keys", False
+        )
+
+        if not isinstance(prefetch_threshold, int):
+            raise ValueError(
+                f"prefetch_threshold must be int, got {type(prefetch_threshold).__name__}"
+            )
+        if not isinstance(prefetch_timeout_base, (int, float)):
+            raise ValueError(
+                f"prefetch_timeout_base must be number, got {type(prefetch_timeout_base).__name__}"
+            )
+        if not isinstance(prefetch_timeout_per_ki_token, (int, float)):
+            raise ValueError(
+                "prefetch_timeout_per_ki_token must be number, got "
+                f"{type(prefetch_timeout_per_ki_token).__name__}"
+            )
+        if not isinstance(hicache_storage_pass_prefix_keys, bool):
+            raise ValueError(
+                "hicache_storage_pass_prefix_keys must be bool, got "
+                f"{type(hicache_storage_pass_prefix_keys).__name__}"
+            )
+
+        return (
+            extra_config,
+            prefetch_threshold,
+            float(prefetch_timeout_base),
+            float(prefetch_timeout_per_ki_token),
+            hicache_storage_pass_prefix_keys,
+        )
+
+    def clear_storage_backend(self) -> bool:
+        if not self.enable_storage:
+            logger.warning("Hierarchical cache storage backend is not enabled.")
+            return False
+        if not hasattr(self.storage_backend, "clear"):
+            logger.warning(
+                "Storage backend %s does not support clear operation.",
+                type(self.storage_backend).__name__,
+            )
+            return False
+        self.storage_backend.clear()
+        return True
+
+    def _init_extra_host_mem_release_queues(self) -> None:
+        self.extra_host_mem_release_queues = {}
+        entries = getattr(self.mem_pool_host, "entries", None) or []
+        anchor_entry = getattr(self.mem_pool_host, "anchor_entry", None)
+        for entry in entries:
+            if entry is anchor_entry or entry.is_primary_index_anchor:
+                continue
+            self.extra_host_mem_release_queues[entry.name] = Queue()
+
+    def _append_host_mem_release_pages(
+        self, release_queue: Queue, host_indices: torch.Tensor, page_size: int
+    ) -> None:
+        if host_indices.numel() == 0:
+            return
+        for page in host_indices.split(page_size):
+            release_queue.put(page)
+
+    def append_host_mem_release(
+        self,
+        host_indices: Optional[torch.Tensor] = None,
+        extra_pools: Optional[list[PoolTransfer]] = None,
+    ):
+        if host_indices is not None:
+            self._append_host_mem_release_pages(
+                self.host_mem_release_queue,
+                host_indices,
+                self.mem_pool_host.page_size,
+            )
+        for transfer in extra_pools or []:
+            if transfer.host_indices is None or transfer.host_indices.numel() == 0:
+                continue
+            entry = self.mem_pool_host.entry_map.get(transfer.name)
+            if (
+                entry is None
+                or entry.is_primary_index_anchor
+                or transfer.indices_from_pool is not None
+            ):
+                continue
+            release_queue = self.extra_host_mem_release_queues.get(transfer.name)
+            if release_queue is None:
+                continue
+            self._append_host_mem_release_pages(
+                release_queue, transfer.host_indices, entry.host_pool.page_size
+            )
+
     def reset(self):
         super().reset()
         if self.enable_storage:
             self.host_mem_release_queue.queue.clear()
+            for release_queue in self.extra_host_mem_release_queues.values():
+                release_queue.queue.clear()
             self.prefetch_tokens_occupied = 0
 
     def write(

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -679,6 +679,11 @@ class StackStrategy:
         load_cache_event,
         attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
         attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        storage_backend: Optional[str] = None,
+        storage_backend_extra_config: Optional[dict] = None,
+        prefetch_threshold: int = 256,
+        model_name: Optional[str] = None,
+        enable_storage_metrics: bool = False,
     ) -> StackBuildResult:
         raise NotImplementedError
 
@@ -704,6 +709,11 @@ class _DeepSeekV4Strategy(StackStrategy):
         load_cache_event,
         attn_cp_group=None,
         attn_tp_group=None,
+        storage_backend=None,
+        storage_backend_extra_config=None,
+        prefetch_threshold=256,
+        model_name=None,
+        enable_storage_metrics=False,
     ):
         from sglang.srt.mem_cache.base_prefix_cache import EvictParams
 
@@ -716,11 +726,15 @@ class _DeepSeekV4Strategy(StackStrategy):
             load_cache_event=load_cache_event,
             attn_cp_group=attn_cp_group,
             attn_tp_group=attn_tp_group,
-            storage_backend=None,
+            storage_backend=storage_backend,
             host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA),
             device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)),
+            prefetch_threshold=prefetch_threshold,
+            model_name=model_name,
+            storage_backend_extra_config=storage_backend_extra_config,
             pp_rank=params.pp_rank,
             pp_size=params.pp_size,
+            enable_storage_metrics=enable_storage_metrics,
         )
         sidecars = [
             SidecarPoolSpec(pool_name=name, indices_from_pool=src)
@@ -766,6 +780,11 @@ class _MambaStrategy(StackStrategy):
         load_cache_event,
         attn_cp_group=None,
         attn_tp_group=None,
+        storage_backend=None,
+        storage_backend_extra_config=None,
+        prefetch_threshold=256,
+        model_name=None,
+        enable_storage_metrics=False,
     ):
         from sglang.srt.mem_cache.base_prefix_cache import EvictParams
 
@@ -783,12 +802,16 @@ class _MambaStrategy(StackStrategy):
             load_cache_event=load_cache_event,
             attn_cp_group=attn_cp_group,
             attn_tp_group=attn_tp_group,
-            storage_backend=None,
+            storage_backend=storage_backend,
             use_mla=kvcache.use_mla,
             host_mamba_evict_fn=lambda n: cache.evict_host(n, ComponentType.MAMBA),
             device_mamba_evict_fn=lambda n: cache.evict(EvictParams(mamba_num=n)),
+            prefetch_threshold=prefetch_threshold,
+            model_name=model_name,
+            storage_backend_extra_config=storage_backend_extra_config,
             pp_rank=params.pp_rank,
             pp_size=params.pp_size,
+            enable_storage_metrics=enable_storage_metrics,
         )
         return StackBuildResult(
             host_pool_group=host_pool_group,
@@ -834,6 +857,11 @@ class _SwaStrategy(StackStrategy):
         load_cache_event,
         attn_cp_group=None,
         attn_tp_group=None,
+        storage_backend=None,
+        storage_backend_extra_config=None,
+        prefetch_threshold=256,
+        model_name=None,
+        enable_storage_metrics=False,
     ):
         from sglang.srt.mem_cache.base_prefix_cache import EvictParams
 
@@ -850,12 +878,16 @@ class _SwaStrategy(StackStrategy):
             load_cache_event=load_cache_event,
             attn_cp_group=attn_cp_group,
             attn_tp_group=attn_tp_group,
-            storage_backend=None,
+            storage_backend=storage_backend,
             use_mla=False,
             host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA),
             device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)),
+            prefetch_threshold=prefetch_threshold,
+            model_name=model_name,
+            storage_backend_extra_config=storage_backend_extra_config,
             pp_rank=params.pp_rank,
             pp_size=params.pp_size,
+            enable_storage_metrics=enable_storage_metrics,
         )
         return StackBuildResult(
             host_pool_group=host_pool_group,
@@ -887,6 +919,11 @@ class _DsaStrategy(StackStrategy):
         load_cache_event,
         attn_cp_group=None,
         attn_tp_group=None,
+        storage_backend=None,
+        storage_backend_extra_config=None,
+        prefetch_threshold=256,
+        model_name=None,
+        enable_storage_metrics=False,
     ):
         from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 
@@ -904,7 +941,7 @@ class _DsaStrategy(StackStrategy):
             load_cache_event=load_cache_event,
             attn_cp_group=attn_cp_group,
             attn_tp_group=attn_tp_group,
-            storage_backend=None,
+            storage_backend=storage_backend,
             use_mla=use_mla,
             override_kv_cache_dim=full_kv_pool.kv_cache_dim,
             sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost(
@@ -913,8 +950,12 @@ class _DsaStrategy(StackStrategy):
                 server_args.hicache_mem_layout,
                 allocator_type=server_args.hicache_storage_backend,
             ),
+            prefetch_threshold=prefetch_threshold,
+            model_name=model_name,
+            storage_backend_extra_config=storage_backend_extra_config,
             pp_rank=params.pp_rank,
             pp_size=params.pp_size,
+            enable_storage_metrics=enable_storage_metrics,
         )
         return StackBuildResult(
             host_pool_group=host_pool_group,
@@ -961,6 +1002,11 @@ class _PlainKvStrategy(StackStrategy):
         load_cache_event,
         attn_cp_group=None,
         attn_tp_group=None,
+        storage_backend=None,
+        storage_backend_extra_config=None,
+        prefetch_threshold=256,
+        model_name=None,
+        enable_storage_metrics=False,
     ):
         from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 
@@ -977,10 +1023,14 @@ class _PlainKvStrategy(StackStrategy):
             load_cache_event=load_cache_event,
             attn_cp_group=attn_cp_group,
             attn_tp_group=attn_tp_group,
-            storage_backend=None,
+            storage_backend=storage_backend,
             use_mla=use_mla,
+            prefetch_threshold=prefetch_threshold,
+            model_name=model_name,
+            storage_backend_extra_config=storage_backend_extra_config,
             pp_rank=params.pp_rank,
             pp_size=params.pp_size,
+            enable_storage_metrics=enable_storage_metrics,
         )
         return StackBuildResult(
             host_pool_group=host_pool_group,
@@ -1057,6 +1107,9 @@ def attach_hybrid_pool_to_unified_cache(
     load_cache_event,
     attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
     attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
+    storage_backend: Optional[str] = None,
+    storage_extra_config: Optional[dict] = None,
+    storage_prefetch_threshold: int = 256,
 ) -> None:
     """Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache."""
     try:
@@ -1071,6 +1124,11 @@ def attach_hybrid_pool_to_unified_cache(
             load_cache_event=load_cache_event,
             attn_cp_group=attn_cp_group,
             attn_tp_group=attn_tp_group,
+            storage_backend=storage_backend,
+            storage_backend_extra_config=storage_extra_config,
+            prefetch_threshold=storage_prefetch_threshold,
+            model_name=server_args.served_model_name,
+            enable_storage_metrics=cache._enable_metrics_flag,
         )
         _apply_stack_result(cache, kvcache, params, result)
     except Exception:

python/sglang/srt/mem_cache/unified_cache_components/full_component.py

@@ -65,7 +65,7 @@ class FullComponent(TreeComponent):
         # last_device_node, summing host_value lengths of evicted nodes.
         ct = self.component_type
         kv_host_hit = 0
-        node = result.last_host_node
+        node = result.best_match_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
@@ -155,9 +155,22 @@ class FullComponent(TreeComponent):
                 heapq.heappush(heap, (x.parent.last_access_time, x.parent))
 
     def acquire_component_lock(
-        self, node: UnifiedTreeNode, result: IncLockRefResult
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
     ) -> IncLockRefResult:
         ct = self.component_type
+
+        # Only the last host node needs to be protected.
+        if lock_host:
+            cd = node.component_data[ct]
+            if cd.host_value is None:
+                return result
+            cd.host_lock_ref += 1
+            self.cache._update_evictable_leaf_sets(node)
+            return result
+
         root = self.cache.root_node
         cur = node
 
@@ -185,9 +198,20 @@ class FullComponent(TreeComponent):
         return result
 
     def release_component_lock(
-        self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
     ) -> None:
         ct = self.component_type
+        if lock_host:
+            cd = node.component_data[ct]
+            if cd.host_value is None or cd.host_lock_ref == 0:
+                return
+            cd.host_lock_ref -= 1
+            self.cache._update_evictable_leaf_sets(node)
+            return
+
         root = self.cache.root_node
         skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
         cur = node
@@ -255,6 +279,7 @@ class FullComponent(TreeComponent):
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
+        **kw,
     ) -> None:
         ct = self.component_type
 

python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py

@@ -13,7 +13,7 @@ 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.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer
 from sglang.srt.mem_cache.unified_cache_components.tree_component import (
     CacheTransferPhase,
     ComponentType,
@@ -213,16 +213,28 @@ class MambaComponent(TreeComponent):
                 x = x_next
 
     def acquire_component_lock(
-        self, node: UnifiedTreeNode, result: IncLockRefResult
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
     ) -> IncLockRefResult:
         ct = self.component_type
+        if node is self.cache.root_node:
+            return result
         cd = node.component_data[ct]
-        value = cd.value
+        value = cd.host_value if lock_host else cd.value
         # A node in skip_lock_node_ids was a tombstone when this lock was acquired.
         if value is None:
             result.skip_lock_node_ids.setdefault(ct, set()).add(node.id)
             return result
 
+        if lock_host:
+            if cd.host_lock_ref == 0:
+                host_lru = self.cache.host_lru_lists[ct]
+                if host_lru.in_list(node):
+                    host_lru.remove_node(node)
+            cd.host_lock_ref += 1
+        else:
             if cd.lock_ref == 0:
                 vlen = len(value)
                 self.cache.component_evictable_size_[ct] -= vlen
@@ -231,16 +243,29 @@ class MambaComponent(TreeComponent):
         return result
 
     def release_component_lock(
-        self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
     ) -> None:
         ct = self.component_type
+        if node is self.cache.root_node:
+            return
         cd = node.component_data[ct]
         skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
         if node.id in skip_lock_node_ids:
             return
 
-        value = cd.value
+        value = cd.host_value if lock_host else cd.value
-        if value is not None and cd.lock_ref > 0:
+        if lock_host:
+            cd.host_lock_ref -= 1
+            if cd.host_lock_ref == 0 and cd.value is None and cd.host_value is not None:
+                host_lru = self.cache.host_lru_lists[ct]
+                if not host_lru.in_list(node):
+                    host_lru.insert_mru(node)
+            return
+
+        if cd.lock_ref > 0:
             if cd.lock_ref == 1:
                 vlen = len(value)
                 self.cache.component_evictable_size_[ct] += vlen
@@ -392,6 +417,35 @@ class MambaComponent(TreeComponent):
 
             return transfers if transfers else None
 
+        if phase == CacheTransferPhase.BACKUP_STORAGE:
+            cd = node.component_data[ct]
+            if cd.host_value is None or not node.hash_value:
+                return None
+            return [
+                PoolTransfer(
+                    name=PoolName.MAMBA,
+                    host_indices=cd.host_value,
+                    keys=[node.hash_value[-1]],
+                    hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                )
+            ]
+
+        if phase == CacheTransferPhase.PREFETCH:
+            host_indices = self._mamba_pool_host.alloc(1)
+            if host_indices is None:
+                self.cache.evict_host(1, ComponentType.MAMBA)
+                host_indices = self._mamba_pool_host.alloc(1)
+            if host_indices is None:
+                return []
+            return [
+                PoolTransfer(
+                    name=PoolName.MAMBA,
+                    host_indices=host_indices,
+                    keys=["__placeholder__"],
+                    hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                )
+            ]
+
         return None
 
     def commit_hicache_transfer(
@@ -399,6 +453,7 @@ class MambaComponent(TreeComponent):
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
+        **kw,
     ) -> None:
         ct = self.component_type
 
@@ -423,6 +478,41 @@ class MambaComponent(TreeComponent):
                 self.cache.lru_lists[ct].insert_mru(node)
                 self.cache.component_evictable_size_[ct] += count
 
+        elif phase == CacheTransferPhase.PREFETCH:
+            if not transfers:
+                return
+            transfer = transfers[0]
+            host_indices = transfer.host_indices
+            insert_result = kw.get("insert_result")
+            pool_storage_result = kw.get("pool_storage_result")
+            loaded = (
+                pool_storage_result is not None
+                and pool_storage_result.extra_pool_hit_pages.get(PoolName.MAMBA, 0) >= 1
+            )
+            target_node = (
+                insert_result.inserted_host_node if insert_result is not None else None
+            )
+            if (
+                host_indices is None
+                or target_node is None
+                or not loaded
+                or target_node.component_data[ct].host_value is not None
+            ):
+                self.cache.cache_controller.append_host_mem_release(
+                    extra_pools=[transfer]
+                )
+                if insert_result is not None:
+                    insert_result.mamba_exist = True
+                return
+
+            target_node.component_data[ct].host_value = host_indices.clone()
+            if target_node.component_data[ct].value is None:
+                host_lru = self.cache.host_lru_lists[ct]
+                if not host_lru.in_list(target_node):
+                    host_lru.insert_mru(target_node)
+            if insert_result is not None:
+                insert_result.mamba_exist = False
+
     def drive_host_eviction(
         self, num_tokens: int, tracker: dict[ComponentType, int]
     ) -> None:
@@ -445,4 +535,5 @@ class MambaComponent(TreeComponent):
                     x, self, target=EvictLayer.HOST, tracker=tracker
                 )
                 self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST)
+                self.cache._update_evictable_leaf_sets(x)
             x = x_next

python/sglang/srt/mem_cache/unified_cache_components/swa_component.py

@@ -350,7 +350,10 @@ class SWAComponent(TreeComponent):
                 x = x_next
 
     def acquire_component_lock(
-        self, node: UnifiedTreeNode, result: IncLockRefResult
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
     ) -> IncLockRefResult:
         ct = self.component_type
         root = self.cache.root_node
@@ -384,7 +387,10 @@ class SWAComponent(TreeComponent):
         return result
 
     def release_component_lock(
-        self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
     ) -> None:
         ct = self.component_type
         root = self.cache.root_node
@@ -484,6 +490,7 @@ class SWAComponent(TreeComponent):
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
+        **kw,
     ) -> None:
         ct = self.component_type
 

python/sglang/srt/mem_cache/unified_cache_components/tree_component.py

@@ -276,9 +276,12 @@ class TreeComponent(ABC):
 
     @abstractmethod
     def acquire_component_lock(
-        self, node: UnifiedTreeNode, result: IncLockRefResult
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
     ) -> IncLockRefResult:
-        """Increment lock_ref for this component, protecting nodes from
+        """Increment component lock refs, protecting nodes from
         eviction. Updates evictable → protected size on first lock.
         - Full: path-lock — walks from node up to root, incrementing
           lock_ref on every ancestor.
@@ -286,21 +289,31 @@ class TreeComponent(ABC):
           sliding window is filled; records a component_uuid at the
           boundary for release_component_lock to know where to stop.
         - Mamba: single-node lock — only increments lock_ref on the
-          node itself (mamba state is per-leaf, not per-path)."""
+          node itself (mamba state is per-leaf, not per-path).
+
+        When ``lock_host`` is True, the lock applies to host-side state:
+        - Full: single-node host lock.
+        - SWA: host window-lock with a dedicated host UUID boundary.
+        - Mamba: single-node host lock with host LRU detach."""
         ...
 
     @abstractmethod
     def release_component_lock(
-        self, node: UnifiedTreeNode, params: Optional[DecLockRefParams]
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
     ) -> None:
-        """Decrement lock_ref for this component, un-protecting nodes.
+        """Decrement component lock refs, un-protecting nodes.
         Updates protected → evictable size when lock_ref drops to 0.
         - Full: path-unlock — walks from node up to root, decrementing
           lock_ref on every ancestor.
         - SWA: path-unlock — walks upward, stopping at the node whose
           component_uuid matches the one recorded during acquire.
         - Mamba: single-node unlock — only decrements lock_ref on the
-          node itself."""
+          node itself.
+
+        When ``lock_host`` is True, the inverse host-side semantics apply."""
         ...
 
     def prepare_for_caching_req(
@@ -351,6 +364,7 @@ class TreeComponent(ABC):
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
+        **kw,
     ) -> None:
         """Post-transfer bookkeeping: store host indices, update LRU, etc."""
         pass

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -5,7 +5,8 @@ import threading
 import time
 from array import array
 from collections import defaultdict
-from functools import partial
+from functools import lru_cache, partial
+from queue import Empty
 from typing import TYPE_CHECKING, Any, Optional
 
 import torch
@@ -28,6 +29,9 @@ from sglang.srt.mem_cache.hicache_storage import (
     PoolTransfer,
     SidecarPoolSpec,
 )
+from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
+    HybridCacheController,
+)
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.unified_cache_components import (
     _NUM_COMPONENT_TYPES,
@@ -42,6 +46,8 @@ from sglang.srt.mem_cache.unified_cache_components import (
     TreeComponent,
     get_and_increase_time_counter,
 )
+from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value
+from sglang.srt.observability.metrics_collector import StorageMetricsCollector
 from sglang.srt.session.streaming_session import StreamingSession
 
 if TYPE_CHECKING:
@@ -93,6 +99,18 @@ class UnifiedTreeNode:
     def __lt__(self, other: UnifiedTreeNode):
         return self.last_access_time < other.last_access_time
 
+    def get_last_hash_value(self) -> Optional[str]:
+        if self.hash_value is None or len(self.hash_value) == 0:
+            return None
+        return self.hash_value[-1]
+
+    @lru_cache(maxsize=1)
+    def get_prefix_hash_values(self, node: UnifiedTreeNode) -> list[str]:
+        if node is None or node.hash_value is None:
+            return []
+
+        return node.get_prefix_hash_values(node.parent) + node.hash_value
+
 
 class UnifiedLRUList:
     def __init__(
@@ -220,6 +238,10 @@ class UnifiedRadixCache(BasePrefixCache):
 
         if params.enable_metrics:
             self.init_metrics_collector()
+        self._enable_metrics_flag = params.enable_metrics
+        self.enable_storage_metrics = False
+        self.storage_metrics_collector: Optional[StorageMetricsCollector] = None
+        self.extra_metric_labels = None
 
         assert params.tree_components is not None
         self.tree_components = tuple(params.tree_components)
@@ -248,6 +270,11 @@ class UnifiedRadixCache(BasePrefixCache):
         # HiCache D↔H defaults (overridden by init_hicache)
         self.cache_controller = None
         self.write_through_threshold = 256
+        self.prefetch_stop_policy = "best_effort"
+        self.prefetch_threshold = 256
+        self.prefetch_timeout_base = 1.0
+        self.prefetch_timeout_per_page = 0.25
+        self.hicache_storage_pass_prefix_keys = False
 
         self.reset()
         logger.info(f"Init Unified RadixTree with components {self.tree_components}")
@@ -260,6 +287,7 @@ class UnifiedRadixCache(BasePrefixCache):
         self.root_node = UnifiedTreeNode(self.tree_components)
         self.root_node.key = RadixKey(array("q"), None)
         self.root_node.component_data[BASE_COMPONENT_TYPE].value = []
+        self.root_node.hash_value = []
         for ct in self.tree_components:
             self.root_node.component_data[ct].lock_ref = 1
         self.component_evictable_size_ = {ct: 0 for ct in self.tree_components}
@@ -281,12 +309,14 @@ class UnifiedRadixCache(BasePrefixCache):
         ] = {}
         self.ongoing_load_back: dict[int, tuple[UnifiedTreeNode, DecLockRefParams]] = {}
         self.enable_storage = False
+        self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
         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()
+            self.enable_storage = self.cache_controller.enable_storage
 
         self._empty_match_result = MatchResult(
             device_indices=torch.empty(
@@ -316,6 +346,26 @@ class UnifiedRadixCache(BasePrefixCache):
 
         self.load_cache_event = threading.Event()
         self.sidecar_pool_specs.clear()
+        self.extra_metric_labels = server_args.extra_metric_labels
+
+        # Parse storage config once, share with assembler and tree
+        storage_backend = server_args.hicache_storage_backend
+        storage_extra_config = None
+        storage_prefetch_threshold = 256
+        prefetch_timeout_base = 1.0
+        prefetch_timeout_per_ki_token = 0.25
+        hicache_storage_pass_prefix_keys = False
+        if storage_backend is not None:
+            (
+                storage_extra_config,
+                storage_prefetch_threshold,
+                prefetch_timeout_base,
+                prefetch_timeout_per_ki_token,
+                hicache_storage_pass_prefix_keys,
+            ) = HybridCacheController.parse_storage_backend_extra_config(
+                server_args.hicache_storage_backend_extra_config
+            )
+
         attach_hybrid_pool_to_unified_cache(
             self,
             params,
@@ -323,6 +373,9 @@ class UnifiedRadixCache(BasePrefixCache):
             load_cache_event=self.load_cache_event,
             attn_cp_group=params.attn_cp_cache_group,
             attn_tp_group=params.attn_tp_cache_group,
+            storage_backend=storage_backend,
+            storage_extra_config=storage_extra_config,
+            storage_prefetch_threshold=storage_prefetch_threshold,
         )
 
         # State initialization
@@ -330,13 +383,18 @@ class UnifiedRadixCache(BasePrefixCache):
             1 if server_args.hicache_write_policy == "write_through" else 2
         )
         self.load_back_threshold = 256
+        self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
 
-        logger.info(
+        if storage_backend is not None:
-            f"HiCache D\u2194H initialized: "
+            self._apply_storage_runtime_config(
-            f"host_pool_size={self.host_pool_group.size}, "
+                storage_backend=storage_backend,
-            f"write_policy={server_args.hicache_write_policy}, "
+                prefetch_threshold=storage_prefetch_threshold,
-            f"tp_world_size={self.tp_world_size}, "
+                prefetch_timeout_base=prefetch_timeout_base,
-            f"transfer_layer_num={self.cache_controller.layer_num}"
+                prefetch_timeout_per_ki_token=prefetch_timeout_per_ki_token,
+                hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys,
+                enable_storage=self.cache_controller.enable_storage,
+                enable_storage_metrics=self._enable_metrics_flag,
+                extra_metric_labels=self.extra_metric_labels,
             )
 
     def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None:
@@ -435,6 +493,29 @@ class UnifiedRadixCache(BasePrefixCache):
         # TODO: delta is not aggregated from components; no caller uses it yet.
         return DecLockRefResult()
 
+    def inc_host_lock_ref(self, node: Any) -> IncLockRefResult:
+        if self.disable:
+            return IncLockRefResult()
+        result = IncLockRefResult()
+        for component in self._components_tuple:
+            result = component.acquire_component_lock(
+                node=node, result=result, lock_host=True
+            )
+
+        self._update_evictable_leaf_sets(node)
+        return result
+
+    def dec_host_lock_ref(
+        self, node: Any, params: Optional[DecLockRefParams] = None
+    ) -> DecLockRefResult:
+        if self.disable:
+            return DecLockRefResult()
+        for component in self._components_tuple:
+            component.release_component_lock(node=node, params=params, lock_host=True)
+
+        self._update_evictable_leaf_sets(node)
+        return DecLockRefResult()
+
     def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs) -> None:
         if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs):
             return
@@ -703,13 +784,15 @@ class UnifiedRadixCache(BasePrefixCache):
             cur_time -= 0.00001
             node_update = node_update.parent
 
-        # Walk up to find last_host_node for full component.
+        # last_host_node will be used as the starting node for the subsequent
-        if self.cache_controller is None:
+        # `prefetch_from_storage` flow. We directly use best_match_node here,
-            last_host_node = best_match_device_node
+        # because best_match_node represents the node where all components
-        else:
+        # have reached consensus on both device & host availability.
-            last_host_node = best_match_node
+        last_host_node = (
-            while last_host_node is not self.root_node and not last_host_node.backuped:
+            best_match_node
-                last_host_node = last_host_node.parent
+            if self.cache_controller is not None
+            else best_match_device_node
+        )
 
         if best_match_device_value_len > 0:
             device_indices = torch.cat(value[:best_match_device_value_len])
@@ -744,6 +827,9 @@ class UnifiedRadixCache(BasePrefixCache):
 
         child.parent = new_node
         child.key = child.key[split_len:]
+        new_node.hash_value, child.hash_value = split_node_hash_value(
+            child.hash_value, split_len, self.page_size
+        )
 
         for component in self._components_tuple:
             component.redistribute_on_node_split(new_parent=new_node, child=child)
@@ -778,6 +864,8 @@ class UnifiedRadixCache(BasePrefixCache):
         new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone()
         parent.children[key.child_key(self.page_size)] = new_node
         self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value)
+        if self.enable_storage:
+            new_node.hash_value = compute_node_hash_values(new_node, self.page_size)
 
         self._update_evictable_leaf_sets(new_node)
         self._update_evictable_leaf_sets(parent)
@@ -894,6 +982,58 @@ class UnifiedRadixCache(BasePrefixCache):
             self._inc_hit_count(target_node, params.chunked)
         return result
 
+    def _insert_helper_host(
+        self,
+        node: UnifiedTreeNode,
+        key: RadixKey,
+        host_value: torch.Tensor,
+        hash_value: list[str],
+    ) -> InsertResult:
+        total_len = len(key)
+        self._touch_node(node)
+        if total_len == 0:
+            return InsertResult(prefix_len=0, mamba_exist=True)
+
+        child_key = key.child_key(self.page_size)
+        matched_length = 0
+        while len(key) > 0 and child_key in node.children:
+            node = node.children[child_key]
+            self._touch_node(node)
+            prefix_len = node.key.match(key, page_size=self.page_size)
+
+            key = key[prefix_len:]
+            host_value = host_value[prefix_len:]
+            hash_value = hash_value[prefix_len // self.page_size :]
+            matched_length += prefix_len
+
+            if prefix_len < len(node.key):
+                node = self._split_node(node.key, node, prefix_len)
+
+            if len(key):
+                child_key = key.child_key(self.page_size)
+
+        result = InsertResult(
+            prefix_len=matched_length,
+        )
+        if len(key) == 0:
+            if (
+                node is not self.root_node
+                and node.component_data[BASE_COMPONENT_TYPE].host_value is not None
+            ):
+                result.inserted_host_node = node
+            return result
+
+        new_node = UnifiedTreeNode(self.tree_components)
+        new_node.parent = node
+        new_node.key = key
+        new_node.hash_value = hash_value
+        new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone()
+        node.children[child_key] = new_node
+        self._update_evictable_leaf_sets(new_node)
+        self._update_evictable_leaf_sets(node)
+        result.inserted_host_node = new_node
+        return result
+
     # ---- Evict Helpers ----
 
     def _cascade_evict(
@@ -1126,6 +1266,7 @@ class UnifiedRadixCache(BasePrefixCache):
                 and self.cache_controller.write_policy == "write_back"
             ):
                 self.write_backup(node, write_back=True)
+                self.writing_check(write_back=True)
                 self._evict_to_host(node, tracker)
                 return
             else:
@@ -1368,6 +1509,511 @@ class UnifiedRadixCache(BasePrefixCache):
         if not node.backuped and node.hit_count >= self.write_through_threshold:
             self.write_backup(node)
 
+    def write_backup_storage(self, node: UnifiedTreeNode) -> None:
+        if (
+            not self.enable_storage
+            or self.cache_controller is None
+            or not node.backuped
+        ):
+            return
+
+        prefix_keys = None
+        if self.hicache_storage_pass_prefix_keys:
+            prefix_keys = node.get_prefix_hash_values(node.parent)
+
+        comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
+        for comp in self._components_tuple:
+            if comp.component_type == BASE_COMPONENT_TYPE:
+                continue
+            transfers = comp.build_hicache_transfers(
+                node,
+                CacheTransferPhase.BACKUP_STORAGE,
+            )
+            if transfers:
+                comp_xfers[comp.component_type] = transfers
+
+        kv_xfer = PoolTransfer(
+            name=PoolName.KV,
+            host_indices=node.component_data[BASE_COMPONENT_TYPE].host_value,
+            keys=node.hash_value,
+        )
+        sidecar_xfers = self._build_sidecar_transfers(
+            CacheTransferPhase.BACKUP_STORAGE, kv_xfer, comp_xfers
+        )
+        aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
+        aux_xfers.extend(sidecar_xfers)
+
+        operation_id = self.cache_controller.write_storage(
+            node.component_data[BASE_COMPONENT_TYPE].host_value,
+            node.key.token_ids,
+            node.hash_value,
+            prefix_keys,
+            extra_pools=aux_xfers or None,
+        )
+        self.ongoing_backup[operation_id] = (
+            node,
+            self.inc_host_lock_ref(node).to_dec_params(),
+        )
+
+    def prefetch_from_storage(
+        self,
+        req_id: str,
+        last_host_node: UnifiedTreeNode,
+        new_input_tokens: list[int],
+        last_hash: Optional[str] = None,
+        prefix_keys: Optional[list[str]] = None,
+    ) -> None:
+        if not self.enable_storage or self.cache_controller is None:
+            return
+
+        extra_key = last_host_node.key.extra_key if last_host_node.key else None
+        prefetch_key = RadixKey(
+            new_input_tokens,
+            extra_key=extra_key,
+            is_bigram=self.is_eagle,
+        ).page_aligned(self.page_size)
+        prefetch_length = len(prefetch_key)
+        if (
+            prefetch_length < self.prefetch_threshold
+            or self.cache_controller.prefetch_rate_limited()
+        ):
+            return
+
+        anchor_lock_params = self.inc_host_lock_ref(last_host_node).to_dec_params()
+        host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
+        if host_indices is None:
+            self.evict_host(prefetch_length)
+            host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
+        if host_indices is None:
+            available_size = self.cache_controller.mem_pool_host.available_size()
+            prefetch_length = available_size - (available_size % self.page_size)
+            if prefetch_length >= self.prefetch_threshold:
+                prefetch_key = prefetch_key[:prefetch_length]
+                host_indices = self.cache_controller.mem_pool_host.alloc(
+                    prefetch_length
+                )
+            else:
+                self.dec_host_lock_ref(last_host_node, anchor_lock_params)
+                return
+        if host_indices is None:
+            self.dec_host_lock_ref(last_host_node, anchor_lock_params)
+            return
+
+        comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
+        alloc_failed = False
+        for comp in self._components_tuple:
+            if comp.component_type == BASE_COMPONENT_TYPE:
+                continue
+            transfers = comp.build_hicache_transfers(
+                last_host_node,
+                CacheTransferPhase.PREFETCH,
+                token_ids=prefetch_key.token_ids,
+                prefetch_tokens=len(prefetch_key),
+                last_hash=last_hash,
+            )
+            if transfers == []:
+                alloc_failed = True
+                break
+            if transfers:
+                comp_xfers[comp.component_type] = transfers
+        kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices)
+        sidecar_xfers = self._build_sidecar_transfers(
+            CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers
+        )
+        if alloc_failed:
+            self.cache_controller.append_host_mem_release(
+                host_indices=host_indices,
+                extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
+            )
+            self.dec_host_lock_ref(last_host_node, anchor_lock_params)
+            return
+
+        aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
+        aux_xfers.extend(sidecar_xfers)
+        operation = self.cache_controller.prefetch(
+            req_id,
+            host_indices,
+            prefetch_key.token_ids,
+            last_hash,
+            prefix_keys,
+            extra_pools=aux_xfers or None,
+        )
+        self.ongoing_prefetch[req_id] = (
+            last_host_node,
+            prefetch_key,
+            host_indices,
+            operation,
+            anchor_lock_params,
+            comp_xfers,
+        )
+        self.cache_controller.prefetch_tokens_occupied += len(prefetch_key)
+
+    def _prefetch_timeout_check_linear_func(self, operation) -> bool:
+        return (
+            time.monotonic() - operation.start_time
+            > self.prefetch_timeout_base
+            + len(operation.hash_value) * self.prefetch_timeout_per_page
+        )
+
+    def can_terminate_prefetch(self, operation) -> bool:
+        if self.prefetch_stop_policy == "best_effort":
+            return True
+
+        if len(operation.hash_value) == 0:
+            completed = False
+        else:
+            completed = (
+                operation.completed_tokens == len(operation.hash_value) * self.page_size
+            )
+
+        if self.prefetch_stop_policy == "wait_complete":
+            can_terminate = completed
+        elif self.prefetch_stop_policy == "timeout":
+            can_terminate = completed or self._prefetch_timeout_check_linear_func(
+                operation
+            )
+        else:
+            return True
+
+        operation_terminated = operation.is_terminated()
+        states = torch.tensor(
+            [1 - int(can_terminate), int(operation_terminated)],
+            dtype=torch.int,
+        )
+        if self.tp_world_size > 1:
+            torch.distributed.all_reduce(
+                states, op=torch.distributed.ReduceOp.MAX, group=self.tp_group
+            )
+        can_terminate = states[0].item() == 0
+        operation_terminated = states[1].item() == 1
+        return can_terminate or operation_terminated
+
+    def check_prefetch_progress(self, req_id: str) -> bool:
+        if req_id not in self.ongoing_prefetch:
+            return True
+
+        (
+            last_host_node,
+            prefetch_key,
+            host_indices,
+            operation,
+            anchor_lock_params,
+            comp_xfers,
+        ) = self.ongoing_prefetch[req_id]
+        if operation.host_indices is None:
+            return True
+        if not self.can_terminate_prefetch(operation):
+            return False
+
+        completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
+            operation
+        )
+        min_completed_tokens = completed_tokens
+        if self.tp_world_size > 1:
+            completed_tokens_tensor = torch.tensor(
+                min_completed_tokens, dtype=torch.int
+            )
+            torch.distributed.all_reduce(
+                completed_tokens_tensor,
+                op=torch.distributed.ReduceOp.MIN,
+                group=self.tp_group,
+            )
+            min_completed_tokens = int(completed_tokens_tensor.item())
+
+        fetched_key = prefetch_key[:min_completed_tokens]
+        insert_result = self._insert_helper_host(
+            last_host_node,
+            fetched_key,
+            host_indices[:min_completed_tokens],
+            hash_value[: min_completed_tokens // self.page_size],
+        )
+
+        for ct, xfers in comp_xfers.items():
+            self.components[ct].commit_hicache_transfer(
+                last_host_node,
+                CacheTransferPhase.PREFETCH,
+                xfers,
+                insert_result=insert_result,
+                pool_storage_result=operation.pool_storage_result,
+            )
+
+        self.cache_controller.mem_pool_host.free(
+            host_indices[: insert_result.prefix_len]
+        )
+        self.cache_controller.append_host_mem_release(
+            host_indices[min_completed_tokens:completed_tokens]
+        )
+        self.dec_host_lock_ref(last_host_node, anchor_lock_params)
+        del self.ongoing_prefetch[req_id]
+        self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
+
+        loaded_from_storage = min_completed_tokens - insert_result.prefix_len
+        self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage
+        logger.info(
+            "HiCache prefetch success req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d tail_release=%d occupied=%d",
+            req_id,
+            completed_tokens,
+            min_completed_tokens,
+            insert_result.prefix_len,
+            loaded_from_storage,
+            completed_tokens - min_completed_tokens,
+            self.cache_controller.prefetch_tokens_occupied,
+        )
+        if self.enable_storage_metrics and self.storage_metrics_collector is not None:
+            self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage)
+        return True
+
+    def terminate_prefetch(self, req_id: str) -> None:
+        if req_id not in self.ongoing_prefetch:
+            return
+        _, _, _, operation, _, _ = self.ongoing_prefetch[req_id]
+        if operation.host_indices is None:
+            return
+        operation.mark_terminate()
+
+    def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
+        return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
+
+    def release_aborted_request(self, rid: str) -> None:
+        self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
+        if rid not in self.ongoing_prefetch:
+            return
+
+        (
+            last_host_node,
+            prefetch_key,
+            host_indices,
+            operation,
+            anchor_lock_params,
+            comp_xfers,
+        ) = self.ongoing_prefetch[rid]
+        if operation.host_indices is None:
+            return
+
+        completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
+        if self.tp_world_size > 1:
+            torch.distributed.barrier(group=self.tp_group)
+        self.dec_host_lock_ref(last_host_node, anchor_lock_params)
+        del self.ongoing_prefetch[rid]
+        self.cache_controller.append_host_mem_release(
+            host_indices=host_indices[:completed_tokens],
+            extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
+        )
+        self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
+
+    def _drain_storage_control_queues_impl(
+        self,
+        n_revoke: Optional[int],
+        n_backup: Optional[int],
+        n_release: Optional[int],
+        extra_release_counts: Optional[dict[PoolName, int]],
+        log_metrics: bool,
+    ) -> None:
+        cc = self.cache_controller
+
+        def _drain_queue(q, limit: Optional[int]):
+            drained = 0
+            while limit is None or drained < limit:
+                try:
+                    item = q.get_nowait()
+                except Empty:
+                    break
+                drained += 1
+                yield item
+
+        def _drain_revoke():
+            drained = 0
+            for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke):
+                info = self.ongoing_prefetch.pop(req_id, None)
+                if info is None:
+                    continue
+                drained += 1
+                (
+                    last_host_node,
+                    prefetch_key,
+                    _host_indices,
+                    _operation,
+                    anchor_lock_params,
+                    comp_xfers,
+                ) = info
+                cc.append_host_mem_release(
+                    extra_pools=[x for xfers in comp_xfers.values() for x in xfers]
+                )
+                self.dec_host_lock_ref(last_host_node, anchor_lock_params)
+                cc.prefetch_tokens_occupied -= len(prefetch_key)
+                if cc.prefetch_tokens_occupied < 0:
+                    cc.prefetch_tokens_occupied = 0
+            return drained
+
+        def _drain_backup():
+            drained = 0
+            for operation in _drain_queue(cc.ack_backup_queue, n_backup):
+                drained += 1
+                entry = self.ongoing_backup.pop(operation.id, None)
+                if entry is not None:
+                    node, lock_params = entry
+                    self.dec_host_lock_ref(node, lock_params)
+                if (
+                    log_metrics
+                    and self.enable_storage_metrics
+                    and self.storage_metrics_collector is not None
+                ):
+                    self.storage_metrics_collector.log_backuped_tokens(
+                        operation.completed_tokens
+                    )
+            return drained
+
+        def _drain_release():
+            host_indices_list = []
+            released_tokens = 0
+            for host_indices in _drain_queue(cc.host_mem_release_queue, n_release):
+                host_indices_list.append(host_indices)
+                released_tokens += len(host_indices)
+            if host_indices_list:
+                cc.mem_pool_host.free(torch.cat(host_indices_list, dim=0))
+            return len(host_indices_list), released_tokens
+
+        def _drain_extra_release():
+            drained: dict[PoolName, tuple[int, int]] = {}
+            if not extra_release_counts:
+                return drained
+            for pool_name, limit in extra_release_counts.items():
+                release_queue = cc.extra_host_mem_release_queues.get(pool_name)
+                if release_queue is None:
+                    continue
+                host_indices_list = []
+                released_tokens = 0
+                for host_indices in _drain_queue(release_queue, limit):
+                    host_indices_list.append(host_indices)
+                    released_tokens += len(host_indices)
+                if host_indices_list:
+                    entry = cc.mem_pool_host.entry_map.get(pool_name)
+                    if entry is not None:
+                        entry.host_pool.free(torch.cat(host_indices_list, dim=0))
+                drained[pool_name] = (len(host_indices_list), released_tokens)
+            return drained
+
+        _drain_revoke()
+        _drain_backup()
+        _drain_release()
+        _drain_extra_release()
+
+    def drain_storage_control_queues(self) -> None:
+        cc = self.cache_controller
+        extra_release_queues = getattr(cc, "extra_host_mem_release_queues", {})
+        extra_pool_names = list(extra_release_queues)
+        local_qsize_list = [
+            cc.prefetch_revoke_queue.qsize(),
+            cc.ack_backup_queue.qsize(),
+            cc.host_mem_release_queue.qsize(),
+            *[
+                extra_release_queues[pool_name].qsize()
+                for pool_name in extra_pool_names
+            ],
+        ]
+        qsizes = torch.tensor(
+            local_qsize_list,
+            dtype=torch.int,
+        )
+        if self.tp_world_size > 1:
+            torch.distributed.all_reduce(
+                qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
+            )
+        qsize_list = list(map(int, qsizes.tolist()))
+        n_revoke, n_backup, n_release = qsize_list[:3]
+        extra_release_counts = {
+            pool_name: count
+            for pool_name, count in zip(extra_pool_names, qsize_list[3:])
+        }
+        self._drain_storage_control_queues_impl(
+            n_revoke=n_revoke,
+            n_backup=n_backup,
+            n_release=n_release,
+            extra_release_counts=extra_release_counts,
+            log_metrics=True,
+        )
+
+    def _apply_storage_runtime_config(
+        self,
+        *,
+        storage_backend: Optional[str],
+        prefetch_threshold: int,
+        prefetch_timeout_base: float,
+        prefetch_timeout_per_ki_token: float,
+        hicache_storage_pass_prefix_keys: bool,
+        enable_storage: bool,
+        enable_storage_metrics: bool,
+        extra_metric_labels: Optional[dict[str, str]],
+    ) -> None:
+        self.enable_storage = enable_storage
+        self.prefetch_threshold = prefetch_threshold
+        self.prefetch_timeout_base = prefetch_timeout_base
+        self.prefetch_timeout_per_page = (
+            self.page_size / 1024 * prefetch_timeout_per_ki_token
+        )
+        self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys
+        self.enable_storage_metrics = enable_storage_metrics
+
+        if self.enable_storage_metrics:
+            attn_cp_rank, attn_cp_size = (
+                self.cache_controller.get_attn_cp_rank_and_size()
+            )
+            labels = {
+                "storage_backend": storage_backend,
+                "tp_rank": self.cache_controller.tp_rank,
+                "dp_rank": self.cache_controller.dp_rank,
+                "pp_rank": self.cache_controller.pp_rank,
+                "pp_size": self.cache_controller.pp_size,
+                "attn_cp_rank": attn_cp_rank,
+                "attn_cp_size": attn_cp_size,
+            }
+            if extra_metric_labels:
+                labels.update(extra_metric_labels)
+            existing_collector = self.storage_metrics_collector
+            if existing_collector is None:
+                self.storage_metrics_collector = StorageMetricsCollector(labels=labels)
+            elif set(existing_collector.labels.keys()) == set(labels.keys()):
+                existing_collector.labels = labels
+            else:
+                logger.warning(
+                    "Storage metrics labels changed (%s -> %s). Keep existing labels to avoid duplicate metric registration.",
+                    sorted(existing_collector.labels.keys()),
+                    sorted(labels.keys()),
+                )
+        else:
+            self.storage_metrics_collector = None
+
+    def attach_storage_backend(
+        self,
+        storage_backend: str,
+        storage_backend_extra_config_json: Optional[str] = None,
+        served_model_name: Optional[str] = None,
+        hicache_storage_prefetch_policy: Optional[str] = None,
+        hicache_write_policy: Optional[str] = None,
+    ) -> tuple[bool, str]:
+        return (
+            False,
+            "UnifiedRadixCache does not support runtime HiCache storage attach yet. "
+            "Configure hicache_storage_backend at startup instead.",
+        )
+
+    def detach_storage_backend(self) -> tuple[bool, str]:
+        return (
+            False,
+            "UnifiedRadixCache does not support runtime HiCache storage detach yet. "
+            "Restart without hicache_storage_backend to disable it.",
+        )
+
+    def clear_storage_backend(self) -> bool:
+        try:
+            ok = self.cache_controller.clear_storage_backend()
+        except Exception as e:
+            logger.error("Failed to clear hierarchical cache storage backend: %s", e)
+            return False
+        if ok:
+            logger.info("Hierarchical cache storage backend cleared successfully!")
+        return ok
+
     # ---- HiCache: Async Event Management ----
 
     def writing_check(self, write_back: bool = False) -> None:
@@ -1387,6 +2033,8 @@ class UnifiedRadixCache(BasePrefixCache):
                             node, params = entry
                             if params is not None:
                                 self.dec_lock_ref(node, params)
+                            if self.enable_storage:
+                                self.write_backup_storage(node)
                 cc.ack_write_queue.clear()
                 assert len(self.ongoing_write_through) == 0
             return
@@ -1415,6 +2063,8 @@ class UnifiedRadixCache(BasePrefixCache):
             for ack_id in ack_list:
                 node, params = self.ongoing_write_through.pop(ack_id)
                 self.dec_lock_ref(node, params)
+                if self.enable_storage:
+                    self.write_backup_storage(node)
             finish_count -= 1
 
     def loading_check(self) -> None:
@@ -1484,6 +2134,12 @@ class UnifiedRadixCache(BasePrefixCache):
         """Called per scheduler step to poll async HiCache events."""
         self.writing_check()
         self.loading_check()
+        if self.enable_storage:
+            self.drain_storage_control_queues()
+        if self.enable_storage_metrics and self.storage_metrics_collector is not None:
+            self.storage_metrics_collector.log_storage_metrics(
+                self.cache_controller.storage_backend.get_stats()
+            )
 
     def flush_write_through_acks(self) -> None:
         """Flush pending write-through acknowledgements."""
@@ -1883,10 +2539,6 @@ class UnifiedRadixCache(BasePrefixCache):
             logger.error(msg)
             self.pretty_print()
             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,

test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py

@@ -24,7 +24,7 @@ MAMBA_TRACK_INTERVAL = 128
 DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
 DSV4_FLASH_LAUNCH_TIMEOUT = 3600
 
-register_cuda_ci(est_time=745, stage="base-c", runner_config="8-gpu-h200")
+register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200")
 
 
 class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):

test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py

@@ -28,8 +28,8 @@ GLM5_LAUNCH_TIMEOUT = 3600
 register_cuda_ci(est_time=900, suite="nightly-8-gpu-h200", nightly=True)
 
 
-class GSM8KTwoPassMixin:
+class AccuracyTwoPassMixin:
-    """Mixin: run GSM8K twice with flush in between, verify accuracy diff.
+    """Mixin: run an eval twice with flush in between, verify accuracy diff.
 
     Subclass must provide:
       - self.base_url
@@ -38,9 +38,14 @@ class GSM8KTwoPassMixin:
 
     gsm8k_threshold: float = 0.90
     num_gsm8k_questions: int = 200
-    max_accuracy_diff: float = 0.02
     gsm8k_parallel: int = 40
 
+    mmlu_threshold: float = 0.75
+    num_mmlu_examples: int = 200
+    mmlu_num_threads: int = 32
+
+    max_accuracy_diff: float = 0.02
+
     def _run_gsm8k(self):
         from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
 
@@ -57,6 +62,19 @@ class GSM8KTwoPassMixin:
         metrics = run_few_shot_gsm8k(args)
         return metrics["accuracy"]
 
+    def _run_mmlu(self):
+        from sglang.test.run_eval import run_eval as run_simple_eval
+
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmlu",
+            num_examples=self.num_mmlu_examples,
+            num_threads=self.mmlu_num_threads,
+        )
+        metrics = run_simple_eval(args)
+        return metrics["score"]
+
     def _flush_cache(self):
         response = requests.post(
             self.base_url + "/flush_cache",
@@ -65,45 +83,52 @@ class GSM8KTwoPassMixin:
         )
         response.raise_for_status()
 
-    def test_gsm8k_two_passes(self):
+    def _two_pass(self, name: str, run_fn, threshold: float):
-        """Run GSM8K twice with flush in between, verify accuracy diff <= max_accuracy_diff."""
         # First pass
-        acc1 = self._run_gsm8k()
+        acc1 = run_fn()
-        print(f"[{self.__class__.__name__}] GSM8K pass 1 accuracy: {acc1:.3f}")
+        print(f"[{self.__class__.__name__}] {name} pass 1 accuracy: {acc1:.3f}")
         self.assertGreaterEqual(
             acc1,
-            self.gsm8k_threshold,
+            threshold,
-            f"Pass 1 accuracy {acc1:.3f} < threshold {self.gsm8k_threshold}",
+            f"{name} pass 1 accuracy {acc1:.3f} < threshold {threshold}",
         )
 
         # Flush cache
         self._flush_cache()
 
         # Second pass
-        acc2 = self._run_gsm8k()
+        acc2 = run_fn()
-        print(f"[{self.__class__.__name__}] GSM8K pass 2 accuracy: {acc2:.3f}")
+        print(f"[{self.__class__.__name__}] {name} pass 2 accuracy: {acc2:.3f}")
         self.assertGreaterEqual(
             acc2,
-            self.gsm8k_threshold,
+            threshold,
-            f"Pass 2 accuracy {acc2:.3f} < threshold {self.gsm8k_threshold}",
+            f"{name} pass 2 accuracy {acc2:.3f} < threshold {threshold}",
         )
 
-        # Verify diff
+        # Verify diff (only fail when 2nd pass regressed)
         if acc1 > acc2:
             diff = abs(acc1 - acc2)
             print(
-                f"[{self.__class__.__name__}] Accuracy diff: {diff:.3f} "
+                f"[{self.__class__.__name__}] {name} accuracy diff: {diff:.3f} "
                 f"(max allowed: {self.max_accuracy_diff})"
             )
             self.assertLessEqual(
                 diff,
                 self.max_accuracy_diff,
-                f"Accuracy diff {diff:.3f} exceeds max {self.max_accuracy_diff} "
+                f"{name} accuracy diff {diff:.3f} exceeds max {self.max_accuracy_diff} "
                 f"(pass1={acc1:.3f}, pass2={acc2:.3f})",
             )
 
+    def test_gsm8k_two_passes(self):
+        """Run GSM8K twice with flush in between, verify accuracy diff <= max_accuracy_diff."""
+        self._two_pass("GSM8K", self._run_gsm8k, self.gsm8k_threshold)
 
-class TestGLM5HiCacheL3GSM8K(GSM8KTwoPassMixin, CustomTestCase):
+    def test_mmlu_two_passes(self):
+        """Run MMLU twice with flush in between, verify accuracy diff <= max_accuracy_diff."""
+        self._two_pass("MMLU", self._run_mmlu, self.mmlu_threshold)
+
+
+class TestGLM5HiCacheL3Accuracy(AccuracyTwoPassMixin, CustomTestCase):
     """GLM-5.1-FP8 + HiCache L3 (file backend), with HiRadixTree."""
 
     @classmethod

test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py

@@ -0,0 +1,77 @@
+import os
+import shutil
+import tempfile
+import unittest
+
+from test_unified_radix_cache_kl_hicache_nightly import AccuracyTwoPassMixin
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
+MAMBA_TRACK_INTERVAL = 128
+
+register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200")
+
+
+class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
+    """Mamba hybrid + HiCache L3 (file backend) + UnifiedRadixCache."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = MAMBA_MODEL
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.hicache_dir = tempfile.mkdtemp(prefix="hicache_l3_mamba_")
+        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",
+                "2",
+                "--hicache-write-policy",
+                "write_through",
+                "--hicache-storage-prefetch-policy",
+                "wait_complete",
+                "--hicache-io-backend",
+                "direct",
+                "--hicache-mem-layout",
+                "page_first_direct",
+                "--hicache-storage-backend",
+                "file",
+                "--max-mamba-cache-size",
+                "500",
+                "--weight-loader-prefetch-checkpoints",
+            ],
+            env={
+                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+                "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
+            },
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        if os.path.isdir(cls.hicache_dir):
+            shutil.rmtree(cls.hicache_dir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -2274,13 +2274,6 @@ class UnifiedRadixCacheSuite:
         tokens = self._swa_anchor_chain_tokens(len(chain))
         return tree, chain, n, y, x, tokens
 
-    def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self):
-        tree, _, n, y, x, tokens = self._swa_anchor_setup()
-        result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
-        self.assertIs(result.best_match_node, x)
-        self.assertIs(result.last_device_node, n.parent)
-        self.assertIs(result.last_host_node, y)
-
     def test_hicache_swa_load_back_anchored_on_best_match_node(self):
         tree, _, _, y, x, _ = self._swa_anchor_setup()
         ps = self.cfg.page_size