diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index f619f17ad..6901521aa 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -96,36 +96,12 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: return RadixCacheCpp(params=params, server_args=server_args) if envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get() or use_mlx(): - from sglang.srt.mem_cache.unified_cache_components import ComponentType - from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache - - tree_components = [ComponentType.FULL] - if ctx.is_hybrid_swa or ctx.is_hybrid_ssm: - tree_components.append( - ComponentType.SWA if ctx.is_hybrid_swa else ComponentType.MAMBA - ) - params.tree_components = tuple(tree_components) - if use_mlx() and ctx.is_hybrid_ssm: - from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import ( - MlxAuxiliaryStateComponent, - ) - - params.component_registry_override = { - ComponentType.MAMBA: MlxAuxiliaryStateComponent, - } - cache = UnifiedRadixCache(params) - if ctx.enable_hierarchical_cache: - cache.init_hicache(server_args, params) - ctx.tp_worker.register_hicache_layer_transfer_counter( - cache.cache_controller.layer_done_counter - ) - return cache + return _create_unified_radix_cache(ctx, server_args, params) if ctx.enable_hierarchical_cache: - if ctx.is_hybrid_ssm: - from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache - - cache = HiMambaRadixCache(params=params, server_args=server_args) + if ctx.is_hybrid_ssm or ctx.is_hybrid_swa: + # HybridModel launches HiCache via UnifiedRadixCache by default. + return _create_unified_radix_cache(ctx, server_args, params) else: from sglang.srt.mem_cache.hiradix_cache import HiRadixCache @@ -163,6 +139,39 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: return RadixCache(params) +def _create_unified_radix_cache( + ctx: TreeCacheBuildContext, + server_args: "ServerArgs", + params: CacheInitParams, +) -> BasePrefixCache: + """Initialize a UnifiedRadixCache with proper components and optional HiCache.""" + from sglang.srt.mem_cache.unified_cache_components import ComponentType + from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache + + tree_components = [ComponentType.FULL] + if ctx.is_hybrid_swa: + tree_components.append(ComponentType.SWA) + if ctx.is_hybrid_ssm: + tree_components.append(ComponentType.MAMBA) + + params.tree_components = tuple(tree_components) + if use_mlx() and ctx.is_hybrid_ssm: + from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import ( + MlxAuxiliaryStateComponent, + ) + + params.component_registry_override = { + ComponentType.MAMBA: MlxAuxiliaryStateComponent, + } + cache = UnifiedRadixCache(params) + if ctx.enable_hierarchical_cache: + cache.init_hicache(server_args, params) + ctx.tp_worker.register_hicache_layer_transfer_counter( + cache.cache_controller.layer_done_counter + ) + return cache + + def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache: """Route to the matching factory to construct Radix Cache.""" name = ctx.server_args.radix_cache_backend diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index 026ee98c4..61b22390c 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -233,20 +233,45 @@ class TestDefaultRadixCacheFactory(CustomTestCase): ctx.tp_worker.register_hicache_layer_transfer_counter.assert_called_once() self.assertIs(result, fake_module.HiRadixCache.return_value) - def test_hi_mamba_radix_cache_when_hierarchical_and_hybrid_ssm(self): + def test_unified_radix_cache_when_hierarchical_and_hybrid_ssm(self): ctx = _make_ctx(enable_hierarchical_cache=True, is_hybrid_ssm=True) - # `hi_mamba_radix_cache` imports `hicache_storage`, which - # transitively loads `sgl_kernel`; inject a stand-in module. - fake_module = MagicMock() + # Hybrid SSM with hierarchical cache now uses UnifiedRadixCache. + fake_components = MagicMock() + fake_radix = MagicMock() with patch.dict( "sys.modules", - {"sglang.srt.mem_cache.hi_mamba_radix_cache": fake_module}, + { + "sglang.srt.mem_cache.unified_cache_components": fake_components, + "sglang.srt.mem_cache.unified_radix_cache": fake_radix, + }, ): result = default_radix_cache_factory(ctx) - fake_module.HiMambaRadixCache.assert_called_once_with( - params=ctx.params, server_args=ctx.server_args + fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params) + fake_radix.UnifiedRadixCache.return_value.init_hicache.assert_called_once_with( + ctx.server_args, ctx.params ) - self.assertIs(result, fake_module.HiMambaRadixCache.return_value) + ctx.tp_worker.register_hicache_layer_transfer_counter.assert_called_once() + self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) + + def test_unified_radix_cache_when_hierarchical_and_hybrid_swa(self): + ctx = _make_ctx(enable_hierarchical_cache=True, is_hybrid_swa=True) + # Hybrid SWA with hierarchical cache also uses UnifiedRadixCache. + fake_components = MagicMock() + fake_radix = MagicMock() + with patch.dict( + "sys.modules", + { + "sglang.srt.mem_cache.unified_cache_components": fake_components, + "sglang.srt.mem_cache.unified_radix_cache": fake_radix, + }, + ): + result = default_radix_cache_factory(ctx) + fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params) + fake_radix.UnifiedRadixCache.return_value.init_hicache.assert_called_once_with( + ctx.server_args, ctx.params + ) + ctx.tp_worker.register_hicache_layer_transfer_counter.assert_called_once() + self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_swa_radix_cache_when_hybrid_swa(self): ctx = _make_ctx(is_hybrid_swa=True)