Using UnifiedRadixTree by default for SWA, Mamba, and DSA models (#30468)

Co-authored-by: ispobock <ispobaoke@gmail.com>

python/sglang/srt/kv_canary/radix_cache_walker.py

@@ -93,7 +93,9 @@ def _walk_radix_subtree(
 
     if unlocked_only:
         emit_slots = not is_root and _node_is_unlocked_for_canary(
-            node=node, radix_cache=radix_cache
+            node=node,
+            radix_cache=radix_cache,
+            swa_resident_only=swa_resident_only,
         )
     else:
         emit_slots = not is_root
@@ -168,6 +170,7 @@ def _node_is_unlocked_for_canary(
     *,
     node: TreeNode | UnifiedTreeNode,
     radix_cache: BasePrefixCache,
+    swa_resident_only: bool,
 ) -> bool:
     if type(radix_cache) is RadixCache:
         return node.lock_ref == 0
@@ -176,6 +179,10 @@ def _node_is_unlocked_for_canary(
         return node.full_lock_ref == 0
 
     if type(radix_cache) is UnifiedRadixCache:
+        if swa_resident_only and radix_cache.supports_swa():
+            # Unified SWA owns an independent component lock. A node can still
+            # hold Full KV for a running request while its SWA slots are unused.
+            return node.component_data[ComponentType.SWA].lock_ref == 0
         return node.component_data[BASE_COMPONENT_TYPE].lock_ref == 0
 
     raise NotImplementedError(

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -23,7 +23,7 @@ class KVCacheBuildResult:
 
 from typing import TYPE_CHECKING
 
-from sglang.srt.configs.model_config import ModelImpl
+from sglang.srt.configs.model_config import ModelImpl, is_deepseek_dsa
 from sglang.srt.environ import envs
 from sglang.srt.managers.mm_utils import init_mm_embedding_cache
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
@@ -155,6 +155,7 @@ def build_kv_cache(
         or tp_worker.model_runner.kimi_linear_config is not None
         or tp_worker.model_runner.hybrid_lightning_config is not None
     )
+    is_dsa = is_deepseek_dsa(model_config.hf_config)
 
     sliding_window_size = None
     if is_hybrid_swa:
@@ -234,6 +235,7 @@ def build_kv_cache(
             is_hybrid_swa=is_hybrid_swa,
             full_tokens_per_layer=full_tokens_per_layer,
             is_hybrid_ssm=is_hybrid_ssm,
+            is_dsa=is_dsa,
             enable_hierarchical_cache=enable_hierarchical_cache,
             disable_radix_cache=disable_radix_cache,
             effective_chunked_prefill_size=effective_chunked_prefill_size,

python/sglang/srt/mem_cache/registry.py

@@ -43,6 +43,7 @@ class TreeCacheBuildContext:
     tp_rank: int
     tp_group: Any
     full_tokens_per_layer: Optional[int] = None
+    is_dsa: bool = False
 
 
 RadixCacheFactory = Callable[[TreeCacheBuildContext], BasePrefixCache]
@@ -103,9 +104,20 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
     if envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get() or use_mlx():
         return _create_unified_radix_cache(ctx, server_args, params)
 
+    if ctx.is_hybrid_swa:
+        if ctx.full_tokens_per_layer == 0:
+            from sglang.srt.mem_cache.pure_swa_radix_cache import PureSWARadixCache
+
+            return PureSWARadixCache(params=params)
+        return _create_unified_radix_cache(ctx, server_args, params)
+
+    if ctx.is_hybrid_ssm:
+        return _create_unified_radix_cache(ctx, server_args, params)
+
     if ctx.enable_hierarchical_cache:
-        if ctx.is_hybrid_ssm or ctx.is_hybrid_swa:
+        if ctx.is_hybrid_ssm or ctx.is_hybrid_swa or ctx.is_dsa:
-            # HybridModel launches HiCache via UnifiedRadixCache by default.
+            # HybridModel and DSA (e.g. DeepSeek V3.2 / GLM-5.1) launch
+            # HiCache via UnifiedRadixCache by default.
             return _create_unified_radix_cache(ctx, server_args, params)
         else:
             from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
@@ -116,20 +128,6 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
         )
         return cache
 
-    if ctx.is_hybrid_swa:
-        if ctx.full_tokens_per_layer == 0:
-            from sglang.srt.mem_cache.pure_swa_radix_cache import PureSWARadixCache
-
-            return PureSWARadixCache(params=params)
-        from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
-
-        return SWARadixCache(params=params)
-
-    if ctx.is_hybrid_ssm:
-        from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
-
-        return MambaRadixCache(params)
-
     if server_args.enable_lmcache:
         from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
             LMCRadixCache,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -312,7 +312,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
         self.page_size = params.page_size
         self.disable = params.disable
-        self.is_eagle = params.is_eagle
         self.enable_kv_cache_events = params.enable_kv_cache_events
         self.kv_event_queue = []
         self.eviction_policy = params.eviction_policy.lower()
@@ -332,6 +331,9 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
 
         assert params.tree_components is not None
         self.tree_components = tuple(params.tree_components)
+        self.is_eagle = (
+            params.is_eagle and ComponentType.MAMBA not in self.tree_components
+        )
         component_registry = COMPONENT_REGISTRY
         if params.component_registry_override:
             component_registry = {

python/sglang/test/scripted_runtime/context/radix.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Callable, Dict
 
 from sglang.srt.mem_cache.swa_radix_cache import TreeNode as SWATreeNode
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode
 
 if TYPE_CHECKING:
     from sglang.test.scripted_runtime.context.api import ScriptedContext
@@ -19,6 +20,8 @@ def get_all_node_lock_refs(ctx: ScriptedContext) -> Dict[int, int]:
 def _node_lock_ref(node: Any) -> int:
     if isinstance(node, SWATreeNode):
         return node.full_lock_ref + node.swa_lock_ref
+    if isinstance(node, UnifiedTreeNode):
+        return sum(cd.lock_ref for cd in node.component_data)
     return node.lock_ref
 
 

test/registered/unit/mem_cache/test_registry.py

@@ -26,6 +26,7 @@ def _make_ctx(
     enable_lmcache=False,
     is_hybrid_swa=False,
     is_hybrid_ssm=False,
+    is_dsa=False,
     enable_hierarchical_cache=False,
     disable_radix_cache=False,
     effective_chunked_prefill_size=None,
@@ -41,6 +42,7 @@ def _make_ctx(
         params=MagicMock(),
         is_hybrid_swa=is_hybrid_swa,
         is_hybrid_ssm=is_hybrid_ssm,
+        is_dsa=is_dsa,
         enable_hierarchical_cache=enable_hierarchical_cache,
         disable_radix_cache=disable_radix_cache,
         effective_chunked_prefill_size=effective_chunked_prefill_size,
@@ -291,13 +293,42 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
             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_dsa(self):
+        ctx = _make_ctx(enable_hierarchical_cache=True, is_dsa=True)
+        # DSA models (e.g. DeepSeek V3.2 / GLM-5.1) with hierarchical cache
+        # use 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)
-        with patch("sglang.srt.mem_cache.swa_radix_cache.SWARadixCache") as SWA:
+        # SWA hybrid models now default to the unified radix tree.
-            SWA.return_value = MagicMock()
+        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)
-            SWA.assert_called_once_with(params=ctx.params)
+            fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params)
-            self.assertIs(result, SWA.return_value)
+            self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
 
     def test_pure_swa_radix_cache_when_all_swa(self):
         ctx = _make_ctx(is_hybrid_swa=True, full_tokens_per_layer=0)
@@ -311,11 +342,19 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
 
     def test_mamba_radix_cache_when_hybrid_ssm(self):
         ctx = _make_ctx(is_hybrid_ssm=True)
-        with patch("sglang.srt.mem_cache.mamba_radix_cache.MambaRadixCache") as Mamba:
+        # Mamba hybrid models now default to the unified radix tree.
-            Mamba.return_value = MagicMock()
+        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)
-            Mamba.assert_called_once_with(ctx.params)
+            fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params)
-            self.assertIs(result, Mamba.return_value)
+            self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
 
     def test_lmc_radix_cache_when_enable_lmcache(self):
         ctx = _make_ctx(enable_lmcache=True)