[UnifiedTree]: Support l3 storage for swa and deepseek v4 (#26881)

Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>

python/sglang/srt/managers/cache_controller.py

@@ -625,7 +625,15 @@ class HiCacheController:
             self.dp_rank = 0
 
         # Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool.
-        is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool)
+        # DeepSeekV4TokenToKVPool has compressed MLA-style rank-replicated cache
+        # data. storage only needs rank 0 to write it back.
+        from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+
+        is_mla_model = isinstance(self.mem_pool_device, MLATokenToKVPool)
+        is_compressed_mla_model = isinstance(
+            self.mem_pool_device, DeepSeekV4TokenToKVPool
+        )
+        is_rank_replicated = is_mla_model or is_compressed_mla_model
         # Least Common Multiple among heterogeneous tp size
         tp_lcm_size = storage_backend_extra_config.pop("tp_lcm_size", None)
         should_split_heads = False
@@ -635,7 +643,7 @@ class HiCacheController:
                 tp_lcm_size % self.tp_size == 0
             ), "tp_lcm_size must be divisible by tp_size."
             should_split_heads = (
-                not is_mla_backend
+                not is_rank_replicated
                 and self.mem_pool_host.layout == "page_head"
                 and tp_lcm_size > self.tp_size
             )
@@ -649,7 +657,8 @@ class HiCacheController:
             pp_size=self.pp_size,
             attn_cp_rank=attn_cp_rank,
             attn_cp_size=attn_cp_size,
-            is_mla_model=is_mla_backend,
+            # TODO(hzh): Rename is_mla_model to is_rank_replicated.
+            is_mla_model=is_rank_replicated,
             enable_storage_metrics=self.enable_storage_metrics,
             is_page_first_layout=self.mem_pool_host.layout == "page_first",
             model_name=model_name,

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -74,6 +74,7 @@ class InsertResult:
     """Result of an insert operation"""
 
     prefix_len: int
+    total_len: int = 0
     mamba_exist: bool = False
     inserted_host_node: Any = None
 

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

@@ -136,6 +136,7 @@ class PrefetchOperation(StorageOperation):
             prefix_keys=prefix_keys,
             pool_transfers=pool_transfers,
         )
+        self.pool_transfers_done = not bool(pool_transfers)
 
     def increment(self, num_tokens: int):
         with self._lock:
@@ -581,9 +582,6 @@ class HybridCacheController(BaseHiCacheController):
         kv_hit_pages = hit_result.kv_hit_pages
         operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages)
 
-        if kv_hit_pages > 0 and operation.pool_transfers:
-            self._sync_trailing_keys(operation.pool_transfers, hash_value, kv_hit_pages)
-
         return (
             hash_value[:kv_hit_pages],
             kv_hit_pages * self.page_size,
@@ -618,14 +616,21 @@ class HybridCacheController(BaseHiCacheController):
         return host_indices, device_indices, resolved_pool_transfers
 
     def _page_transfer(self, operation):
-        # Transfer extra pools
-        if operation.pool_transfers and not operation.is_terminated():
+        # KV pools first — determines actual completed page count
+        super()._page_transfer(operation)
+
+        # Extra pools only after KV fully completes. If KV terminated early
+        # (IO failure, timeout, TP mismatch), skip extra IO entirely to avoid
+        # data misalignment.
+        kv_completed_pages = operation.completed_tokens // self.page_size
+        if operation.pool_transfers and kv_completed_pages == len(operation.hash_value):
+            self._sync_trailing_keys(
+                operation.pool_transfers, operation.hash_value, kv_completed_pages
+            )
             self._resolve_sidecar_derived_pool_transfers(operation)
             results = self.storage_backend.batch_get_v2(operation.pool_transfers)
             operation.pool_storage_result.update_extra_pool_hit_pages(results)
-
-        # Transfer kv pools
-        super()._page_transfer(operation)
+        operation.pool_transfers_done = True
 
     def _page_backup(self, operation):
         # Backup extra pools
@@ -642,14 +647,27 @@ class HybridCacheController(BaseHiCacheController):
             if transfer.indices_from_pool is None:
                 continue
             if transfer.indices_from_pool != PoolName.KV:
-                # TODO(hzh): Support storage sidecar derived pools from other sources
-                raise AssertionError(
-                    "Storage sidecar derived pool currently only supports KV-shared "
-                    f"indices, got {transfer.name} from {transfer.indices_from_pool}."
+                source = next(
+                    (
+                        t
+                        for t in operation.pool_transfers
+                        if t.indices_from_pool is None
+                        and t.name == transfer.indices_from_pool
+                    ),
+                    None,
                 )
-            transfer.host_indices = operation.host_indices
-            if transfer.keys is None:
-                transfer.keys = operation.hash_value
+                if source is None:
+                    raise AssertionError(
+                        "Storage sidecar derived pool source missing: "
+                        f"{transfer.name} from {transfer.indices_from_pool}."
+                    )
+                transfer.host_indices = source.host_indices
+                if transfer.keys is None:
+                    transfer.keys = source.keys
+            else:
+                transfer.host_indices = operation.host_indices
+                if transfer.keys is None:
+                    transfer.keys = operation.hash_value
 
     def _sync_trailing_keys(
         self,

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

@@ -4,7 +4,11 @@ import logging
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
-from sglang.srt.mem_cache.hicache_storage import PoolName, SidecarPoolSpec
+from sglang.srt.mem_cache.hicache_storage import (
+    PoolHitPolicy,
+    PoolName,
+    SidecarPoolSpec,
+)
 from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
     HybridCacheController,
 )
@@ -737,7 +741,15 @@ class _DeepSeekV4Strategy(StackStrategy):
             enable_storage_metrics=enable_storage_metrics,
         )
         sidecars = [
-            SidecarPoolSpec(pool_name=name, indices_from_pool=src)
+            SidecarPoolSpec(
+                pool_name=name,
+                indices_from_pool=src,
+                hit_policy=(
+                    PoolHitPolicy.TRAILING_PAGES
+                    if src == PoolName.SWA
+                    else PoolHitPolicy.ALL_PAGES
+                ),
+            )
             for name, src in (
                 (PoolName.DEEPSEEK_V4_C4, PoolName.KV),
                 (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV),

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -2078,6 +2078,7 @@ class DeepSeekV4PagedHostPool(HostKVCache):
             return
         host_rows = self._to_page_indices(host_indices)
         device_rows = self._to_page_indices(device_indices)
+
         if io_backend == "kernel" and self.layout == "layer_first":
             transfer_kv_per_layer_mla(
                 src=self.data_refs[layer_id],

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -601,6 +601,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
 
     def register_mem_pool_host(self, mem_pool_host: HostKVCache):
         super().register_mem_pool_host(mem_pool_host)
+        if getattr(self.mem_pool_host, "kv_buffer", None) is None:
+            # Hybrid logical anchors only own allocation indices. Their physical
+            # tensors are registered through register_mem_host_pool_v2().
+            return
         assert self.mem_pool_host.layout in [
             "page_first",
             "page_first_direct",
@@ -631,37 +635,41 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
         # the corresponding host pool implementation at runtime.
         self.registered_pools[host_pool_name] = host_pool
 
-        # Hybrid pools expose the tensors that Mooncake needs for zero-copy I/O.
-        # The storage backend only depends on this accessor, not concrete fields.
-        buf_list = host_pool.get_hybrid_pool_buffer()
-        for buf in buf_list:
+        # Non-anchor pools are either sidecar-specific pools with their own
+        # accessor, or ordinary KV-like host pools used as SWA side pools.
+        get_buffers = getattr(
+            host_pool,
+            "get_hybrid_pool_buffer",
+            lambda: [getattr(host_pool, "kv_buffer", None)],
+        )
+        for buf in get_buffers():
+            if buf is None:
+                continue
             super().register_buffer(buf)
 
     def _tag_keys(self, keys: List[str]) -> List[str]:
         if self.extra_backend_tag is None:
             return keys
-        return [f"{ self.extra_backend_tag}_{key}" for key in keys]
+        return [f"{self.extra_backend_tag}_{key}" for key in keys]
 
     def _get_hybrid_page_component_keys(
         self, page_keys: List[str], transfer: PoolTransfer
     ) -> Tuple[List[str], int]:
-        # A logical "page" may map to multiple physical objects in storage.
-        # - INDEXER: one key per page
-        # - MAMBA  : one temporal key + N conv keys per page
-        # key_multiplier records how many component keys are generated per page.
-        name = transfer.name
+        host_pool = getattr(self, "registered_pools", {}).get(transfer.name)
+        if host_pool is None:
+            raise ValueError(f"Unregistered Mooncake hybrid pool: {transfer.name}")
+
+        # Suffix order must match get_page_buffer_meta() for one page, because
+        # Mooncake zips object keys with registered buffer pointers.
+        pool_name = transfer.name
         suffixes = []
-        if name == PoolName.INDEXER:
-            suffixes = [f"_{self.mla_suffix}_{PoolName.INDEXER}"]
-        elif name == PoolName.MAMBA:
-            pools = getattr(self, "registered_pools", {})
-            mamba_pool = pools.get(PoolName.MAMBA)
-            conv_num = len(getattr(mamba_pool, "conv_buffer", None) or [])
-            base_suffix = f"_{self.mha_suffix}"
-            suffixes = [f"{base_suffix}_temporal"] + [
-                f"{base_suffix}_conv_{i}" for i in range(conv_num)
+        if pool_name == PoolName.MAMBA:
+            # Mamba stores one temporal object plus one object per conv state.
+            conv_num = len(getattr(host_pool, "conv_buffer", None) or [])
+            suffixes = [f"_{self.mha_suffix}_temporal"] + [
+                f"_{self.mha_suffix}_conv_{i}" for i in range(conv_num)
             ]
-        elif name == PoolName.DRAFT:
+        elif pool_name == PoolName.DRAFT:
             # Draft pool's MLA/MHA layout is independent from the target
             # (e.g. EAGLE-MHA draft on top of an MLA target), so pick the
             # suffix scheme from the draft pool's own class. The `_draft`
@@ -675,6 +683,33 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
                     f"_{self.mha_suffix}_{PoolName.DRAFT}_k",
                     f"_{self.mha_suffix}_{PoolName.DRAFT}_v",
                 ]
+        elif pool_name in (
+            PoolName.INDEXER,
+            PoolName.DEEPSEEK_V4_C4,
+            PoolName.DEEPSEEK_V4_C4_INDEXER,
+            PoolName.DEEPSEEK_V4_C128,
+            PoolName.DEEPSEEK_V4_C4_STATE,
+            PoolName.DEEPSEEK_V4_C4_INDEXER_STATE,
+            PoolName.DEEPSEEK_V4_C128_STATE,
+        ):
+            # DSA indexer and DeepSeek V4 side pools are page-packed
+            # single-object pools.
+            suffixes = [f"_{self.mla_suffix}_{pool_name}"]
+        elif pool_name == PoolName.SWA:
+            if not self.is_mla_backend and hasattr(host_pool, "v_buffer"):
+                # Ordinary MHA SWA mirrors a K/V pool.
+                suffixes = [
+                    f"_{self.mha_suffix}_{pool_name}_k",
+                    f"_{self.mha_suffix}_{pool_name}_v",
+                ]
+            elif self.is_mla_backend:
+                suffixes = [f"_{self.mla_suffix}_{pool_name}"]
+
+        if not suffixes:
+            raise ValueError(
+                f"Unsupported Mooncake hybrid pool name: {pool_name}, "
+                f"host_pool={type(host_pool)}"
+            )
         key_multiplier = len(suffixes)
         component_keys = [
             f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes
@@ -687,7 +722,12 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
         pool_transfers: Optional[List[PoolTransfer]] = None,
         extra_info: Optional[HiCacheStorageExtraInfo] = None,
     ) -> PoolTransferResult:
-        kv_pages = self.batch_exists(keys, extra_info)
+        if self.mem_pool_host.kv_buffer is None:
+            # Logical anchor: no physical KV object exists in Mooncake, so the
+            # usable prefix is determined entirely by required sidecar objects.
+            kv_pages = len(keys)
+        else:
+            kv_pages = self.batch_exists(keys, extra_info)
 
         hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {}
         final_pages = kv_pages
@@ -863,6 +903,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
         host_indices: torch.Tensor,
         extra_info: Optional[HiCacheStorageExtraInfo] = None,
     ) -> List[bool]:
+        if self.mem_pool_host.kv_buffer is None:
+            # DeepSeek V4's KV anchor is logical only; v2 side pools carry data.
+            return [True] * len(keys)
+
         # Apply extra_backend_tag prefix if available
         keys = self._tag_keys(keys)
 
@@ -888,6 +932,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
         host_indices: torch.Tensor,
         extra_info: Optional[HiCacheStorageExtraInfo] = None,
     ) -> List[bool]:
+        if self.mem_pool_host.kv_buffer is None:
+            # DeepSeek V4's KV anchor is logical only; v2 side pools carry data.
+            return [True] * len(keys)
+
         # Apply extra_backend_tag prefix if available
         keys = self._tag_keys(keys)
 

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

@@ -14,6 +14,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchResult,
 )
 from sglang.srt.mem_cache.hicache_storage import (
+    PoolHitPolicy,
     PoolName,
     PoolTransfer,
     PoolTransferResult,
@@ -415,7 +416,9 @@ class SWAComponent(TreeComponent):
         root = self.cache.root_node
         sliding_window_size = self.sliding_window_size
         swa_lock_size = 0
-        swa_uuid_for_lock = None
+        swa_uuid = None
+        uuid_key = "host_uuid" if lock_host else "uuid"
+        lru = self.cache.host_lru_lists[ct] if lock_host else self.cache.lru_lists[ct]
 
         # Tombstoned nodes (cd.value is None) have no SWA chunk to protect
         # skip them and keep walking up. This path is hit when HiCache
@@ -423,23 +426,36 @@ class SWAComponent(TreeComponent):
         cur = node
         while cur != root and swa_lock_size < sliding_window_size:
             comp = cur.component_data[ct]
-            if comp.value is None:
+            value = comp.host_value if lock_host else comp.value
+            if value is None:
                 result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
                 cur = cur.parent
                 continue
-            if comp.lock_ref == 0:
-                key_len = len(cur.key)
-                self.cache.component_evictable_size_[ct] -= key_len
-                self.cache.component_protected_size_[ct] += key_len
-            comp.lock_ref += 1
-            swa_lock_size += len(cur.key)
+
+            ref = comp.host_lock_ref if lock_host else comp.lock_ref
+            if ref == 0:
+                if lock_host:
+                    if lru.in_list(cur):
+                        lru.remove_node(cur)
+                else:
+                    key_len = len(cur.key)
+                    self.cache.component_evictable_size_[ct] -= key_len
+                    self.cache.component_protected_size_[ct] += key_len
+            if lock_host:
+                comp.host_lock_ref = ref + 1
+            else:
+                comp.lock_ref = ref + 1
+            swa_lock_size += len(value)
             if swa_lock_size >= sliding_window_size:
-                if comp.metadata.get("uuid") is None:
-                    comp.metadata["uuid"] = next_component_uuid()
-                swa_uuid_for_lock = comp.metadata["uuid"]
+                if comp.metadata.get(uuid_key) is None:
+                    comp.metadata[uuid_key] = next_component_uuid()
+                swa_uuid = comp.metadata[uuid_key]
             cur = cur.parent
 
-        result.swa_uuid_for_lock = swa_uuid_for_lock
+        if lock_host:
+            result.swa_uuid_for_host_lock = swa_uuid
+        else:
+            result.swa_uuid_for_lock = swa_uuid
         return result
 
     def release_component_lock(
@@ -450,9 +466,14 @@ class SWAComponent(TreeComponent):
     ) -> None:
         ct = self.component_type
         root = self.cache.root_node
-        swa_uuid_for_lock = params.swa_uuid_for_lock if params else None
+        swa_uuid_for_lock = (
+            (params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock)
+            if params
+            else None
+        )
         skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
         dec_swa = True
+        uuid_key = "host_uuid" if lock_host else "uuid"
 
         # A node in skip_lock_node_ids was a tombstone when this lock was acquired.
         cur = node
@@ -461,15 +482,25 @@ class SWAComponent(TreeComponent):
             if cur.id in skip_lock_node_ids:
                 cur = cur.parent
                 continue
-            if comp.lock_ref == 0:
+            ref = comp.host_lock_ref if lock_host else comp.lock_ref
+            if ref == 0:
                 cur = cur.parent
                 continue
-            if comp.lock_ref == 1:
-                key_len = len(cur.key)
-                self.cache.component_evictable_size_[ct] += key_len
-                self.cache.component_protected_size_[ct] -= key_len
-            comp.lock_ref -= 1
-            if swa_uuid_for_lock and comp.metadata.get("uuid") == swa_uuid_for_lock:
+            if ref == 1:
+                if lock_host:
+                    if comp.value is None and comp.host_value is not None:
+                        host_lru = self.cache.host_lru_lists[ct]
+                        if not host_lru.in_list(cur):
+                            host_lru.insert_mru(cur)
+                else:
+                    key_len = len(comp.value)
+                    self.cache.component_evictable_size_[ct] += key_len
+                    self.cache.component_protected_size_[ct] -= key_len
+            if lock_host:
+                comp.host_lock_ref = ref - 1
+            else:
+                comp.lock_ref = ref - 1
+            if swa_uuid_for_lock and comp.metadata.get(uuid_key) == swa_uuid_for_lock:
                 dec_swa = False
             cur = cur.parent
 
@@ -546,6 +577,46 @@ class SWAComponent(TreeComponent):
                 )
             ]
 
+        if phase == CacheTransferPhase.BACKUP_STORAGE:
+            cd = node.component_data[ct]
+            if cd.host_value is None or not node.hash_value:
+                return None
+            num_pages = len(cd.host_value) // self.cache.page_size
+            if num_pages == 0:
+                return None
+            return [
+                PoolTransfer(
+                    name=PoolName.SWA,
+                    host_indices=cd.host_value[-num_pages * self.cache.page_size :],
+                    keys=node.hash_value[-num_pages:],
+                    hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                )
+            ]
+
+        if phase == CacheTransferPhase.PREFETCH:
+            num_pages = min(
+                prefetch_tokens // self.cache.page_size,
+                (self.sliding_window_size + self.cache.page_size - 1)
+                // self.cache.page_size,
+            )
+            if num_pages == 0:
+                return None
+            num_tokens = num_pages * self.cache.page_size
+            host_indices = self._swa_kv_pool_host.alloc(num_tokens)
+            if host_indices is None:
+                self.cache.evict_host(num_tokens, ComponentType.SWA)
+                host_indices = self._swa_kv_pool_host.alloc(num_tokens)
+            if host_indices is None:
+                return []
+            return [
+                PoolTransfer(
+                    name=PoolName.SWA,
+                    host_indices=host_indices,
+                    keys=["__placeholder__"] * num_pages,
+                    hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                )
+            ]
+
         return None
 
     def commit_hicache_transfer(
@@ -586,6 +657,98 @@ class SWAComponent(TreeComponent):
             assert offset == len(xfer.host_indices)
             return
 
+        if phase == CacheTransferPhase.PREFETCH:
+            self._commit_prefetch(
+                node,
+                transfers,
+                insert_result=insert_result,
+                pool_storage_result=pool_storage_result,
+            )
+            return
+
+    def _release_swa_host(self, host_indices: torch.Tensor) -> None:
+        if host_indices is not None and host_indices.numel() > 0:
+            self.cache.cache_controller.append_host_mem_release(
+                extra_pools=[PoolTransfer(name=PoolName.SWA, host_indices=host_indices)]
+            )
+
+    def _attach_swa_host_value(
+        self, node: "UnifiedTreeNode", host_indices: torch.Tensor
+    ) -> None:
+        """Write host_indices into node's SWA host_value and refresh tree state."""
+        ct = self.component_type
+        cd = node.component_data[ct]
+        cd.host_value = host_indices.clone()
+        host_lru = self.cache.host_lru_lists[ct]
+        if cd.value is None and not host_lru.in_list(node):
+            host_lru.insert_mru(node)
+        self.cache._update_evictable_leaf_sets(node)
+        if node.parent:
+            self.cache._update_evictable_leaf_sets(node.parent)
+
+    def _commit_prefetch(
+        self,
+        anchor,
+        transfers: list[PoolTransfer],
+        *,
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        """Distribute the prefetched SWA buffer onto the leaf→anchor path.
+
+        The buffer holds the trailing ``loaded_pages`` of the completed KV
+        prefix, mapped to token range ``[loaded_start, total_len)``. We walk
+        upward from ``inserted_host_node`` to ``anchor`` and, for each node
+        whose token range overlaps the buffer:
+          - SWA tombstone (host_value is None) → fill from buffer (split if
+            the node only partially overlaps at the buffer's left edge)
+          - already has SWA host_value → release the corresponding slice
+        Any leftover buffer beyond the walked range is also released.
+        """
+        if not transfers:
+            return
+        ct = self.component_type
+        host_indices = transfers[0].host_indices
+        loaded_pages = (
+            pool_storage_result.extra_pool_hit_pages.get(PoolName.SWA, 0)
+            if pool_storage_result
+            else 0
+        )
+        target = insert_result.inserted_host_node if insert_result else None
+        if not loaded_pages or target is None:
+            self._release_swa_host(host_indices)
+            return
+
+        # Buffer covers token range [loaded_start, total_len).
+        loaded_start = insert_result.total_len - loaded_pages * self.cache.page_size
+
+        # Walk leaf → anchor; ``pos`` is the right edge of ``cur`` in tokens.
+        pos, cur = insert_result.total_len, target
+        while cur is not anchor and pos > loaded_start:
+            node_start = pos - len(cur.key)
+            # Intersection of cur's range and the buffer.
+            fill_start = max(node_start, loaded_start)
+            fill_len = pos - fill_start
+            buf_off = fill_start - loaded_start
+            slice_ = host_indices[buf_off : buf_off + fill_len]
+
+            cd = cur.component_data[ct]
+            if cd.host_value is None and fill_len > 0:
+                # Tombstone: split off the in-buffer tail if needed, then fill.
+                if fill_start > node_start:
+                    self.cache._split_node(cur.key, cur, fill_start - node_start)
+                self._attach_swa_host_value(cur, slice_)
+            else:
+                # Already has SWA (or empty overlap): drop this slice.
+                self._release_swa_host(slice_)
+
+            pos = node_start
+            cur = cur.parent
+
+        # Buffer prefix that fell outside the anchor→leaf path.
+        if pos > loaded_start:
+            self._release_swa_host(host_indices[: pos - loaded_start])
+
     def drive_host_eviction(
         self, num_tokens: int, tracker: dict[ComponentType, int]
     ) -> None:

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1121,9 +1121,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             if len(key):
                 child_key = key.child_key(self.page_size)
 
-        result = InsertResult(
-            prefix_len=matched_length,
-        )
+        result = InsertResult(prefix_len=matched_length, total_len=total_len)
         if len(key) == 0:
             if (
                 node is not self.root_node
@@ -1662,6 +1660,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             transfers.append(
                 PoolTransfer(
                     name=spec.pool_name,
+                    keys=indices_source.keys,
                     hit_policy=spec.hit_policy,
                     indices_from_pool=spec.indices_from_pool,
                 )
@@ -1850,6 +1849,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             )
         else:
             return True
+        if (
+            completed
+            and getattr(operation, "pool_transfers", None)
+            and not getattr(operation, "pool_transfers_done", True)
+        ):
+            can_terminate = False
 
         operation_terminated = operation.is_terminated()
         states = torch.tensor(

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -1,5 +1,9 @@
 """Unit tests for UnifiedRadixCache"""
 
+import json
+import shutil
+import tempfile
+import time
 import unittest
 from array import array
 from dataclasses import dataclass, replace
@@ -1840,6 +1844,146 @@ class UnifiedRadixCacheSuite:
     # HiCache Unit Tests (real cache_controller D<->H backup/load)
     # ================================================================
 
+    # ---------- L3 storage (file backend) helpers ----------
+
+    def _path_chain(self, tree, node):
+        """Return root->node node chain (excluding root)."""
+        chain = []
+        cur = node
+        while cur is not tree.root_node:
+            chain.append(cur)
+            cur = cur.parent
+        chain.reverse()
+        return chain
+
+    def _write_path_to_l3(self, tree, node):
+        """Offload every node on root->node path from host to L3 storage."""
+        for n in self._path_chain(tree, node):
+            tree.write_backup_storage(n)
+
+    def _flush_l3_backups(self, tree, timeout: float = 10.0):
+        """Wait for backup threads to finish, then drain acks (release locks)."""
+        deadline = time.time() + timeout
+        while tree.ongoing_backup and time.time() < deadline:
+            tree.drain_storage_control_queues()
+            if tree.ongoing_backup:
+                time.sleep(0.01)
+        tree.drain_storage_control_queues()
+        self.assertFalse(tree.ongoing_backup, "L3 backups did not complete in time")
+
+    def _run_prefetch_to_completion(self, tree, req_id, timeout: float = 10.0):
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            if tree.check_prefetch_progress(req_id):
+                return
+            time.sleep(0.01)
+        self.fail(f"prefetch {req_id} did not complete in time")
+
+    def _all_page_hashes(self, tree, node):
+        hashes = []
+        for n in self._path_chain(tree, node):
+            hashes.extend(list(n.hash_value))
+        return hashes
+
+    def test_hicache_l3_write_storage(self):
+        """D->H->L3 offload: every KV page lands in the file storage backend."""
+        if self._skip_unsupported_hicache_test():
+            return
+        if self.cfg.has_mamba:
+            self.skipTest("mamba L3 offload is out of scope for this unit fixture")
+
+        storage_dir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
+
+        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+        self._init_hicache(
+            tree,
+            storage_backend="file",
+            storage_dir=storage_dir,
+            prefetch_threshold=1,
+        )
+
+        seq = self._make_seq(1, 4)
+        self._insert(tree, allocator, req_to_token_pool, seq)
+        m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
+        leaf = m.last_device_node
+
+        # D->H first, then H->L3.
+        self._backup_node(tree, leaf)
+        self.assertTrue(leaf.hash_value)
+        self._write_path_to_l3(tree, leaf)
+        self._flush_l3_backups(tree)
+
+        # Every KV page hash on the path must now exist in storage.
+        backend = tree.cache_controller.storage_backend
+        page_hashes = self._all_page_hashes(tree, leaf)
+        self.assertEqual(len(page_hashes), len(seq) // self.cfg.page_size)
+        self.assertEqual(backend.batch_exists(page_hashes), len(page_hashes))
+        tree.sanity_check()
+
+    def test_hicache_l3_prefetch(self):
+        """L3 round trip: write with one tree, prefetch into a fresh tree.
+
+        Uses two independent trees that share the same file storage dir so the
+        prefetch path genuinely reloads from L3 (no host/device residue).
+        """
+        if self._skip_unsupported_hicache_test():
+            return
+        if self.cfg.has_mamba:
+            self.skipTest("mamba L3 prefetch is out of scope for this unit fixture")
+
+        storage_dir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
+        seq = self._make_seq(1, 4)
+
+        # --- Producer tree: fill KV, backup D->H, offload H->L3. ---
+        prod, prod_alloc, prod_rtp = build_fixture(self.cfg)
+        self._init_hicache(
+            prod,
+            storage_backend="file",
+            storage_dir=storage_dir,
+            prefetch_threshold=1,
+        )
+        self._insert(prod, prod_alloc, prod_rtp, seq)
+        mp = prod.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
+        prod_leaf = mp.last_device_node
+        self._fill_full_kv(prod_alloc, mp.device_indices, marker=7)
+        expected_k, expected_v = self._snapshot_full_kv(prod_alloc, mp.device_indices)
+        self._backup_node(prod, prod_leaf)
+        self._write_path_to_l3(prod, prod_leaf)
+        self._flush_l3_backups(prod)
+
+        # --- Consumer tree: prefetch the same tokens straight from L3. ---
+        cons, cons_alloc, cons_rtp = build_fixture(self.cfg)
+        self._init_hicache(
+            cons,
+            storage_backend="file",
+            storage_dir=storage_dir,
+            prefetch_threshold=1,
+        )
+        req_id = "l3-prefetch-req"
+        cons.prefetch_from_storage(req_id, cons.root_node, array("q", seq), None, None)
+        self._run_prefetch_to_completion(cons, req_id)
+        cons.drain_storage_control_queues()
+
+        # The full prefix must now be a host hit (loaded from L3).
+        mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
+        self.assertEqual(mc.host_hit_length, len(seq))
+        host_node = mc.last_host_node
+        self.assertIsNot(host_node, cons.root_node)
+        self.assertTrue(host_node.evicted)
+
+        # Load the reloaded host prefix back to device and verify KV bytes.
+        self._load_back_node(cons, host_node)
+        loaded_indices = cons.match_prefix(
+            MatchPrefixParams(key=RadixKey(array("q", seq)))
+        ).device_indices
+        self.assertEqual(len(loaded_indices), len(seq))
+        loaded_k, loaded_v = self._snapshot_full_kv(cons_alloc, loaded_indices)
+        self.assertTrue(torch.equal(loaded_k, expected_k))
+        self.assertTrue(torch.equal(loaded_v, expected_v))
+        cons.sanity_check()
+
     def _skip_unsupported_hicache_test(self):
         if self.cfg.has_swa and self.cfg.has_mamba:
             self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks")
@@ -1869,7 +2013,16 @@ class UnifiedRadixCacheSuite:
                 self._simulate_backup(tree, node)
             stack.extend(node.children.values())
 
-    def _init_hicache(self, tree, *, write_policy: str = "write_through"):
+    def _init_hicache(
+        self,
+        tree,
+        *,
+        write_policy: str = "write_through",
+        storage_backend: Optional[str] = None,
+        storage_dir: Optional[str] = None,
+        prefetch_threshold: Optional[int] = None,
+        prefetch_policy: str = "wait_complete",
+    ):
         import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
 
         orig_kv_host_pool = assembler.MHATokenToKVPoolHost
@@ -1899,11 +2052,42 @@ class UnifiedRadixCacheSuite:
             patcher.start()
             self.addCleanup(patcher.stop)
 
+        storage_extra_config = None
+        if storage_backend == "file":
+            import sglang.srt.managers.cache_controller as cache_controller
+
+            # The file-backend storage config records TP rank/size.  These unit
+            # fixtures run without initializing distributed parallel state, so
+            # provide the local single-rank values that the fixture represents.
+            tp_rank_patcher = mock.patch.object(
+                cache_controller, "get_tensor_model_parallel_rank", return_value=0
+            )
+            tp_size_patcher = mock.patch.object(
+                cache_controller, "get_tensor_model_parallel_world_size", return_value=1
+            )
+            tp_rank_patcher.start()
+            tp_size_patcher.start()
+            self.addCleanup(tp_rank_patcher.stop)
+            self.addCleanup(tp_size_patcher.stop)
+
+            assert storage_dir is not None, "file backend needs a storage_dir"
+            # HiCacheFile reads the directory from this env var.
+            cm = envs.SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR.override(storage_dir)
+            cm.__enter__()
+            self.addCleanup(cm.__exit__, None, None, None)
+            extra = {}
+            if prefetch_threshold is not None:
+                extra["prefetch_threshold"] = prefetch_threshold
+            storage_extra_config = json.dumps(extra) if extra else None
+
         server_args = ServerArgs(
             model_path="dummy",
             page_size=self.cfg.page_size,
             hicache_io_backend="direct",
             hicache_write_policy=write_policy,
+            hicache_storage_backend=storage_backend,
+            hicache_storage_backend_extra_config=storage_extra_config,
+            hicache_storage_prefetch_policy=prefetch_policy,
         )
         # See build_fixture for why _mamba_cache_chunk_size is preset.
         server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size)
@@ -1911,6 +2095,17 @@ class UnifiedRadixCacheSuite:
         tree.init_hicache(server_args, tree.cache_init_params)
         tree.write_through_threshold = 1 << 30
         tree.load_back_threshold = 0
+        if storage_backend is not None:
+            # Unit fixtures size host/device pools equally, which makes the
+            # production prefetch capacity limit (host - device) zero.  Keep the
+            # L3 tests focused on storage round trips by allowing one fixture
+            # worth of prefetch tokens.
+            tree.cache_controller.prefetch_capacity_limit = max(
+                tree.cache_controller.prefetch_capacity_limit,
+                tree.cache_controller.mem_pool_host.size,
+            )
+            # Background prefetch/backup threads are daemon; stop them per-test.
+            self.addCleanup(tree.cache_controller._stop_storage_threads)
 
     def _build_hicache_fixture(self):
         fixture = build_fixture(self.cfg)