From 70ac0c4b0e8e832945c245d9fb0696f368e4aa19 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Thu, 23 Jul 2026 03:48:31 -0700 Subject: [PATCH] [UnifiedRadixCache][mamba] Fix mamba state corruption and slot leak when load_back aborts (#30986) Co-authored-by: hzh0425 --- .../unified_cache_components/__init__.py | 2 + .../mamba_component.py | 42 ++- .../tree_component.py | 24 ++ .../srt/mem_cache/unified_radix_cache.py | 36 ++- .../test_unified_radix_cache_unittest.py | 250 ++++++++++++++++++ 5 files changed, 345 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py b/python/sglang/srt/mem_cache/unified_cache_components/__init__.py index 144b27e51..60d27d6d4 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py +++ b/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", diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index 4cee69d2f..58af3f9f3 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index 425eea226..43b723ed0 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 6e80925f0..3694620b6 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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: diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index db6c47843..3245b42e0 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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")