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