[Mamba] Add a per-path cap for cached states (#31230)

python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py

@@ -54,6 +54,7 @@ class MambaComponent(TreeComponent):
         super().__init__(cache, params)
         self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer
         self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy
+        self.mamba_max_states_per_path = get_server_args().mamba_max_states_per_path
         # HiCache state
         self._mamba_pool_host = None  # set to host mamba pool when HiCache enabled
 
@@ -161,6 +162,7 @@ class MambaComponent(TreeComponent):
             self.cache.component_evictable_size_[self.component_type] += len(
                 params.mamba_value
             )
+            self._evict_excess_path_states(node)
             return
         if node.component_data[self.component_type].value is None:
             node.component_data[self.component_type].value = params.mamba_value
@@ -173,11 +175,49 @@ class MambaComponent(TreeComponent):
                 params.mamba_value
             )
             node.last_access_time = get_and_increase_time_counter()
+            self._evict_excess_path_states(node)
             return
         self.cache.lru_lists[self.component_type].reset_node_mru(node)
         node.last_access_time = get_and_increase_time_counter()
         result.mamba_exist = True
 
+    def _evict_excess_path_states(self, tail: UnifiedTreeNode) -> None:
+        """Evict shallow eligible device checkpoints beyond the path cap.
+
+        Full KV and any existing host backup are retained. The tail, forks,
+        locked nodes, and device leaves are preserved, so the cap is a
+        best-effort soft limit.
+        """
+        cap = self.mamba_max_states_per_path
+        if cap < 0:
+            return
+
+        ct = self.component_type
+        holders = []
+        node = tail
+        while node is not None and node is not self.cache.root_node:
+            if node.component_data[ct].value is not None:
+                holders.append(node)
+            node = node.parent
+
+        excess = len(holders) - cap
+        if excess <= 0:
+            return
+
+        tracker = {component: 0 for component in self.cache.tree_components}
+        for node in reversed(holders):
+            if excess <= 0 or node is tail:
+                break
+            if node.component_data[ct].lock_ref > 0 or len(node.children) != 1:
+                continue
+            if node in self.cache.evictable_device_leaves:
+                continue
+            self.cache._evict_component_and_detach_lru(
+                node, self, target=EvictLayer.DEVICE, tracker=tracker
+            )
+            self.cache._cascade_evict(node, self, tracker)
+            excess -= 1
+
     def redistribute_on_node_split(
         self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
     ):

python/sglang/srt/server_args.py

@@ -2383,6 +2383,13 @@ class ServerArgs:
         ),
         NS("exec.mamba"),
     ] = None
+    mamba_max_states_per_path: A[
+        int,
+        "Maximum number of cached Mamba states retained per root-to-tail path "
+        "(-1 means unlimited). When enabled, after each insert the shallowest eligible "
+        "interior states beyond the cap are removed while their full KV remains. "
+        "Tail, fork, and locked nodes are preserved. Must be -1 or a positive integer.",
+    ] = -1
     enable_mamba_cache_stochastic_rounding: A[
         bool,
         "Enable stochastic rounding when writing FP16 Mamba SSM cache states. Requires --mamba-ssm-dtype float16 and CUDA. With --mamba-backend triton, requires SM100.",
@@ -3310,6 +3317,8 @@ class ServerArgs:
         # _handle_model_specific_adjustments never runs.
         self._resolved_overrides = []
 
+        self._validate_mamba_max_states_per_path()
+
         if self.model_path.lower() in ["none", "dummy"]:
             return
 
@@ -3487,6 +3496,14 @@ class ServerArgs:
 
         materialize_declarations(self)
 
+    def _validate_mamba_max_states_per_path(self):
+        value = self.mamba_max_states_per_path
+        if value == 0 or value < -1:
+            raise ValueError(
+                "--mamba-max-states-per-path must be -1 (unlimited) or a positive "
+                f"integer, got {value}."
+            )
+
     def _handle_model_capability_adjustments(self):
         if parse_connector_type(self.model_path) == ConnectorType.INSTANCE:
             return

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

@@ -60,6 +60,8 @@ class TestUnifiedMambaRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
                 "extra_buffer",
                 "--mamba-track-interval",
                 str(MAMBA_TRACK_INTERVAL),
+                "--mamba-max-states-per-path",
+                "3",
             ],
             env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
         )

test/registered/unit/mem_cache/test_mamba_path_state_cap.py

@@ -0,0 +1,179 @@
+"""CPU-only unit tests for the per-path Mamba checkpoint cap."""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+import argparse
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.mem_cache.unified_cache_components.mamba_component import (
+    MambaComponent,
+)
+from sglang.srt.mem_cache.unified_cache_components.tree_component import (
+    ComponentType,
+)
+from sglang.srt.mem_cache.unified_radix_cache import (
+    UnifiedLRUList,
+    UnifiedRadixCache,
+    UnifiedTreeNode,
+)
+from sglang.srt.server_args import ServerArgs
+
+
+class _RecordingAllocator:
+    def __init__(self):
+        self.freed = []
+
+    def free(self, value):
+        self.freed.extend(value.tolist())
+
+
+class _FakeUnifiedCache:
+    tree_components = (ComponentType.FULL, ComponentType.MAMBA)
+
+    def __init__(self):
+        self.root_node = UnifiedTreeNode(self.tree_components)
+        self.evictable_device_leaves = set()
+        self.req_to_token_pool = SimpleNamespace(mamba_allocator=_RecordingAllocator())
+        self.component_evictable_size_ = {ComponentType.MAMBA: 0}
+        self.component_protected_size_ = {ComponentType.MAMBA: 0}
+        self.lru_lists = {
+            ComponentType.MAMBA: UnifiedLRUList(
+                ComponentType.MAMBA, self.tree_components
+            )
+        }
+        self.host_lru_lists = {
+            ComponentType.MAMBA: UnifiedLRUList(
+                ComponentType.MAMBA, self.tree_components
+            )
+        }
+        self.evicted = []
+        self.cascaded = []
+
+    def _evict_component_and_detach_lru(self, node, component, **kwargs):
+        self.evicted.append(node)
+        return UnifiedRadixCache._evict_component_and_detach_lru(
+            self, node, component, **kwargs
+        )
+
+    def _cascade_evict(self, node, component, tracker):
+        self.cascaded.append(node)
+
+
+def _build_unified_chain(cap, length=3):
+    cache = _FakeUnifiedCache()
+    component = object.__new__(MambaComponent)
+    component.cache = cache
+    component.mamba_max_states_per_path = cap
+
+    nodes = []
+    parent = cache.root_node
+    for index in range(length):
+        node = UnifiedTreeNode(cache.tree_components)
+        node.parent = parent
+        node.component_data[ComponentType.FULL].value = torch.tensor([100 + index])
+        node.component_data[ComponentType.MAMBA].value = torch.tensor([index])
+        parent.children[index] = node
+        cache.component_evictable_size_[ComponentType.MAMBA] += 1
+        cache.lru_lists[ComponentType.MAMBA].insert_mru(node)
+        nodes.append(node)
+        parent = node
+    return component, nodes, cache
+
+
+class TestMambaPathStateCap(unittest.TestCase):
+    def test_server_arg_defaults_to_unlimited(self):
+        self.assertEqual(
+            ServerArgs(model_path="dummy").mamba_max_states_per_path,
+            -1,
+        )
+
+    def test_server_arg_cli(self):
+        parser = argparse.ArgumentParser()
+        ServerArgs.add_cli_args(parser)
+
+        args = parser.parse_args(
+            ["--model-path", "dummy", "--mamba-max-states-per-path", "3"]
+        )
+
+        self.assertEqual(args.mamba_max_states_per_path, 3)
+
+    def test_server_arg_rejects_zero_and_values_below_negative_one(self):
+        for value in (0, -2):
+            with self.subTest(value=value), self.assertRaisesRegex(
+                ValueError,
+                "must be -1 \\(unlimited\\) or a positive integer",
+            ):
+                ServerArgs(
+                    model_path="dummy",
+                    mamba_max_states_per_path=value,
+                )
+
+    def test_unified_cache_removes_only_shallow_mamba_state(self):
+        component, nodes, cache = _build_unified_chain(cap=2)
+
+        component._evict_excess_path_states(nodes[-1])
+
+        self.assertEqual(cache.evicted, [nodes[0]])
+        self.assertEqual(cache.cascaded, [nodes[0]])
+        self.assertIsNone(nodes[0].component_data[ComponentType.MAMBA].value)
+        self.assertIsNotNone(nodes[-1].component_data[ComponentType.MAMBA].value)
+        self.assertEqual(
+            cache.req_to_token_pool.mamba_allocator.freed,
+            [0],
+        )
+        self.assertEqual(cache.component_evictable_size_[ComponentType.MAMBA], 2)
+        self.assertFalse(cache.lru_lists[ComponentType.MAMBA].in_list(nodes[0]))
+        self.assertTrue(
+            all(
+                node.component_data[ComponentType.FULL].value is not None
+                for node in nodes
+            )
+        )
+
+    def test_unified_cache_cap_is_soft_for_fork_and_locked_nodes(self):
+        component, nodes, cache = _build_unified_chain(cap=1, length=4)
+        fork_child = UnifiedTreeNode(cache.tree_components)
+        fork_child.parent = nodes[0]
+        nodes[0].children["fork"] = fork_child
+        nodes[1].component_data[ComponentType.MAMBA].lock_ref = 1
+
+        component._evict_excess_path_states(nodes[-1])
+
+        self.assertEqual(cache.evicted, [nodes[2]])
+        self.assertIsNotNone(nodes[0].component_data[ComponentType.MAMBA].value)
+        self.assertIsNotNone(nodes[1].component_data[ComponentType.MAMBA].value)
+        self.assertIsNone(nodes[2].component_data[ComponentType.MAMBA].value)
+        self.assertIsNotNone(nodes[3].component_data[ComponentType.MAMBA].value)
+
+    def test_unified_cache_preserves_existing_host_backup(self):
+        component, nodes, cache = _build_unified_chain(cap=2)
+        mamba_data = nodes[0].component_data[ComponentType.MAMBA]
+        mamba_data.host_value = torch.tensor([10])
+
+        component._evict_excess_path_states(nodes[-1])
+
+        self.assertIsNone(mamba_data.value)
+        self.assertIsNotNone(mamba_data.host_value)
+        self.assertTrue(cache.host_lru_lists[ComponentType.MAMBA].in_list(nodes[0]))
+
+    def test_unified_cache_negative_one_disables_cap(self):
+        component, nodes, cache = _build_unified_chain(cap=-1)
+
+        component._evict_excess_path_states(nodes[-1])
+
+        self.assertEqual(cache.evicted, [])
+        self.assertTrue(
+            all(
+                node.component_data[ComponentType.MAMBA].value is not None
+                for node in nodes
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()