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