[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