[UnifiedRadixCache][mamba] Fix mamba state corruption and slot leak when load_back aborts (#30986)

Co-authored-by: hzh0425 <hzh0425@apache.org>

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

@@ -9,6 +9,7 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import (
     ComponentType,
     EvictLayer,
     LRURefreshPhase,
+    PrepareLoadBackResult,
     TreeComponent,
     get_and_increase_time_counter,
     next_component_uuid,
@@ -23,6 +24,7 @@ __all__ = [
     "CacheTransferPhase",
     "LRURefreshPhase",
     "MambaComponent",
+    "PrepareLoadBackResult",
     "SWAComponent",
     "TreeComponent",
     "_NUM_COMPONENT_TYPES",

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

@@ -24,6 +24,7 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import (
     ComponentType,
     EvictLayer,
     LRURefreshPhase,
+    PrepareLoadBackResult,
     TreeComponent,
     get_and_increase_time_counter,
 )
@@ -511,6 +512,37 @@ class MambaComponent(TreeComponent):
 
     # ---- HiCache Hooks ----
 
+    def prepare_load_back(
+        self,
+        node: UnifiedTreeNode,
+        *,
+        req: Optional[Req] = None,
+    ) -> PrepareLoadBackResult:
+        cd = node.component_data[self.component_type]
+        # skip unless the node needs a load-back (device value absent), like build_hicache_transfers
+        if (
+            req is None
+            or req.mamba_pool_idx is not None
+            or cd.host_value is None
+            or cd.value is not None
+        ):
+            return PrepareLoadBackResult()
+        dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+        if dst is None:
+            self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
+            dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+            assert dst is not None, "Cannot alloc mamba for load_back"
+        req.mamba_pool_idx = dst[0]
+        return PrepareLoadBackResult(allocated_mamba_slot=dst)
+
+    def finalize_load_back(
+        self, req: Optional[Req], prep: PrepareLoadBackResult, success: bool
+    ) -> None:
+        # A called-off load-back returns the slot prepare allocated and clears req (the H->D copy never ran).
+        if not success and prep.allocated_mamba_slot is not None:
+            self.cache.req_to_token_pool.mamba_allocator.free(prep.allocated_mamba_slot)
+            req.mamba_pool_idx = None
+
     def build_hicache_transfers(
         self,
         node: UnifiedTreeNode,
@@ -551,16 +583,10 @@ class MambaComponent(TreeComponent):
                     )
                 )
 
-            # Per-request mamba CoW (H→D copy into request's device slot)
+            # Per-request mamba CoW: H→D copy into the request's device slot allocated by prepare_load_back.
             cd = node.component_data[ct]
             if req is not None and cd.host_value is not None:
-                if req.mamba_pool_idx is None:
-                    dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
-                    if dst is None:
-                        self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
-                        dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
-                        assert dst is not None, "Cannot alloc mamba for load_back"
-                    req.mamba_pool_idx = dst[0]
+                assert req.mamba_pool_idx is not None
                 transfers.append(
                     PoolTransfer(
                         name=PoolName.MAMBA,

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

@@ -75,6 +75,14 @@ class EvictLayer(IntFlag):
     ALL = DEVICE | HOST
 
 
+@dataclasses.dataclass(frozen=True)
+class PrepareLoadBackResult:
+    """Outcome of prepare_load_back; default = nothing to prepare."""
+
+    # Freshly allocated device mamba slot, recovered on failure.
+    allocated_mamba_slot: Optional[torch.Tensor] = None
+
+
 class CacheTransferPhase(str, Enum):
 
     BACKUP_HOST = "backup_host"  # D→H
@@ -380,6 +388,22 @@ class TreeComponent(ABC):
 
     # ---- HiCache Hooks ----
 
+    def prepare_load_back(
+        self,
+        node: UnifiedTreeNode,
+        *,
+        req: Optional[Req] = None,
+    ) -> PrepareLoadBackResult:
+        """Cache-level pre-allocation before a load-back builds its transfers."""
+        return PrepareLoadBackResult()
+
+    def finalize_load_back(
+        self, req: Optional[Req], prep: PrepareLoadBackResult, success: bool
+    ) -> None:
+        """Release state populated by prepare_load_back when the load-back did
+        not go through."""
+        pass
+
     def build_hicache_transfers(
         self,
         node: UnifiedTreeNode,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -49,6 +49,7 @@ from sglang.srt.mem_cache.unified_cache_components import (
     FullComponent,
     LRURefreshPhase,
     MambaComponent,
+    PrepareLoadBackResult,
     SWAComponent,
     TreeComponent,
     get_and_increase_time_counter,
@@ -1683,8 +1684,41 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         # Lock path & pre-evict if device pool is insufficient
         result = self.inc_lock_ref(best_match_node)
         ancestor_lock_params = result.to_dec_params()
-        kv_tokens = len(kv_xfer.host_indices)
 
+        # Let each component pre-allocate per-request state for the load-back;
+        # the finally below lets components recover it unless the load succeeds.
+        preps: dict[ComponentType, PrepareLoadBackResult] = {
+            comp.component_type: comp.prepare_load_back(best_match_node, req=req)
+            for comp in self._components_tuple
+        }
+        success = False
+        try:
+            success = self._load_back_transfers(
+                best_match_node=best_match_node,
+                mem_quota=mem_quota,
+                req=req,
+                kv_xfer=kv_xfer,
+                result=result,
+                ancestor_lock_params=ancestor_lock_params,
+                host_anchor_params=host_anchor_params,
+            )
+            return success
+        finally:
+            for comp in self._components_tuple:
+                comp.finalize_load_back(req, preps[comp.component_type], success)
+
+    def _load_back_transfers(
+        self,
+        *,
+        best_match_node: UnifiedTreeNode,
+        mem_quota: Optional[int],
+        req,
+        kv_xfer: PoolTransfer,
+        result: IncLockRefResult,
+        ancestor_lock_params: Optional[DecLockRefParams],
+        host_anchor_params: Optional[DecLockRefParams],
+    ) -> bool:
+        kv_tokens = len(kv_xfer.host_indices)
         # Build aux transfers, keyed per component.
         comp_xfers: dict[ComponentType, list] = {}
         for comp in self._components_tuple: