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 670881e2d..4cee69d2f 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 @@ -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 ): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1d2444762..e7de7981b 100644 --- a/python/sglang/srt/server_args.py +++ b/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 diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py index e1f150b65..5465cb32f 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py +++ b/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"}, ) diff --git a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py new file mode 100644 index 000000000..be0eddc1f --- /dev/null +++ b/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()