Using unified radix tree by default for all case (#35081)

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."
+    ),
 }
 
 

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)
 

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():
+    if ctx.is_hybrid_swa and ctx.full_tokens_per_layer == 0:
-        return _create_unified_radix_cache(ctx, server_args, params)
+        from sglang.srt.mem_cache.pure_swa_radix_cache import PureSWARadixCache
 
-    if ctx.is_hybrid_swa:
+        return PureSWARadixCache(params=params)
-        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
 
     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 _create_unified_radix_cache(ctx, server_args, params)
-
-    return RadixCache(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 "
+            f"tree_cache is {type(cache).__name__}. Drop the flag or the "
-            "SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 (or remove "
+            "option that selected another tree cache for this model."
-            "--enable-session-radix-cache)."
         )
 
     hicache_attached = cache.cache_controller is not None

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)`

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.

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:
+            # MiMoV2 hierarchical cache runs on the unified radix tree, which
-                if not envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get():
+            # is the default tree cache now. MiMoV2 has head_dim != v_head_dim,
-                    raise ValueError(
+            # so the host KV pool uses asymmetric K/V allocation. Both
-                        "Hierarchical cache for MiMoV2 requires the unified "
+            # kernel/page_first and direct/page_first_direct have split K/V
-                        "radix tree. Set SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 "
+            # transfer paths.
-                        "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.
         elif (
             "Step3p5ForCausalLM" in model_arch
             or "Step3p7ForConditionalGeneration" in model_arch

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]:

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]:

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)

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 (
+        with patch.dict(
-            patch(
+            "sys.modules",
-                "sglang.srt.mem_cache.registry.envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get",
+            {
-                return_value=True,
+                "sglang.srt.mem_cache.unified_cache.components": fake_components,
-            ),
+                "sglang.srt.mem_cache.unified_radix_cache": fake_radix,
-            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
+        # Full attention with hierarchical cache also uses UnifiedRadixCache.
-        # `memory_pool_host`, both of which transitively load
+        fake_components = MagicMock()
-        # `sgl_kernel`; inject a stand-in module.
+        fake_radix = MagicMock()
-        fake_module = 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(
+            fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params)
-                params=ctx.params, server_args=ctx.server_args
+            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()