minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Zhangheng
2026-07-14 15:17:08 +08:00
committed by GitHub
co-authored by ispobock
parent 7e0b29ac03
commit afa3c06d1f
6 changed files with 78 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/kv_canary/radix_cache_walker.py
+8 -1
View File
@@ -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
+3 -1
View File
@@ -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
+14 -16
View File
@@ -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,
python/sglang/srt/mem_cache/unified_radix_cache.py
+3 -1
View File
@@ -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
View File
@@ -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
+47 -8
View File
@@ -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)