Reset only the used mamba state on unified radix cache (#31648)

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]:

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),