minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
Jialin Ouyang
2026-07-23 18:48:31 +08:00
committed by GitHub
co-authored by hzh0425
parent c18919f8f3
commit 70ac0c4b0e
5 changed files with 345 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_cache_components/__init__.py
+2
View File
@@ -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
+34 -8
View File
@@ -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
+24
View File
@@ -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
+35 -1
View File
@@ -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
+250
View File
@@ -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")