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

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

Browse Source
This commit is contained in:
Zhangheng
2026-06-11 12:03:36 +08:00
committed by GitHub
parent 9788c8e867
commit 5e0271536a
2 changed files with 70 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/registry.py
+37 -28
View File
@@ -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
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
return _create_unified_radix_cache(ctx, server_args, params)
if ctx.enable_hierarchical_cache:
if ctx.is_hybrid_ssm:
from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache
cache = HiMambaRadixCache(params=params, server_args=server_args)
if ctx.is_hybrid_ssm or ctx.is_hybrid_swa:
# HybridModel launches HiCache via UnifiedRadixCache by default.
return _create_unified_radix_cache(ctx, server_args, params)
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
+33 -8
View File
@@ -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
# transitively loads `sgl_kernel`; inject a stand-in module.
fake_module = MagicMock()
# Hybrid SSM with hierarchical cache now uses UnifiedRadixCache.
fake_components = 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(
params=ctx.params, server_args=ctx.server_args
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
)
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)