diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 817b388ce..497bb6750 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -595,6 +595,9 @@ class Envs: SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS = EnvBool(True) # Decode batches between SWA out-of-window evictions. SGLANG_SWA_EVICTION_INTERVAL = EnvInt(128) + # Deprecated: the unified radix tree is the default tree cache now, so the + # registry no longer reads this. Kept because a few model/arch call sites + # still assert on it; do not use in new code. SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False) # Registered TreeCore backend serving the unified radix cache. SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND = EnvStr("python") @@ -1624,6 +1627,10 @@ _DEPRECATED_ENVS: Dict[str, _DeprecatedEnv] = { note="DFlash now auto-enables the min-free-slots delay; unset this env. " "To override the threshold, use '--min-free-slots-delay'." ), + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": _DeprecatedEnv( + note="The unified radix tree is the default tree cache now; unset this " + "env. The field is still defined for legacy call sites." + ), } diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 2bb22431b..3ec5c0fc6 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1045,12 +1045,20 @@ class HostPoolGroup: def get_ksize_per_token(self): return self.anchor_entry.host_pool.get_ksize_per_token() + def get_size_per_token(self): + return self.anchor_entry.host_pool.get_size_per_token() + def get_pool(self, name: PoolName): return self.entry_map[name].host_pool def get_page_buffer_meta(self, indices): return self.anchor_entry.host_pool.get_page_buffer_meta(indices) + def get_split_heads_page_buffer_meta(self, indices, split_factor: int): + return self.anchor_entry.host_pool.get_split_heads_page_buffer_meta( + indices, split_factor + ) + def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool: return self.anchor_entry.host_pool.is_stride_page_aligned(page_size_bytes) diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index fb1a0b253..5f3980b8a 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -108,32 +108,10 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: logger.info("Using experimental C++ radix tree implementation.") return RadixCacheCpp(params=params, server_args=server_args) - 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 and ctx.full_tokens_per_layer == 0: + from sglang.srt.mem_cache.pure_swa_radix_cache import PureSWARadixCache - 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 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 - - cache = HiRadixCache(params=params, server_args=server_args) - ctx.tp_worker.register_hicache_layer_transfer_counter( - cache.cache_controller.layer_done_counter - ) - return cache + return PureSWARadixCache(params=params) if get_memory().enable_lmcache: from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import ( @@ -162,9 +140,7 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: os.environ["FLEXKV_CONFIG_PATH"] = get_memory().flexkv_config_file return _flexkv_factory(ctx) - from sglang.srt.mem_cache.radix_cache import RadixCache - - return RadixCache(params) + return _create_unified_radix_cache(ctx, server_args, params) def _create_unified_radix_cache( @@ -254,9 +230,8 @@ def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache: ): raise ValueError( "--enable-session-radix-cache requires UnifiedRadixCache, but " - f"tree_cache is {type(cache).__name__}. Set " - "SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 (or remove " - "--enable-session-radix-cache)." + f"tree_cache is {type(cache).__name__}. Drop the flag or the " + "option that selected another tree cache for this model." ) hicache_attached = cache.cache_controller is not None diff --git a/python/sglang/srt/mem_cache/unified_cache/components/README.md b/python/sglang/srt/mem_cache/unified_cache/components/README.md index 14f2f15c1..5429f1a90 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache/components/README.md @@ -329,7 +329,7 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d ## Construction -`UnifiedRadixCache` is constructed directly from `mem_cache/registry.py` when `SGLANG_ENABLE_UNIFIED_RADIX_TREE` is enabled. The registry sets `params.tree_components` before construction: +`UnifiedRadixCache` is the default tree cache and is constructed directly from `mem_cache/registry.py`. The registry sets `params.tree_components` before construction: - Regular full-attention models → `(ComponentType.FULL,)` - Hybrid SWA models → `(ComponentType.FULL, ComponentType.SWA)` diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index a5c3b9765..25d61251b 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1540,6 +1540,50 @@ class UnifiedRadixCache(BasePrefixCache): def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: return self.tree_core.get_prefix_hash_values(node_id) + def query_storage_hit_length( + self, + last_host_node_id: NodeId, + new_input_tokens: list[int], + last_hash: Optional[str] = None, + prefix_keys: Optional[list[str]] = None, + ) -> int: + """Synchronously probe L3 storage for the reusable prefix length.""" + if ( + not self.enable_storage + or self.cache_controller is None + or self.cache_controller.prefetch_rate_limited() + ): + return 0 + + extra_key, cache_salt = self.tree_core.prefetch_anchor_info(last_host_node_id) + prefetch_key = RadixKey( + new_input_tokens, + extra_key=extra_key, + is_bigram=self.tree_core.is_eagle, + cache_salt=cache_salt, + ).page_aligned(self.page_size) + if len(prefetch_key) < self.prefetch_threshold: + return 0 + + from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + PrefetchOperation, + ) + + operation = PrefetchOperation( + "__storage_hit_query__", + prefetch_key, + last_hash, + prefix_keys, + ) + _, storage_hit_count = self.cache_controller._storage_hit_query(operation) + storage_hit_count_tensor = torch.tensor(storage_hit_count, dtype=torch.int) + self._all_reduce_attn_groups( + storage_hit_count_tensor, torch.distributed.ReduceOp.MIN + ) + storage_hit_count = storage_hit_count_tensor.item() + storage_hit_count -= storage_hit_count % self.page_size + return storage_hit_count + def prefetch_from_storage( self, req_id: str, @@ -2593,6 +2637,27 @@ class UnifiedRadixCache(BasePrefixCache): return self.cache_controller.start_loading() return 0 + def is_load_back_event_done(self, consumer_index: int) -> bool: + """Return True after the local load-back event is complete. + + Mirrors ``HiRadixCache`` so the disagg decode restore state machine + (``DecodeHiCacheTransferMixin``) can gate on load-back completion; the + controller-level ``layer_done_counter`` event is shared across cache + implementations, while the tree-side bookkeeping runs in + ``loading_check``. + """ + if consumer_index < 0 or self.cache_controller is None: + return True + + finish_event = self.cache_controller.layer_done_counter.events[ + consumer_index + ].finish_event + if not finish_event.query(): + return False + + self.loading_check() + return True + # ---- Query / Inspection APIs ---- # These APIs exist for compatibility with other RadixTree implementations. # TODO: simplify and consolidate in a future refactor. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e7c8d4c05..a97ca3cac 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5656,17 +5656,11 @@ class ServerArgs: # enable_multi_layer_eagle for EAGLE moved to the override registry # (arg_groups/overrides.py: _mimo_v2_overrides). - if self.enable_hierarchical_cache: - if not envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get(): - raise ValueError( - "Hierarchical cache for MiMoV2 requires the unified " - "radix tree. Set SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 " - "to enable --enable-hierarchical-cache for this model." - ) - - # MiMoV2 has head_dim != v_head_dim, so the host KV pool uses - # asymmetric K/V allocation. Both kernel/page_first and - # direct/page_first_direct have split K/V transfer paths. + # MiMoV2 hierarchical cache runs on the unified radix tree, which + # is the default tree cache now. MiMoV2 has head_dim != v_head_dim, + # so the host KV pool uses asymmetric K/V allocation. Both + # kernel/page_first and direct/page_first_direct have split K/V + # transfer paths. elif ( "Step3p5ForCausalLM" in model_arch or "Step3p7ForConditionalGeneration" in model_arch diff --git a/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py index 67368f55a..2530348fd 100644 --- a/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py +++ b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, List -from sglang.test.scripted_runtime.context.radix import _node_lock_ref +from sglang.test.scripted_runtime.context.radix import _node_lock_ref, to_node_handle if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler @@ -25,7 +25,7 @@ class ScriptedLockRefExhauster: return target = evictable[0] - tree_cache.inc_lock_ref(target) + tree_cache.inc_lock_ref(to_node_handle(tree_cache, target)) newly_locked = [node for node in evictable if _node_lock_ref(node) > 0] if not newly_locked: @@ -35,7 +35,7 @@ class ScriptedLockRefExhauster: def release(self) -> None: tree_cache = self.scheduler.tree_cache for node in self._locked: - tree_cache.dec_lock_ref(node) + tree_cache.dec_lock_ref(to_node_handle(tree_cache, node)) self._locked.clear() def _evictable_nodes(self) -> List[Any]: diff --git a/python/sglang/test/scripted_runtime/context/radix.py b/python/sglang/test/scripted_runtime/context/radix.py index 333ac6957..9ca8c0f2a 100644 --- a/python/sglang/test/scripted_runtime/context/radix.py +++ b/python/sglang/test/scripted_runtime/context/radix.py @@ -25,6 +25,30 @@ def _node_lock_ref(node: Any) -> int: return node.lock_ref +def resolve_node(tree_cache: Any, node_handle: Any) -> Any: + """Resolve whatever a req or match result carries into a tree node. + + ``UnifiedRadixCache`` hands out NodeIds (ints); every other cache hands out + the node object itself. Returns None when the handle no longer maps to a + live node, which only happens once the node has been freed -- and a freed + node holds no locks. + """ + resolve = getattr(tree_cache, "resolve_node_handle", None) + if resolve is None: + return node_handle + try: + return resolve(node_handle) + except KeyError: + return None + + +def to_node_handle(tree_cache: Any, node: Any) -> Any: + """Inverse of `resolve_node`: what the tree_cache lock APIs accept.""" + if isinstance(node, UnifiedTreeNode): + return node.id + return node + + def _collect_node_attr( ctx: ScriptedContext, get_value: Callable[[Any], int] ) -> Dict[int, int]: diff --git a/python/sglang/test/scripted_runtime/req_handle.py b/python/sglang/test/scripted_runtime/req_handle.py index 24f2cfd88..5638ceb51 100644 --- a/python/sglang/test/scripted_runtime/req_handle.py +++ b/python/sglang/test/scripted_runtime/req_handle.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass from typing import TYPE_CHECKING, Optional -from sglang.test.scripted_runtime.context.radix import _node_lock_ref +from sglang.test.scripted_runtime.context.radix import _node_lock_ref, resolve_node if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -52,7 +52,7 @@ class ScriptedReqHandle: req = self.req if req is None: return 0 - node = req.last_node + node = resolve_node(self.context.scheduler.tree_cache, req.last_node) if node is None: return 0 return _node_lock_ref(node) diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index 9e5cc31d7..80a613ada 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -236,46 +236,43 @@ class TestDefaultRadixCacheFactory(CustomTestCase): ) self.assertIs(result, fake_module.RadixCacheCpp.return_value) - def test_unified_radix_cache_when_env_flag_set(self): + def test_unified_radix_cache_is_the_default(self): ctx = _make_ctx( self, ) # Shim both factory imports — each transitively loads sgl_kernel. fake_components = MagicMock() fake_radix = MagicMock() - with ( - patch( - "sglang.srt.mem_cache.registry.envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get", - return_value=True, - ), - patch.dict( - "sys.modules", - { - "sglang.srt.mem_cache.unified_cache.components": fake_components, - "sglang.srt.mem_cache.unified_radix_cache": fake_radix, - }, - ), + 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) self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) - def test_hi_radix_cache_when_hierarchical(self): + def test_unified_radix_cache_when_hierarchical(self): ctx = _make_ctx(self, enable_hierarchical_cache=True) - # `hiradix_cache` imports `hicache_storage` and - # `memory_pool_host`, both of which transitively load - # `sgl_kernel`; inject a stand-in module. - fake_module = MagicMock() + # Full attention with hierarchical cache also uses UnifiedRadixCache. + fake_components = MagicMock() + fake_radix = MagicMock() with patch.dict( "sys.modules", - {"sglang.srt.mem_cache.hiradix_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.HiRadixCache.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 ) ctx.tp_worker.register_hicache_layer_transfer_counter.assert_called_once() - self.assertIs(result, fake_module.HiRadixCache.return_value) + self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_unified_radix_cache_when_hierarchical_and_hybrid_ssm(self): ctx = _make_ctx(self, enable_hierarchical_cache=True, is_hybrid_ssm=True) @@ -400,16 +397,6 @@ class TestDefaultRadixCacheFactory(CustomTestCase): ) self.assertIs(result, fake_module.LMCRadixCache.return_value) - def test_fallback_to_radix_cache(self): - ctx = _make_ctx( - self, - ) - with patch("sglang.srt.mem_cache.radix_cache.RadixCache") as RadixCache: - RadixCache.return_value = MagicMock() - result = default_radix_cache_factory(ctx) - RadixCache.assert_called_once_with(ctx.params) - self.assertIs(result, RadixCache.return_value) - if __name__ == "__main__": unittest.main()