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 edfbaa052..670881e2d 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 @@ -23,6 +23,7 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, EvictLayer, + LRURefreshPhase, TreeComponent, get_and_increase_time_counter, ) @@ -56,6 +57,30 @@ class MambaComponent(TreeComponent): # HiCache state self._mamba_pool_host = None # set to host mamba pool when HiCache enabled + def refresh_lru( + self, + phase: LRURefreshPhase, + node: UnifiedTreeNode, + root_node: UnifiedTreeNode, + ) -> None: + # A match consumes only best_match_node's mamba state (cf. inc_lock_ref, + # which locks just this node's mamba value), unlike Full whose whole matched + # path is reused as prefix. Refreshing ancestors would keep a whole session's + # states adjacent in the mamba LRU and evict cold sessions wholesale, so touch + # only the used state. New leaf states enter the LRU via + # commit_insert_component_data, so the insert walk (WALKDOWN) is a no-op here. + ct = self.component_type + match phase: + case LRURefreshPhase.WALKDOWN: + return + case LRURefreshPhase.MATCH_END: + if node.component_data[ct].value is not None: + self.cache.lru_lists[ct].reset_node_mru(node) + case LRURefreshPhase.INSERT_END: + return + case _: + raise ValueError(f"Unknown LRURefreshPhase: {phase}") + def create_match_validator( self, match_device_only: bool = False ) -> Callable[[UnifiedTreeNode], bool]: 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 d613275dd..9a1407bf6 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 @@ -4053,6 +4053,78 @@ class UnifiedLRUListBoundedRefreshTest(CustomTestCase): self.assertEqual(self._lru_order(lru), before) +class TestUnifiedMambaLRUMatchRefresh(CustomTestCase): + """A prefix-cache hit must refresh only best_match_node's mamba state in the + mamba LRU, not its ancestors. The base TreeComponent refresh is whole-chain + (reset_node_and_parents_mru), which clusters a session's states so that under + mamba-pool pressure eviction drops whole cold sessions instead of the + intermediate states reuse never needs. Guards MambaComponent.refresh_lru's + single-node MATCH_END override. + """ + + cfg = CacheConfig(page_size=1, components=(ComponentType.FULL, ComponentType.MAMBA)) + + def _mamba_lru_mru_to_lru(self, cache): + lru = cache.lru_lists[ComponentType.MAMBA] + pt = lru._pt + out, cur = [], lru.head.lru_next[pt] + while cur is not lru.tail: + out.append(cur) + cur = cur.lru_next[pt] + return out + + def _make_req(self, req_to_token_pool): + req = Req( + rid=0, + origin_input_text="", + origin_input_ids=array("q"), + sampling_params=SamplingParams(temperature=0, max_new_tokens=1), + ) + req_to_token_pool.alloc([req]) + return req + + def test_match_refreshes_only_used_node(self): + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + + def insert(tokens): + value = allocator.alloc(len(tokens)) + req = self._make_req(req_to_token_pool) + cache.insert( + InsertParams( + key=RadixKey(array("q", tokens)), + value=value[: len(tokens)], + mamba_value=req.mamba_pool_idx.unsqueeze(0), + ) + ) + + def match_leaf(tokens): + return cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).best_match_node + + # Two independent sessions, each a 2-node mamba chain: + # root -> a1 -> b1 and root -> a2 -> b2 + insert([1, 2, 3]) + insert([1, 2, 3, 4, 5, 6]) + insert([7, 8, 9]) + insert([7, 8, 9, 10, 11, 12]) + + b1 = match_leaf([1, 2, 3, 4, 5, 6]) + a1 = b1.parent + b2 = match_leaf([7, 8, 9, 10, 11, 12]) + a2 = b2.parent + # Session 2 matched last, so session 1's ancestor a1 is older than a2. + order = self._mamba_lru_mru_to_lru(cache) + self.assertGreater(order.index(a1), order.index(a2)) + + # Re-access session 1: only its consumed leaf b1 moves to MRU; ancestor a1 + # must stay put -- whole-chain reset would bump a1 ahead of a2. + self.assertIs(match_leaf([1, 2, 3, 4, 5, 6]), b1) + order = self._mamba_lru_mru_to_lru(cache) + self.assertIs(order[0], b1) + self.assertGreater(order.index(a1), order.index(a2)) + + class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase): cfg = CacheConfig( components=(ComponentType.FULL, ComponentType.MAMBA),