[UnifiedTree]: HybridModel launches HiCache via UnifiedTree by default. (#27759)

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
+        return _create_unified_radix_cache(ctx, server_args, params)
-        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
 
     if ctx.enable_hierarchical_cache:
-        if ctx.is_hybrid_ssm:
+        if ctx.is_hybrid_ssm or ctx.is_hybrid_swa:
-            from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache
+            # HybridModel launches HiCache via UnifiedRadixCache by default.
-
+            return _create_unified_radix_cache(ctx, server_args, params)
-            cache = HiMambaRadixCache(params=params, server_args=server_args)
         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

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
+        # Hybrid SSM with hierarchical cache now uses UnifiedRadixCache.
-        # transitively loads `sgl_kernel`; inject a stand-in module.
+        fake_components = MagicMock()
-        fake_module = 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(
+            fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params)
-                params=ctx.params, server_args=ctx.server_args
+            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)