[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:
+                assert req.mamba_pool_idx is not 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]
                 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:

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -3317,6 +3317,256 @@ class UnifiedRadixCacheSuite:
         self._finish_pending_loads(cache)
         self._release_ongoing_load_back_locks(cache)
 
+    def test_load_back_abort_frees_unpublished_mamba_slot(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+
+        self._backup_node(cache, leaf)
+        cache.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+        self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].host_value)
+
+        # A request whose mamba slot was released: load_back's CoW arm allocates one.
+        req = self._make_req(req_to_token_pool)
+        req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0))
+        req.mamba_pool_idx = None
+        mamba_avail = req_to_token_pool.mamba_allocator.available_size()
+
+        # Impossible quota -> load_back aborts after building the transfers.
+        loaded = cache.load_back(leaf, mem_quota=-(10**9), req=req)
+
+        self.assertFalse(loaded)
+        # the aborted call must return its slot and not leave req pointing at it
+        self.assertIsNone(req.mamba_pool_idx)
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), mamba_avail
+        )
+        self._release_ongoing_load_back_locks(cache)
+
+    def test_load_back_load_failure_frees_unpublished_mamba_slot(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+
+        self._backup_node(cache, leaf)
+        cache.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        req = self._make_req(req_to_token_pool)
+        req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0))
+        req.mamba_pool_idx = None
+        mamba_avail = req_to_token_pool.mamba_allocator.available_size()
+
+        # cache_controller.load() failing (device alloc / transfer resolution)
+        # must also return the slot this call allocated.
+        with mock.patch.object(cache.cache_controller, "load", return_value=None):
+            loaded = cache.load_back(leaf, req=req)
+
+        self.assertFalse(loaded)
+        self.assertIsNone(req.mamba_pool_idx)
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), mamba_avail
+        )
+        self._release_ongoing_load_back_locks(cache)
+
+    def test_load_back_abort_keeps_preexisting_mamba_slot(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+
+        self._backup_node(cache, leaf)
+        cache.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        # The request already owns its slot: an aborted load-back must not free it.
+        req = self._make_req(req_to_token_pool)
+        preexisting_slot = req.mamba_pool_idx
+        self.assertIsNotNone(preexisting_slot)
+        mamba_avail = req_to_token_pool.mamba_allocator.available_size()
+
+        loaded = cache.load_back(leaf, mem_quota=-(10**9), req=req)
+
+        self.assertFalse(loaded)
+        self.assertIs(req.mamba_pool_idx, preexisting_slot)
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), mamba_avail
+        )
+        self._release_ongoing_load_back_locks(cache)
+
+    def test_load_back_success_publishes_fresh_mamba_slot(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+
+        self._backup_node(cache, leaf)
+        cache.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        req = self._make_req(req_to_token_pool)
+        req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0))
+        req.mamba_pool_idx = None
+        mamba_avail = req_to_token_pool.mamba_allocator.available_size()
+
+        loaded = cache.load_back(leaf, req=req)
+
+        self.assertTrue(loaded)
+        # the successful load must keep the freshly allocated slot published
+        self.assertIsNotNone(req.mamba_pool_idx)
+        self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].value)
+        # one slot restores the node's mamba value, one is the request's CoW slot
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), mamba_avail - 2
+        )
+        self._finish_pending_loads(cache)
+        self._release_ongoing_load_back_locks(cache)
+
+    def test_load_back_success_copies_mamba_state_into_request_slot(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+
+        # Stamp the node's mamba state so the host backup carries it.
+        node_mamba_indices = leaf.component_data[ComponentType.MAMBA].value.clone()
+        self._fill_mamba_state(req_to_token_pool, node_mamba_indices, marker=11)
+        expected_temporal, expected_conv = self._snapshot_mamba_state(
+            req_to_token_pool, node_mamba_indices
+        )
+
+        self._backup_node(cache, leaf)
+        cache.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        req = self._make_req(req_to_token_pool)
+        req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0))
+        req.mamba_pool_idx = None
+
+        loaded = cache.load_back(leaf, req=req)
+        self.assertTrue(loaded)
+        self.assertIsNotNone(req.mamba_pool_idx)
+        self._finish_pending_loads(cache)
+
+        # The CoW slot must actually hold the backed-up mamba state, not merely exist.
+        actual_temporal, actual_conv = self._snapshot_mamba_state(
+            req_to_token_pool, req.mamba_pool_idx.unsqueeze(0)
+        )
+        self.assertTrue(torch.equal(actual_temporal, expected_temporal))
+        self.assertEqual(len(actual_conv), len(expected_conv))
+        for actual, expected in zip(actual_conv, expected_conv):
+            self.assertTrue(torch.equal(actual, expected))
+        self._release_ongoing_load_back_locks(cache)
+
+    def test_prepare_load_back_mamba(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        comp = cache.components[ComponentType.MAMBA]
+
+        self._backup_node(cache, leaf)
+        cache.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        # a request that already owns a slot -> nothing to prepare
+        req = self._make_req(req_to_token_pool)
+        self.assertIsNone(comp.prepare_load_back(leaf, req=req).allocated_mamba_slot)
+
+        # no request -> nothing to prepare
+        self.assertIsNone(comp.prepare_load_back(leaf, req=None).allocated_mamba_slot)
+
+        # fresh request + host-backed mamba -> allocates and publishes onto req
+        req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0))
+        req.mamba_pool_idx = None
+        prep = comp.prepare_load_back(leaf, req=req)
+        self.assertIsNotNone(prep.allocated_mamba_slot)
+        self.assertEqual(int(req.mamba_pool_idx), int(prep.allocated_mamba_slot[0]))
+
+        # node without host-backed mamba -> nothing to prepare
+        req2 = self._make_req(req_to_token_pool)
+        req_to_token_pool.mamba_allocator.free(req2.mamba_pool_idx.unsqueeze(0))
+        req2.mamba_pool_idx = None
+        root = cache.root_node
+        self.assertIsNone(root.component_data[ComponentType.MAMBA].host_value)
+        self.assertIsNone(comp.prepare_load_back(root, req=req2).allocated_mamba_slot)
+        self.assertIsNone(req2.mamba_pool_idx)
+
+    def test_prepare_load_back_skips_device_present_node(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        comp = cache.components[ComponentType.MAMBA]
+
+        # Back up without evicting: device value stays and a host copy is added, so build_hicache_transfers no-ops and prepare must not allocate a dead slot.
+        self._backup_node(cache, leaf)
+        cd = leaf.component_data[ComponentType.MAMBA]
+        self.assertIsNotNone(cd.value)
+        self.assertIsNotNone(cd.host_value)
+
+        req = self._make_req(req_to_token_pool)
+        req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0))
+        req.mamba_pool_idx = None
+        mamba_avail = req_to_token_pool.mamba_allocator.available_size()
+
+        self.assertIsNone(comp.prepare_load_back(leaf, req=req).allocated_mamba_slot)
+        self.assertIsNone(req.mamba_pool_idx)
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), mamba_avail
+        )
+
+    def test_prepare_load_back_mamba_pool_exhausted(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        comp = cache.components[ComponentType.MAMBA]
+
+        self._backup_node(cache, leaf)
+        cache.evict(EvictParams(num_tokens=len(leaf.key)))
+
+        req = self._make_req(req_to_token_pool)
+        req_to_token_pool.mamba_allocator.free(req.mamba_pool_idx.unsqueeze(0))
+        req.mamba_pool_idx = None
+        retry_slot = req_to_token_pool.mamba_allocator.alloc(1)
+
+        # first alloc fails -> prepare must evict a mamba slot and retry
+        with mock.patch.object(
+            req_to_token_pool.mamba_allocator, "alloc", side_effect=[None, retry_slot]
+        ), mock.patch.object(cache, "evict", autospec=True) as evict:
+            prep = comp.prepare_load_back(leaf, req=req)
+        evict.assert_called_once_with(EvictParams(num_tokens=0, mamba_num=1))
+        self.assertIs(prep.allocated_mamba_slot, retry_slot)
+        self.assertEqual(int(req.mamba_pool_idx), int(retry_slot[0]))
+
     def test_scheduler_hicache_aux_only_load_back_appends_full_device_indices(self):
         if self.cfg.page_size != 1:
             self.skipTest("page_size=1 keeps the expected suffix precise")