From afa3c06d1fc9396b13235cfe358ae5a96170feb6 Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Tue, 14 Jul 2026 15:17:08 +0800 Subject: [PATCH] Using UnifiedRadixTree by default for SWA, Mamba, and DSA models (#30468) Co-authored-by: ispobock --- .../srt/kv_canary/radix_cache_walker.py | 9 ++- .../sglang/srt/mem_cache/kv_cache_builder.py | 4 +- python/sglang/srt/mem_cache/registry.py | 30 +++++----- .../srt/mem_cache/unified_radix_cache.py | 4 +- .../test/scripted_runtime/context/radix.py | 3 + .../unit/mem_cache/test_registry.py | 55 ++++++++++++++++--- 6 files changed, 78 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/kv_canary/radix_cache_walker.py b/python/sglang/srt/kv_canary/radix_cache_walker.py index f70b15869..1b155ff5d 100644 --- a/python/sglang/srt/kv_canary/radix_cache_walker.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 24bb32796..09c1b74ff 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index e4fc6a8c4..744134e01 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/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: - # HybridModel launches HiCache via UnifiedRadixCache by default. + if ctx.is_hybrid_ssm or ctx.is_hybrid_swa or ctx.is_dsa: + # 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, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 9bb07beb4..9677de9e6 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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 = { diff --git a/python/sglang/test/scripted_runtime/context/radix.py b/python/sglang/test/scripted_runtime/context/radix.py index 45faf2868..333ac6957 100644 --- a/python/sglang/test/scripted_runtime/context/radix.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index 6947b851f..1f25923da 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/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.return_value = MagicMock() + # SWA hybrid models now default to the unified radix tree. + 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) - self.assertIs(result, SWA.return_value) + fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params) + 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.return_value = MagicMock() + # Mamba hybrid models now default to the unified radix tree. + 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) - self.assertIs(result, Mamba.return_value) + fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params) + self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_lmc_radix_cache_when_enable_lmcache(self): ctx = _make_ctx(enable_lmcache=True)