[Radix Cache] Add Rust TreeCore backend with shared parity tests (#32710)

Co-authored-by: alphabetc1 <2508695655@qq.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>

python/pyproject.toml

@@ -1,5 +1,11 @@
 [build-system]
-requires = ["setuptools>=61.0", "setuptools-rust>=1.10", "setuptools-scm>=8.0", "wheel"]
+requires = [
+  "setuptools>=61.0",
+  "setuptools-rust>=1.11",
+  "setuptools-scm>=8.0",
+  "torch==2.13.0",
+  "wheel",
+]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -214,6 +220,7 @@ killall_sglang = "sglang.cli.killall:main"
 "sglang" = [
   "kernels/aot/*",
   "kernels/aot/**/*",
+  "srt/mem_cache/rust_tree_core/mem_cache_inspection*.so",
 ]
 
 [tool.setuptools.packages.find]
@@ -247,8 +254,9 @@ git_describe_command = ["python3", "scripts/release/get_version_tag.py"]
 # Allow editable installs even when .git metadata is not available.
 fallback_version = "0.0.0.dev0"
 
-# Rust extension modules are auto-discovered by setup.py from the cargo
-# workspace in ../rust ([package.metadata.sglang] python-module in each crate).
+# Rust extension modules are auto-discovered by setup.py from the Cargo
+# workspace in ../rust and its declared extension manifests
+# ([package.metadata.sglang] python-module in each crate).
 # This CUDA pyproject builds all of them; platform variants restrict the set
 # via [tool.sglang] rust-extensions (see pyproject_other.toml).
 

python/setup.py

@@ -1,7 +1,8 @@
 """sglang build hooks.
 
-Rust extensions are auto-discovered from the cargo workspace in ../rust: every
-crate whose Cargo.toml declares
+Rust extensions are auto-discovered from the Cargo workspace in ../rust and
+the extension manifests declared by its workspace metadata. Every crate whose
+Cargo.toml declares
 
     [package.metadata.sglang]
     python-module = "sglang.srt.<pkg>._core"   # import path inside the wheel
@@ -28,6 +29,7 @@ Two filters can narrow the discovered set:
 import json
 import os
 import re
+import runpy
 import subprocess
 from pathlib import Path
 
@@ -45,14 +47,17 @@ except ModuleNotFoundError as exc:
 _BUILD_RUST_EXTS_ENV = "SGLANG_BUILD_RUST_EXTS"
 _PYTHON_DIR = Path(__file__).resolve().parent
 _RUST_WORKSPACE_DIR = _PYTHON_DIR.parent / "rust"
+_RUST_BUILD_HELPERS = runpy.run_path(
+    os.fspath(_PYTHON_DIR / "sglang" / "srt" / "rust_extensions" / "torch_build.py")
+)
+_torch_build_configuration = _RUST_BUILD_HELPERS["torch_build_configuration"]
 
 
-def _cargo_workspace_metadata():
-    """The rust/ cargo workspace as JSON, straight from cargo's own parser."""
-    manifest_path = _RUST_WORKSPACE_DIR / "Cargo.toml"
+def _cargo_metadata(manifest_path):
+    """One Cargo workspace/package manifest as Cargo's own JSON metadata."""
     if not manifest_path.is_file():
         raise RuntimeError(
-            f"no cargo workspace at {manifest_path} (building outside a repo "
+            f"no Cargo manifest at {manifest_path} (building outside a repo "
             f"checkout?); set {_BUILD_RUST_EXTS_ENV}=none to build without "
             "Rust extensions"
         )
@@ -83,6 +88,27 @@ def _cargo_workspace_metadata():
     return json.loads(out.stdout)
 
 
+def _cargo_workspace_metadata():
+    """Root workspace metadata plus explicitly declared extension workspaces."""
+    root_manifest = _RUST_WORKSPACE_DIR / "Cargo.toml"
+    document = _cargo_metadata(root_manifest)
+    external_manifests = (
+        (document.get("metadata") or {})
+        .get("sglang", {})
+        .get("extension-manifests", [])
+    )
+    packages = list(document["packages"])
+    for relative_manifest in external_manifests:
+        external = (_RUST_WORKSPACE_DIR / relative_manifest).resolve()
+        if _RUST_WORKSPACE_DIR not in external.parents:
+            raise RuntimeError(
+                f"external Rust extension manifest escapes rust/: {relative_manifest}"
+            )
+        packages.extend(_cargo_metadata(external)["packages"])
+    document["packages"] = packages
+    return document
+
+
 def _match_by_substring(declared, tokens, source):
     """Match tokens as case-insensitive substrings of extension names."""
     matched = set()
@@ -111,17 +137,22 @@ def _discovered_rust_extensions():
         sglang_meta = (package["metadata"] or {}).get("sglang", {})
         if "python-module" not in sglang_meta:
             continue
-        extensions.append(
-            RustExtension(
-                target=sglang_meta["python-module"],
-                path=package["manifest_path"],
-                binding=Binding.PyO3,
-                debug=sglang_meta.get("debug"),
-                # Crates that gate their PyO3 bindings behind a non-default
-                # feature (so the pure-Rust core stays pyo3-free) declare it here.
-                features=sglang_meta.get("features"),
-            )
+        extension = RustExtension(
+            target=sglang_meta["python-module"],
+            path=package["manifest_path"],
+            binding=Binding.PyO3,
+            debug=sglang_meta.get("debug"),
+            # Crates that gate their PyO3 bindings behind a non-default
+            # feature (so the pure-Rust core stays pyo3-free) declare it here.
+            features=sglang_meta.get("features"),
+            cargo_manifest_args=["--locked"],
         )
+        # Preserve Cargo metadata until the selected extension is actually
+        # built. Alternate platform pyprojects filter mem-cache out before
+        # this point and therefore do not need torch as a build dependency.
+        extension._sglang_metadata = sglang_meta
+        extension._sglang_manifest_path = package["manifest_path"]
+        extensions.append(extension)
     if not extensions:
         raise RuntimeError(
             f"no crate under {_RUST_WORKSPACE_DIR} declares "
@@ -188,6 +219,19 @@ if build_rust is not None:
     class BuildRust(build_rust):
         """Build only the Rust extensions selected by SGLANG_BUILD_RUST_EXTS."""
 
+        def run_for_extension(self, extension) -> None:
+            metadata = extension._sglang_metadata
+            compat_header = metadata.get("torch-compat-header")
+            if compat_header is not None:
+                manifest = Path(extension._sglang_manifest_path)
+                build = _torch_build_configuration(
+                    compat_header=manifest.parent / compat_header,
+                    python_module=extension.name,
+                    include_absolute_rpath=False,
+                )
+                extension.env.env = build.environment
+            super().run_for_extension(extension)
+
         def run(self) -> None:
             rust_extensions = _selected_rust_extensions(self.extensions or [])
             self.extensions = rust_extensions

python/sglang/srt/arg_groups/pd_disaggregation_hook.py

@@ -22,6 +22,7 @@ logger = logging.getLogger(__name__)
 def handle_pd_disaggregation(server_args: ServerArgs) -> None:
     """Validate and normalize PD-disaggregation server args."""
     cfg = resolving_view(server_args)
+
     # "mooncake_tcp" is mooncake with the TCP transport forced: set MC_FORCE_TCP
     # so mooncake installs TcpTransport instead of RDMA, rewrite the backend to
     # mooncake, and skip RDMA HCA selection. Must run before backend-name checks.

python/sglang/srt/disaggregation/decode_hicache_mixin.py

@@ -71,13 +71,15 @@ class DecodeHiCachePreallocMixin:
         l3_storage_hit_length = 0
         last_host_node = None
         if self.scheduler.enable_decode_hicache:
-            last_host_node = self.tree_cache.resolve_node_handle(result.last_host_node)
-            if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
+            last_host_node = result.last_host_node
+            if self.tree_cache.is_backuped(last_host_node) or self.tree_cache.is_root(
+                last_host_node
+            ):
                 matched_len = l1_prefix_len + l2_host_hit_length
                 suffix_tokens = req.origin_input_ids[matched_len:]
-                last_hash = last_host_node.get_last_hash_value()
+                last_hash = self.tree_cache.get_last_hash_value(last_host_node)
                 prefix_keys = (
-                    last_host_node.get_prefix_hash_values(last_host_node.parent)
+                    self.tree_cache.get_prefix_hash_values(last_host_node)
                     if self.tree_cache.hicache_storage_pass_prefix_keys
                     else None
                 )
@@ -112,14 +114,13 @@ class DecodeHiCachePreallocMixin:
         ):
             return
         try:
-            node = self.tree_cache.resolve_node_handle(prefix_match.last_host_node)
             matched_len = prefix_match.l1_prefix_len + prefix_match.l2_host_hit_length
             suffix = req.origin_input_ids[
                 matched_len : matched_len + prefix_match.l3_storage_hit_length
             ]
-            last_hash = node.get_last_hash_value()
+            last_hash = self.tree_cache.get_last_hash_value(prefix_match.last_host_node)
             prefix_keys = (
-                node.get_prefix_hash_values(node.parent)
+                self.tree_cache.get_prefix_hash_values(prefix_match.last_host_node)
                 if self.tree_cache.hicache_storage_pass_prefix_keys
                 else None
             )

python/sglang/srt/managers/schedule_policy.py

@@ -31,7 +31,7 @@ logger = logging.getLogger(__name__)
 
 import os
 import random
-from collections import Counter, defaultdict
+from collections import Counter
 from contextlib import contextmanager
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
@@ -397,23 +397,8 @@ class SchedulePolicy:
         waiting_queue: List[Req], tree_cache: BasePrefixCache
     ) -> None:
         """Sorts the waiting queue based on a depth-first search weighting."""
-        last_node_to_reqs = defaultdict(list)
-        for req in waiting_queue:
-            last_node = tree_cache.resolve_node_handle(req.last_node)
-            last_node_to_reqs[last_node].append(req)
-
-        node_to_weight = defaultdict(int)
-        for node in last_node_to_reqs:
-            node_to_weight[node] = len(last_node_to_reqs[node])
-        SchedulePolicy._calc_weight(tree_cache.root_node, node_to_weight)
-
-        waiting_queue.clear()
-        SchedulePolicy._get_dfs_priority(
-            tree_cache.root_node,
-            node_to_weight,
-            last_node_to_reqs,
-            waiting_queue,
-        )
+        order = tree_cache.dfs_weight_order([req.last_node for req in waiting_queue])
+        waiting_queue[:] = [waiting_queue[index] for index in order]
 
     @staticmethod
     def _sort_by_longest_output(
@@ -482,27 +467,6 @@ class SchedulePolicy:
             waiting_keys_after = [r.routing_key for r in waiting_queue]
             logger.info(f"waiting_keys_after={waiting_keys_after}")
 
-    @staticmethod
-    def _calc_weight(cur_node: TreeNode, node_to_weight: Dict[TreeNode, int]) -> None:
-        for child in cur_node.children.values():
-            SchedulePolicy._calc_weight(child, node_to_weight)
-            node_to_weight[cur_node] += node_to_weight[child]
-
-    @staticmethod
-    def _get_dfs_priority(
-        cur_node: TreeNode,
-        node_to_priority: Dict[TreeNode, int],
-        last_node_to_reqs: Dict[TreeNode, List[Req]],
-        q: List,
-    ) -> None:
-        children = [child for child in cur_node.children.values()]
-        children.sort(key=lambda x: -node_to_priority[x])
-        for child in children:
-            SchedulePolicy._get_dfs_priority(
-                child, node_to_priority, last_node_to_reqs, q
-            )
-        q.extend(last_node_to_reqs[cur_node])
-
 
 class AddReqResult(Enum):
     CONTINUE = auto()  # Continue to add requests

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     NamedTuple,
     Optional,
     Protocol,
@@ -244,6 +245,42 @@ def zero_match_result(
     )
 
 
+def _dfs_weight_order(
+    root_node: Any,
+    node_handles: Sequence[Any],
+    resolve_node_handle: Callable[[Any], Any],
+) -> list[int]:
+    last_node_to_indices: dict[Any, list[int]] = {}
+    for index, node_handle in enumerate(node_handles):
+        node = resolve_node_handle(node_handle)
+        last_node_to_indices.setdefault(node, []).append(index)
+
+    node_to_weight: dict[Any, int] = {
+        node: len(indices) for node, indices in last_node_to_indices.items()
+    }
+
+    def calc_weight(node: Any) -> None:
+        for child in node.children.values():
+            calc_weight(child)
+            node_to_weight[node] = node_to_weight.get(node, 0) + node_to_weight.get(
+                child, 0
+            )
+
+    calc_weight(root_node)
+
+    order: list[int] = []
+
+    def append_dfs(node: Any) -> None:
+        children = list(node.children.values())
+        children.sort(key=lambda child: -node_to_weight.get(child, 0))
+        for child in children:
+            append_dfs(child)
+        order.extend(last_node_to_indices.get(node, ()))
+
+    append_dfs(root_node)
+    return order
+
+
 class BasePrefixCache(ABC, PrefixCacheTrait):
     """Cache can be indexed by either rid or key."""
 
@@ -289,6 +326,10 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
     def supports_fast_match_prefix(self) -> bool:
         return False
 
+    def dfs_weight_order(self, node_handles: Sequence[Any]) -> list[int]:
+        """Return request indices in depth-first, subtree-weight order."""
+        return _dfs_weight_order(self.root_node, node_handles, self.resolve_node_handle)
+
     def resolve_node_handle(self, node_handle: Any) -> Any:
         """Map a node handle to its node -- e.g. UnifiedRadixCache looks up the
         node object from its NodeId. Temporary API for the Unified Radix Cache

python/sglang/srt/mem_cache/buffer_mode/pipeline.py

@@ -46,12 +46,13 @@ from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTr
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.unified_cache.cache_action import RebuildFullToSWAMapping
 from sglang.srt.mem_cache.unified_cache.components import (
-    BASE_COMPONENT_TYPE,
+    CacheTransferPhase,
     ComponentType,
 )
-from sglang.srt.mem_cache.unified_cache.unified_tree_core import (
+from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
+    BufferBackupSnapshot,
+    BufferBackupState,
     NodeId,
-    UnifiedTreeNode,
 )
 
 if TYPE_CHECKING:
@@ -65,16 +66,12 @@ class _UnifiedBackupIntent(msgspec.Struct):
     """Buffer-mode backup intent, unpinned while queued.
 
     Snapshots node identity at enqueue time: a split rewrites the node's
-    key/hash in place while these copies stay intact, so
-    ``node.hash_value != hash_values`` doubles as split detection and a None
-    FULL device value as eviction detection (``_backup_intent_stale``).
+    key/hash in place while these copies stay intact, so a key-length change
+    detects a split and a missing FULL device value detects eviction
+    (``_validate_backup_intent``).
     """
 
-    node: UnifiedTreeNode
-    node_id: int
-    hash_values: list[str]
-    key: RadixKey
-    prefix_keys: Optional[list[str]] = None
+    snapshot: BufferBackupSnapshot
 
 
 class _UnifiedBufferBackupEntry(msgspec.Struct):
@@ -289,18 +286,19 @@ class BufferModePipeline:
 
     # ---- backup pipeline (device -> staging -> storage) ----
 
-    def _backup_parent_covered(self, node: UnifiedTreeNode) -> bool:
+    def _backup_parent_covered(self, state: BufferBackupState) -> bool:
         """Only admit a node whose parent is stored/in-flight: writing above
         a dropped parent creates a permanent longest-prefix hole."""
-        parent = node.parent
         if (
-            parent is self._cache.root_node
-            or parent.id in self.inflight_backup_node_ids
+            state.parent_is_root
+            or state.parent_node_id in self.inflight_backup_node_ids
         ):
             return True
-        last_hash = parent.get_last_hash_value()
-        return last_hash is not None and self._cache.storage_existence_cache.contains(
-            PoolName.KV, last_hash
+        return (
+            state.parent_last_hash is not None
+            and self._cache.storage_existence_cache.contains(
+                PoolName.KV, state.parent_last_hash
+            )
         )
 
     def _log_backup_dropped(self, num_tokens: int) -> None:
@@ -308,22 +306,29 @@ class BufferModePipeline:
         if cache.enable_storage_metrics and cache.storage_metrics_collector is not None:
             cache.storage_metrics_collector.log_backup_dropped_tokens(num_tokens)
 
-    def enqueue_backup_intent(self, node: UnifiedTreeNode) -> None:
+    def enqueue_backup_intent(self, node_id: NodeId) -> None:
         """Snapshot a backup intent and commit it to the write queue.
         Admission gates: belief skip, parent-cover, backlog cap, oversize.
         Drops are silent; the node re-triggers on a later hit."""
-        if not self._cache.enable_storage or not node.hash_value:
+        if not self._cache.enable_storage:
             return
-        if node.id in self.inflight_backup_node_ids:
+        if node_id in self.inflight_backup_node_ids:
+            return
+        snapshot = self._cache.tree_core.snapshot_buffer_backup(
+            node_id, self._cache.hicache_storage_pass_prefix_keys
+        )
+        if snapshot is None:
             return
         # Admission cover: beliefs plus content past its D2H launch. The
         # launched cover keeps republished content (fill inserts under new
         # node ids) from re-writing while the original write drains.
         if self._cache.storage_existence_cache.covers_all(
-            PoolName.KV, node.hash_value, extra_cover=self.inflight_backup_hashes
+            PoolName.KV,
+            snapshot.hash_values,
+            extra_cover=self.inflight_backup_hashes,
         ):
             return
-        intent_tokens = len(node.hash_value) * self._cache.page_size
+        intent_tokens = len(snapshot.hash_values) * self._cache.page_size
         if self.write_backlog_tokens_ >= self.write_backlog_cap:
             # The cap sits at 2x the intrinsic live-backlog ceiling (see
             # init_hicache), so reaching it means leaked accounting or a
@@ -344,51 +349,59 @@ class BufferModePipeline:
             return
         # A span larger than any pool's whole staging capacity can never
         # stage; admitting it would wedge the head-of-line queue forever.
-        if not self._backup_parent_covered(node) or self._backup_oversize(
-            node, intent_tokens
+        state = BufferBackupState(
+            parent_node_id=snapshot.parent_node_id,
+            parent_is_root=snapshot.parent_is_root,
+            parent_last_hash=snapshot.parent_last_hash,
+        )
+        if not self._backup_parent_covered(state) or self._backup_oversize(
+            snapshot.node_id, snapshot.hash_values, intent_tokens
         ):
             self._log_backup_dropped(intent_tokens)
             return
 
-        prefix_keys = (
-            node.get_prefix_hash_values(node.parent)
-            if self._cache.hicache_storage_pass_prefix_keys
-            else None
-        )
-        intent = _UnifiedBackupIntent(
-            node=node,
-            node_id=node.id,
-            hash_values=list(node.hash_value),
-            key=node.key,
-            prefix_keys=prefix_keys,
-        )
+        intent = _UnifiedBackupIntent(snapshot=snapshot)
         self.pending_write_queue.append(intent)
-        self.inflight_backup_node_ids.add(node.id)
+        self.inflight_backup_node_ids.add(snapshot.node_id)
         self.write_backlog_tokens_ += intent_tokens
 
     def _build_aux_staging_transfers(
-        self, node: UnifiedTreeNode
-    ) -> Optional[list[PoolTransfer]]:
+        self,
+        node_id: NodeId,
+        hash_values: list[str],
+        comp_xfers: Optional[dict[ComponentType, list[PoolTransfer]]] = None,
+    ) -> list[PoolTransfer]:
         """Keys-only aux transfers mirroring what BACKUP_STORAGE would write;
         sizes the per-pool oversize gate (beliefs do not consult these)."""
         transfers: list[PoolTransfer] = []
         if ComponentType.SWA in self._cache.components:
-            cd = node.component_data[ComponentType.SWA]
-            if cd.value is not None:
-                num_pages = len(cd.value) // self._cache.page_size
+            current = (
+                comp_xfers.get(ComponentType.SWA)
+                if comp_xfers is not None
+                else self._cache.tree_core.build_hicache_transfers(
+                    ComponentType.SWA,
+                    node_id,
+                    CacheTransferPhase.BACKUP_HOST,
+                )
+            )
+            for transfer in current or ():
+                if transfer.device_indices is None:
+                    continue
+                num_pages = len(transfer.device_indices) // self._cache.page_size
                 if num_pages > 0:
                     transfers.append(
                         PoolTransfer(
                             name=PoolName.SWA,
-                            keys=node.hash_value[-num_pages:],
+                            keys=hash_values[-num_pages:],
                             hit_policy=PoolHitPolicy.TRAILING_PAGES,
                         )
                     )
-        return transfers or None
+        return transfers
 
     def _backup_oversize(
         self,
-        node: UnifiedTreeNode,
+        node_id: NodeId,
+        hash_values: list[str],
         intent_tokens: int,
         aux_xfers: Optional[list[PoolTransfer]] = None,
     ) -> bool:
@@ -400,7 +413,7 @@ class BufferModePipeline:
         if intent_tokens > cc.mem_pool_host.size:
             return True
         if aux_xfers is None:
-            aux_xfers = self._build_aux_staging_transfers(node)
+            aux_xfers = self._build_aux_staging_transfers(node_id, hash_values)
         for t in aux_xfers or ():
             entry = cc.mem_pool_host.entry_map.get(t.name)
             if entry is not None and (
@@ -420,35 +433,37 @@ class BufferModePipeline:
             host_pool.size // 10,
         )
 
-    def _backup_intent_stale(self, intent: _UnifiedBackupIntent) -> bool:
-        # Arena-lookup failure = deleted, hash mismatch vs the enqueue-time
-        # snapshot = split, a None FULL device value = evicted. Stale
+    def _validate_backup_intent(
+        self, intent: _UnifiedBackupIntent
+    ) -> Optional[BufferBackupState]:
+        # Arena-lookup failure = deleted, key-length mismatch vs the snapshot
+        # = split, a None FULL device value = evicted. Stale
         # intents drop silently; the node re-triggers on a later hit.
-        node = intent.node
-        try:
-            self._cache.tree_core.node_by_id(intent.node_id)
-        except KeyError:
-            return True
-        return (
-            node.component_data[BASE_COMPONENT_TYPE].value is None
-            or node.hash_value != intent.hash_values
+        snapshot = intent.snapshot
+        return self._cache.tree_core.validate_buffer_backup(
+            snapshot.node_id, len(snapshot.key)
         )
 
-    def _sweep_stale_backup_intents(self) -> None:
+    def _sweep_stale_backup_intents(self) -> dict[NodeId, BufferBackupState]:
         """Cancel stale intents anywhere in the queue, not just at the head:
         a dead intent would otherwise inflate the backlog accounting and
         hold FIFO position ahead of live segments."""
         if not self.pending_write_queue:
-            return
+            return {}
         page_size = self._cache.page_size
         survivors: deque[_UnifiedBackupIntent] = deque()
+        states: dict[NodeId, BufferBackupState] = {}
         for intent in self.pending_write_queue:
-            if self._backup_intent_stale(intent):
-                self.inflight_backup_node_ids.discard(intent.node_id)
-                self.write_backlog_tokens_ -= len(intent.hash_values) * page_size
+            snapshot = intent.snapshot
+            state = self._validate_backup_intent(intent)
+            if state is None:
+                self.inflight_backup_node_ids.discard(snapshot.node_id)
+                self.write_backlog_tokens_ -= len(snapshot.hash_values) * page_size
                 continue
             survivors.append(intent)
+            states[snapshot.node_id] = state
         self.pending_write_queue = survivors
+        return states
 
     def flush_pending_writes(self) -> None:
         """Launch D2H transfers for admitted intents, head-of-line: device
@@ -456,7 +471,7 @@ class BufferModePipeline:
         if not self.pending_write_queue:
             return
         cc = self._cache.cache_controller
-        self._sweep_stale_backup_intents()
+        states = self._sweep_stale_backup_intents()
         # Loads have priority (writes are deferrable): the write window is
         # the pool minus prefetch occupancy minus a 10% margin, floored at
         # the configured fraction.
@@ -467,34 +482,56 @@ class BufferModePipeline:
         )
         while self.pending_write_queue:
             intent = self.pending_write_queue[0]
-            intent_tokens = len(intent.hash_values) * self._cache.page_size
-            if not self._backup_parent_covered(intent.node) or self._backup_oversize(
-                intent.node, intent_tokens
-            ):
-                # Unwritable intent (dropped parent or unstageable size):
-                # cascade the drop down the chain rather than creating a
-                # permanent storage hole / stalling the head-of-line queue.
+            snapshot = intent.snapshot
+            state = states[snapshot.node_id]
+            intent_tokens = len(snapshot.hash_values) * self._cache.page_size
+            if not self._backup_parent_covered(state):
+                # Cascade a dropped parent down the chain rather than creating
+                # a permanent storage hole.
                 self.pending_write_queue.popleft()
-                self.inflight_backup_node_ids.discard(intent.node_id)
+                self.inflight_backup_node_ids.discard(snapshot.node_id)
                 self.write_backlog_tokens_ -= intent_tokens
                 self._log_backup_dropped(intent_tokens)
                 continue
             if self.write_staged_tokens_ >= live_cap:
                 # Yield to live fetch demand; retry next round.
                 break
-            if self._aux_budget_blocked(intent):
+            device_value, comp_xfers = self._cache.tree_core.build_backup_spec(
+                snapshot.node_id
+            )
+            sizing_xfers = self._build_aux_staging_transfers(
+                snapshot.node_id, snapshot.hash_values, comp_xfers
+            )
+            if self._backup_oversize(
+                snapshot.node_id,
+                snapshot.hash_values,
+                intent_tokens,
+                sizing_xfers,
+            ):
+                # A permanently unstageable head must not block the queue.
+                self.pending_write_queue.popleft()
+                self.inflight_backup_node_ids.discard(snapshot.node_id)
+                self.write_backlog_tokens_ -= intent_tokens
+                self._log_backup_dropped(intent_tokens)
+                continue
+            if self._aux_budget_blocked(intent, sizing_xfers):
                 # An aux pool lacks staging headroom: yield at the gate
                 # instead of failing the alloc inside cc.write; acks free
                 # aux staging, retry next round.
                 break
-            if not self._launch_backup_intent(intent):
+            if not self._launch_backup_intent(intent, device_value, comp_xfers):
                 # Pool full of in-flight staging and nothing reclaimable
                 # (the tree never holds host values in buffer mode):
                 # defer, head-of-line; pending acks will free slots.
                 break
             self.pending_write_queue.popleft()
 
-    def _launch_backup_intent(self, intent: _UnifiedBackupIntent) -> bool:
+    def _launch_backup_intent(
+        self,
+        intent: _UnifiedBackupIntent,
+        device_value: torch.Tensor,
+        comp_xfers: dict[ComponentType, list[PoolTransfer]],
+    ) -> bool:
         """Launch one admitted intent's D2H (staging alloc + device lock +
         async copy); the caller removes it from pending_write_queue. Returns
         False when staging cannot be allocated. From a successful launch the
@@ -502,33 +539,34 @@ class BufferModePipeline:
         LAUNCHED cover consulted by admission."""
         cache = self._cache
         cc = cache.cache_controller
-        node = intent.node
-        # Build aux transfers from the node's CURRENT state: a SWA span
-        # tombstoned since admission backs up FULL-only, as in cache mode.
-        device_value, comp_xfers = cache.tree_core.build_backup_spec(node.id)
+        snapshot = intent.snapshot
         aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
         host_indices = cc.write(
             device_value,
-            node_id=node.id,
+            node_id=snapshot.node_id,
             extra_pools=aux_xfers or None,
         )
         if host_indices is None:
             return False
-        _track_content_refs(self.inflight_backup_hashes, intent.hash_values)
+        _track_content_refs(self.inflight_backup_hashes, snapshot.hash_values)
         # NOTE: no commit_backup — the node must never appear
         # host-resident in buffer mode; staging slots live in the entry.
-        lock_params = cache.inc_lock_ref(node.id).to_dec_params()
-        self.ongoing_write_through[node.id] = _UnifiedBufferBackupEntry(
+        lock_params = cache.inc_lock_ref(snapshot.node_id).to_dec_params()
+        self.ongoing_write_through[snapshot.node_id] = _UnifiedBufferBackupEntry(
             intent=intent,
             host_indices=host_indices,
             aux_xfers=aux_xfers,
             lock_params=lock_params,
         )
         self.write_staged_tokens_ += len(host_indices)
-        self.write_backlog_tokens_ -= len(intent.hash_values) * cache.page_size
+        self.write_backlog_tokens_ -= len(snapshot.hash_values) * cache.page_size
         return True
 
-    def _aux_budget_blocked(self, intent: _UnifiedBackupIntent) -> bool:
+    def _aux_budget_blocked(
+        self,
+        intent: _UnifiedBackupIntent,
+        aux: Optional[list[PoolTransfer]] = None,
+    ) -> bool:
         """True when an aux pool cannot stage this intent right now (free
         minus the loads-priority margin falls short of the need): defer at
         the gate instead of failing the alloc inside cc.write and blocking
@@ -536,7 +574,11 @@ class BufferModePipeline:
         loads-have-priority on aux pools the way live_cap does on the KV
         pool; avail already reflects prefetch-held slots, so no occupancy
         subtraction here."""
-        aux = self._build_aux_staging_transfers(intent.node)
+        snapshot = intent.snapshot
+        if aux is None:
+            aux = self._build_aux_staging_transfers(
+                snapshot.node_id, snapshot.hash_values
+            )
         if not aux:
             return False
         cc = self._cache.cache_controller
@@ -574,14 +616,15 @@ class BufferModePipeline:
         (which reads from the staging copy, so device eviction may proceed)."""
         entry = self.ongoing_write_through.pop(ack_id)
         intent = entry.intent
-        self._cache.dec_lock_ref(intent.node_id, entry.lock_params)
+        snapshot = intent.snapshot
+        self._cache.dec_lock_ref(snapshot.node_id, entry.lock_params)
 
         # Every aux pool writes a trailing snapshot keyed by the last KV page
         # hashes it covers: the SWA window spans page_size-sized pages, the
         # Mamba state is a single slot (host pool page_size 1 -> one key).
         storage_xfers: list[PoolTransfer] = []
         for staged in entry.aux_xfers:
-            keys = self._aux_window_keys(intent.hash_values, staged)
+            keys = self._aux_window_keys(snapshot.hash_values, staged)
             if keys is None:
                 continue
             storage_xfers.append(
@@ -594,9 +637,9 @@ class BufferModePipeline:
             )
         operation_id = self._cache.cache_controller.write_storage(
             entry.host_indices,
-            intent.key.token_ids,
-            intent.hash_values,
-            intent.prefix_keys,
+            snapshot.key.token_ids,
+            snapshot.hash_values,
+            snapshot.prefix_keys,
             extra_pools=storage_xfers or None,
         )
         self.ongoing_backup[operation_id] = entry
@@ -611,11 +654,12 @@ class BufferModePipeline:
         if entry is None:
             return
         intent = entry.intent
-        self._cache.storage_existence_cache.add(PoolName.KV, intent.hash_values)
+        snapshot = intent.snapshot
+        self._cache.storage_existence_cache.add(PoolName.KV, snapshot.hash_values)
         self._free_staging_now(entry.host_indices, entry.aux_xfers)
         self.write_staged_tokens_ -= len(entry.host_indices)
-        self.inflight_backup_node_ids.discard(entry.intent.node_id)
-        _untrack_content_refs(self.inflight_backup_hashes, intent.hash_values)
+        self.inflight_backup_node_ids.discard(snapshot.node_id)
+        _untrack_content_refs(self.inflight_backup_hashes, snapshot.hash_values)
 
     def _free_staging_now(
         self, host_indices: torch.Tensor, aux_xfers: list[PoolTransfer]
@@ -671,10 +715,18 @@ class BufferModePipeline:
                 )
             return "cap_skip"
         cache = self._cache
+        anchor_tokens = array("q", prefix_tokens)
+        if cache.tree_core.is_eagle:
+            # The suffix owns the boundary token shared with the last matched
+            # bigram, so include it when rebuilding the anchor key.
+            info = cache.ongoing_prefetch.get(req_id)
+            if info is None or not info.prefetch_key.token_ids:
+                return "anchor_lost"
+            anchor_tokens.append(info.prefetch_key.token_ids[0])
         match = cache.match_prefix(
             MatchPrefixParams(
                 key=RadixKey(
-                    array("q", prefix_tokens),
+                    anchor_tokens,
                     extra_key=extra_key,
                     is_bigram=cache.tree_core.is_eagle,
                     cache_salt=cache_salt,

python/sglang/srt/mem_cache/rust_tree_core/.gitignore

@@ -0,0 +1 @@
+mem_cache.so

python/sglang/srt/mem_cache/rust_tree_core/__init__.py

@@ -0,0 +1 @@
+"""The in-tree Rust TreeCore backend; the factory lives in tree_core_registry."""

python/sglang/srt/mem_cache/rust_tree_core/adapter.py

@@ -0,0 +1,971 @@
+"""The Rust TreeCore adapter: satisfies ``UnifiedTreeCoreInterface`` over the
+``mem_cache`` extension's ``RustUnifiedTreeCoreBinding``."""
+
+from __future__ import annotations
+
+from array import array
+from typing import TYPE_CHECKING, Optional, Sequence
+
+import torch
+
+from sglang.srt.disaggregation.kv_events import (
+    AllBlocksCleared,
+    BlockRemoved,
+    BlockStored,
+    BlockStoredMetadata,
+    BlockStoredWithMetadata,
+    StorageMedium,
+)
+from sglang.srt.mem_cache.base_prefix_cache import (
+    DecLockRefParams,
+    DecLockRefResult,
+    IncLockRefResult,
+    InsertParams,
+    InsertResult,
+    MatchPrefixParams,
+    MatchResult,
+)
+from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer
+from sglang.srt.mem_cache.radix_cache import RadixKey
+from sglang.srt.mem_cache.rust_tree_core.extension import bindings
+from sglang.srt.mem_cache.unified_cache.cache_action import (
+    BackupKV,
+    FreeComponentDeviceSlot,
+    FreeComponentHostSlot,
+    FreeDeviceKV,
+    FreeDeviceKVFullOnly,
+    MambaEvictExcessPathStates,
+    RebuildFullToSWAMapping,
+    RecoverSWAWithLockedFull,
+    ReplaceWriteThroughOnNodeSplit,
+    SWARebuild,
+)
+from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
+from sglang.srt.mem_cache.unified_cache.unified_tree_core import StorageBackupSpec
+from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
+    BufferBackupSnapshot,
+    BufferBackupState,
+    DecSwaLockOnlyResult,
+    DemoteResult,
+    DriveHostEvictionResult,
+    DropSubtreeNoHostResult,
+    EvictDeviceLeafResult,
+    EvictDeviceNextNodeResult,
+    InsertStepResult,
+    NodeId,
+    RadixCacheWalkResult,
+    UnifiedTreeCoreInterface,
+)
+from sglang.srt.runtime_context import get_exec, mamba_cache_chunk_size
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+    from sglang.srt.mem_cache.hicache_storage import PoolTransferResult
+    from sglang.srt.mem_cache.unified_cache.cache_action import (
+        CacheAction,
+        ComponentAction,
+    )
+    from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase
+    from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeNode
+
+
+def _radix_key_buffer(key: RadixKey) -> array:
+    """The key's token ids honoring `limit`; view-independent since the
+    binding derives its own atoms."""
+    token_ids = key.raw_token_ids()
+    assert (
+        isinstance(token_ids, array) and token_ids.typecode == "q"
+    ), f"tree keys must carry array('q') token ids, got {type(token_ids).__name__}"
+    return token_ids
+
+
+def _kv_event_from_tagged(event: tuple):
+    """Build the Python KV cache event for one of the binding's tagged tuples."""
+    tag = event[0]
+    if tag == "block_stored":
+        event_args = dict(
+            block_hashes=event[1],
+            parent_block_hash=event[2],
+            token_ids=event[3],
+            block_size=event[4],
+            lora_id=None,
+            medium=StorageMedium(event[5]),
+        )
+        if event[6] is None:
+            return BlockStored(**event_args)
+        return BlockStoredWithMetadata(
+            **event_args,
+            metadata=BlockStoredMetadata(cache_salt=event[6]),
+        )
+    if tag == "block_removed":
+        return BlockRemoved(block_hashes=event[1], medium=StorageMedium(event[2]))
+    if tag == "all_blocks_cleared":
+        return AllBlocksCleared()
+    raise ValueError(f"unknown kv event tag: {tag}")
+
+
+def _cache_action_from_tagged(action: tuple) -> CacheAction:
+    """Build the Python CacheAction for one of the binding's tagged tuples."""
+    tag = action[0]
+    if tag == "free_device_kv":
+        return FreeDeviceKV(indices=list(action[1]))
+    if tag == "free_device_kv_full_only":
+        return FreeDeviceKVFullOnly(indices=list(action[1]))
+    if tag == "backup_kv":
+        return BackupKV(node_ids=list(action[1]))
+    if tag == "mamba_evict_excess_path_states":
+        return MambaEvictExcessPathStates(tail_node_id=action[1])
+    if tag == "replace_write_through_on_node_split":
+        return ReplaceWriteThroughOnNodeSplit(
+            ack_id=action[1],
+            old_node_id=action[2],
+            new_node_id=action[3],
+            new_child_node_id=action[4],
+        )
+    if tag == "free_component_device_slot":
+        return FreeComponentDeviceSlot(
+            component_type=ComponentType(action[1]), indices=list(action[2])
+        )
+    if tag == "free_component_host_slot":
+        return FreeComponentHostSlot(
+            component_type=ComponentType(action[1]), host_indices=list(action[2])
+        )
+    if tag == "rebuild_full_to_swa_mapping":
+        return RebuildFullToSWAMapping(
+            full_indices=list(action[1]), swa_indices=list(action[2])
+        )
+    if tag == "recover_swa_with_locked_full":
+        return RecoverSWAWithLockedFull(
+            node_id=action[1], kept_full=action[2], incoming_full=action[3]
+        )
+    if tag == "swa_rebuild":
+        return SWARebuild(node_id=action[1], source_value=action[2])
+    raise ValueError(f"unknown cache action tag: {tag}")
+
+
+def _cache_actions_from_tagged(actions: Sequence[tuple]) -> list[CacheAction]:
+    """Build the Python CacheActions for the binding's tagged tuples, in order."""
+    return [_cache_action_from_tagged(action) for action in actions]
+
+
+def _inc_lock_ref_result_from_binding(result) -> IncLockRefResult:
+    return IncLockRefResult(
+        delta=result.delta,
+        swa_uuid_for_lock=result.swa_uuid_for_lock,
+        swa_uuid_for_host_lock=result.swa_uuid_for_host_lock,
+        skip_lock_node_ids=_skip_lock_node_ids_from_binding(result.skip_lock_node_ids),
+    )
+
+
+def _transfer_to_binding(transfer: PoolTransfer) -> tuple:
+    """The binding's (name, host_indices, device_indices, nodes_to_load, keys,
+    hit_policy) tuple."""
+    return (
+        transfer.name.value,
+        transfer.host_indices,
+        transfer.device_indices,
+        transfer.nodes_to_load,
+        transfer.keys,
+        transfer.hit_policy.value,
+    )
+
+
+def _transfer_from_binding(transfer: tuple) -> PoolTransfer:
+    """Build the Python PoolTransfer for one of the binding's transfer tuples."""
+    name, host_indices, device_indices, nodes_to_load, keys, hit_policy = transfer
+    return PoolTransfer(
+        name=PoolName(name),
+        host_indices=host_indices,
+        device_indices=device_indices,
+        keys=keys,
+        hit_policy=PoolHitPolicy(hit_policy),
+        nodes_to_load=nodes_to_load,
+    )
+
+
+def _comp_xfers_to_binding(
+    comp_xfers: dict[ComponentType, list[PoolTransfer]],
+) -> dict[int, list[tuple]]:
+    """Rekey per-component transfers by the binding's component values."""
+    return {
+        int(ct): [_transfer_to_binding(x) for x in xfers]
+        for ct, xfers in comp_xfers.items()
+    }
+
+
+def _comp_xfers_from_binding(
+    comp_xfers: dict[int, list[tuple]],
+) -> dict[ComponentType, list[PoolTransfer]]:
+    """Rekey the binding's per-component transfer tuples by ComponentType."""
+    return {
+        ComponentType(ct): [_transfer_from_binding(x) for x in xfers]
+        for ct, xfers in comp_xfers.items()
+    }
+
+
+def _insert_step_from_binding(step) -> InsertStepResult:
+    """Build the interface step for the binding's step (result on the final one)."""
+    result = None
+    if step.result is not None:
+        # A stepped insert delivers all actions through steps, never the result.
+        assert not step.result.cache_actions
+        result = InsertResult(
+            prefix_len=step.result.prefix_len,
+            last_device_node=step.result.last_device_node,
+            mamba_exist=step.result.mamba_exist,
+            host_insert_dropped=step.result.host_insert_dropped,
+            adopted_ranges=(
+                {
+                    ComponentType(component_type): list(ranges)
+                    for component_type, ranges in step.result.adopted_ranges.items()
+                }
+                if step.result.adopted_ranges is not None
+                else None
+            ),
+        )
+    return InsertStepResult(
+        actions=_cache_actions_from_tagged(step.actions), result=result
+    )
+
+
+def _match_result_from_binding(result) -> MatchResult:
+    """Build the Python MatchResult for the binding's match result."""
+    return MatchResult(
+        device_indices=result.device_indices,
+        last_device_node=result.last_device_node_id,
+        last_host_node=result.last_host_node_id,
+        best_match_node=result.best_match_node_id,
+        host_hit_length=result.host_hit_length,
+        swa_host_hit_length=result.swa_host_hit_length,
+        mamba_host_hit_length=result.mamba_host_hit_length,
+        mamba_branching_seqlen=result.mamba_branching_seqlen,
+        full_kv_hit_length=result.full_kv_hit_length,
+        cache_actions=_cache_actions_from_tagged(result.cache_actions),
+    )
+
+
+def _skip_lock_node_ids_from_binding(
+    skip_lock_node_ids: dict[int, set[int]],
+) -> dict[ComponentType, set[int]]:
+    """Rekey the binding's component-value skip map by ComponentType."""
+    return {
+        ComponentType(component): set(node_ids)
+        for component, node_ids in skip_lock_node_ids.items()
+    }
+
+
+def _skip_lock_node_ids_to_binding(
+    skip_lock_node_ids: dict[ComponentType, set[int]],
+) -> dict[int, set[int]]:
+    """Rekey a ComponentType skip map by the binding's component values."""
+    return {
+        int(component): set(node_ids)
+        for component, node_ids in skip_lock_node_ids.items()
+    }
+
+
+def _tracker_to_binding(tracker: dict[ComponentType, int]) -> dict[int, int]:
+    """Rekey a ComponentType tracker by the binding's component values."""
+    return {int(component): freed for component, freed in tracker.items()}
+
+
+def _fill_evict_result(binding_result, result):
+    """Map a binding eviction step into an interface step result; both carry
+    this step's per-component deltas and freed tensors."""
+    for component, delta in binding_result.tracker.items():
+        result.tracker[ComponentType(component)] = delta
+    for component, tensors in binding_result.new_device_frees.items():
+        result.device_frees[ComponentType(component)].extend(tensors)
+    for component, tensors in binding_result.new_host_frees.items():
+        result.host_frees[ComponentType(component)].extend(tensors)
+    return result
+
+
+class _RustKVCacheEventRecorder:
+    """Expose the Rust event queue through the Python recorder interface."""
+
+    def __init__(self, binding, enabled: bool):
+        self._binding = binding
+        self.enabled = enabled
+
+    def record_all_cleared(self) -> None:
+        self._binding.record_all_cleared_event()
+
+    def take(self) -> list:
+        return [_kv_event_from_tagged(event) for event in self._binding.take_events()]
+
+
+class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
+    """A TreeCore backed by the Rust extension binding."""
+
+    _bindings = bindings
+
+    def __init__(self, params: CacheInitParams):
+        assert params.tree_components is not None
+        self.tree_components = tuple(params.tree_components)
+
+        # TODO(Jialin): Port session-reference-aware TreeCore support from #29173.
+        if params.enable_session_radix_cache:
+            raise ValueError(
+                "--enable-session-radix-cache is not supported by the Rust TreeCore"
+            )
+
+        # TODO(Jialin): Port custom component registration from #25754 and
+        # C128 support from #33676.
+        unsupported_components = set(self.tree_components) - {
+            ComponentType.FULL,
+            ComponentType.SWA,
+            ComponentType.MAMBA,
+        }
+        if unsupported_components:
+            names = ", ".join(
+                sorted(component.name for component in unsupported_components)
+            )
+            raise ValueError(f"Rust TreeCore does not support components: {names}")
+        if params.component_registry_override:
+            raise ValueError(
+                "Rust TreeCore does not support component_registry_override"
+            )
+
+        self._page_size = params.page_size
+        self.is_eagle = (
+            params.is_eagle and ComponentType.MAMBA not in self.tree_components
+        )
+
+        # ``device`` is derived from the construction-time allocator; the
+        # allocator/pool themselves are owned by the cache, not the tree.
+        if params.token_to_kv_pool_allocator:
+            device = torch.device(params.token_to_kv_pool_allocator.device)
+            # A bare "cuda" means the process's current device, not cuda:0.
+            if device.type == "cuda" and device.index is None:
+                device = torch.device("cuda", torch.cuda.current_device())
+            self.device = device
+        else:
+            self.device = torch.device("cpu")
+
+        self.enable_kv_cache_events = params.enable_kv_cache_events
+        has_mamba = ComponentType.MAMBA in self.tree_components
+        mamba_max_states_per_path = (
+            get_exec().mamba.mamba_max_states_per_path if has_mamba else -1
+        )
+
+        self._binding = self._binding_class()(
+            self._bindings.TreeCoreInitParamsBinding(
+                eviction_policy=params.eviction_policy,
+                page_size=params.page_size,
+                is_write_back=False,
+                enable_hicache=False,
+                write_through_threshold=256,
+                device=str(self.device),
+                swa_sliding_window_size=params.sliding_window_size,
+                enable_kv_cache_events=params.enable_kv_cache_events,
+                mamba_cache_chunk_size=(
+                    mamba_cache_chunk_size() if has_mamba else None
+                ),
+                mamba_max_states_per_path=(
+                    mamba_max_states_per_path
+                    if mamba_max_states_per_path >= 0
+                    else None
+                ),
+            ),
+            [int(component) for component in self.tree_components],
+        )
+        self.kv_events = _RustKVCacheEventRecorder(
+            self._binding, params.enable_kv_cache_events
+        )
+        # The default-root empty result, prebuilt once from the binding.
+        self._empty_match_result = _match_result_from_binding(
+            self._binding.empty_match_result()
+        )
+
+    def _binding_class(self) -> type:
+        """The extension binding class this core constructs."""
+        if self.is_eagle:
+            return self._bindings.RustBigramUnifiedTreeCoreBinding
+        return self._bindings.RustUnifiedTreeCoreBinding
+
+    # ==== Tree API ====
+
+    def reset(self) -> None:
+        self._binding.reset()
+        # Node handles are never re-minted, so the fresh root gets a new one.
+        self._empty_match_result = _match_result_from_binding(
+            self._binding.empty_match_result()
+        )
+
+    def node_by_id(self, node_id: NodeId) -> UnifiedTreeNode:
+        # TODO(Jialin): Move the remaining Python-node consumers to
+        # backend-neutral APIs: sessions (#29173), C128 (#33676).
+        raise NotImplementedError("node_by_id: not yet ported to the Rust tree core")
+
+    @property
+    def root_node(self) -> UnifiedTreeNode:
+        raise NotImplementedError("root_node: not yet ported to the Rust tree core")
+
+    def inc_lock_ref(
+        self,
+        node_id: NodeId,
+        skip_lock_components: Sequence[ComponentType] = (),
+    ) -> IncLockRefResult:
+        result = self._binding.inc_lock_ref(
+            node_id, [int(component) for component in skip_lock_components]
+        )
+        return _inc_lock_ref_result_from_binding(result)
+
+    def dec_lock_ref(
+        self,
+        node_id: NodeId,
+        params: Optional[DecLockRefParams] = None,
+        skip_swa: bool = False,
+    ) -> DecLockRefResult:
+        binding_params = (
+            self._bindings.DecLockRefParamsBinding(
+                swa_uuid_for_lock=params.swa_uuid_for_lock,
+                swa_uuid_for_host_lock=params.swa_uuid_for_host_lock,
+                skip_lock_node_ids=_skip_lock_node_ids_to_binding(
+                    params.skip_lock_node_ids
+                ),
+            )
+            if params is not None
+            else None
+        )
+        self._binding.dec_lock_ref(node_id, binding_params, skip_swa)
+        return DecLockRefResult()
+
+    def dec_swa_lock_only(
+        self,
+        node_id: NodeId,
+        swa_uuid_for_lock: Optional[int],
+        skip_lock_node_ids: Optional[dict] = None,
+    ) -> DecSwaLockOnlyResult:
+        result = DecSwaLockOnlyResult()
+        new_device_frees, new_host_frees = self._binding.dec_swa_lock_only(
+            node_id,
+            swa_uuid_for_lock,
+            (
+                _skip_lock_node_ids_to_binding(skip_lock_node_ids)
+                if skip_lock_node_ids
+                else None
+            ),
+        )
+        for component, tensors in new_device_frees.items():
+            result.device_frees[ComponentType(component)].extend(tensors)
+        for component, tensors in new_host_frees.items():
+            result.host_frees[ComponentType(component)].extend(tensors)
+        return result
+
+    # ==== Device eviction (driven step-wise by the Controller's evict()) ====
+
+    def evict_device_start(
+        self, component_type: ComponentType, request_cnt: int
+    ) -> None:
+        self._binding.evict_device_start(int(component_type), request_cnt)
+
+    def evict_device_next_node(
+        self, component_type: ComponentType, tracker: dict[ComponentType, int]
+    ) -> EvictDeviceNextNodeResult:
+        binding_result = self._binding.evict_device_next_node(
+            int(component_type), _tracker_to_binding(tracker)
+        )
+        result = EvictDeviceNextNodeResult(
+            node_id=binding_result.node_id,
+            made_progress=binding_result.made_progress,
+        )
+        return _fill_evict_result(binding_result, result)
+
+    def evict_device_leaf(
+        self, node_id: NodeId, is_write_back: bool
+    ) -> EvictDeviceLeafResult:
+        # The binding reads is_write_back from the core's construction config.
+        assert (
+            is_write_back == self.is_write_back
+        ), "is_write_back must match the core's construction config"
+        binding_result = self._binding.evict_device_leaf(node_id)
+        backup = binding_result.backup_kv
+        result = EvictDeviceLeafResult(
+            backup_kv=_cache_action_from_tagged(backup) if backup is not None else None
+        )
+        return _fill_evict_result(binding_result, result)
+
+    def demote(self, node_id: NodeId) -> DemoteResult:
+        binding_result = self._binding.demote(node_id)
+        return _fill_evict_result(binding_result, DemoteResult())
+
+    def evict_device_end(self, component_type: ComponentType) -> None:
+        self._binding.evict_device_end(int(component_type))
+
+    def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
+        result = self._binding.inc_host_lock_ref(node_id)
+        return IncLockRefResult(
+            delta=result.delta,
+            swa_uuid_for_lock=result.swa_uuid_for_lock,
+            swa_uuid_for_host_lock=result.swa_uuid_for_host_lock,
+            skip_lock_node_ids=_skip_lock_node_ids_from_binding(
+                result.skip_lock_node_ids
+            ),
+        )
+
+    def dec_host_lock_ref(
+        self, node_id: NodeId, params: Optional[DecLockRefParams] = None
+    ) -> DecLockRefResult:
+        binding_params = (
+            self._bindings.DecLockRefParamsBinding(
+                swa_uuid_for_lock=params.swa_uuid_for_lock,
+                swa_uuid_for_host_lock=params.swa_uuid_for_host_lock,
+                skip_lock_node_ids=_skip_lock_node_ids_to_binding(
+                    params.skip_lock_node_ids
+                ),
+            )
+            if params is not None
+            else None
+        )
+        self._binding.dec_host_lock_ref(node_id, binding_params)
+        return DecLockRefResult()
+
+    def evictable_size(self) -> int:
+        return self._binding.evictable_size()
+
+    def protected_size(self) -> int:
+        return self._binding.protected_size()
+
+    def component_evictable_size(self, component_type: ComponentType) -> int:
+        return self._binding.component_evictable_size(int(component_type))
+
+    def full_evictable_size(self) -> int:
+        return self._binding.full_evictable_size()
+
+    def full_protected_size(self) -> int:
+        return self._binding.full_protected_size()
+
+    def swa_evictable_size(self) -> int:
+        return self._binding.component_evictable_size(int(ComponentType.SWA))
+
+    def mamba_evictable_size(self) -> int:
+        return self._binding.component_evictable_size(int(ComponentType.MAMBA))
+
+    def swa_protected_size(self) -> int:
+        return self._binding.component_protected_size(int(ComponentType.SWA))
+
+    def mamba_protected_size(self) -> int:
+        return self._binding.component_protected_size(int(ComponentType.MAMBA))
+
+    def total_size(self) -> tuple[int, int]:
+        return self._binding.total_size()
+
+    def all_values_flatten(self) -> torch.Tensor:
+        return self._binding.all_values_flatten()
+
+    def walk_for_kv_canary(
+        self, unlocked_only: bool, swa_resident_only: bool
+    ) -> RadixCacheWalkResult:
+        result = self._binding.walk_for_kv_canary(unlocked_only, swa_resident_only)
+        return RadixCacheWalkResult(
+            slot_indices=result.slot_indices,
+            positions=result.positions,
+            prev_slot_indices=result.prev_slot_indices,
+        )
+
+    def _record_all_cleared_event(self) -> None:
+        self.kv_events.record_all_cleared()
+
+    def take_events(self) -> list:
+        return self.kv_events.take()
+
+    def all_mamba_values_flatten(self) -> torch.Tensor:
+        return self._binding.all_mamba_values_flatten()
+
+    def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+        key = params.key
+        result = self._binding.match_prefix(
+            self._bindings.MatchParamsBinding(
+                key=_radix_key_buffer(key),
+                extra_key=key.extra_key,
+                cache_salt=key.cache_salt,
+            )
+        )
+        return _match_result_from_binding(result)
+
+    @property
+    def empty_match_result(self) -> MatchResult:
+        return self._empty_match_result
+
+    def is_full_device_evicted(self, node_id: NodeId) -> bool:
+        return self._binding.is_full_device_evicted(node_id)
+
+    def collect_full_device_indices(
+        self, from_node_id: NodeId, until_node_id: NodeId
+    ) -> torch.Tensor:
+        return self._binding.collect_full_device_indices(from_node_id, until_node_id)
+
+    def begin_insert(self, params: InsertParams) -> InsertStepResult:
+        key = params.key
+        key_buffer = _radix_key_buffer(key)
+        value = params.value
+        if value is None:
+            # The binding always receives a value tensor; fall back to the
+            # token ids materialized on the core's device.
+            value = torch.tensor(key_buffer, dtype=torch.int64, device=self.device)
+        step = self._binding.begin_insert(
+            self._bindings.InsertParamsBinding(
+                key=key_buffer,
+                value=value,
+                extra_key=key.extra_key,
+                cache_salt=key.cache_salt,
+                mamba_value=params.mamba_value,
+                prev_prefix_len=params.prev_prefix_len,
+                swa_evicted_seqlen=params.swa_evicted_seqlen,
+                chunked=params.chunked,
+                priority=0 if params.priority is None else params.priority,
+                track_adopted_ranges=params.track_adopted_ranges,
+            )
+        )
+        return _insert_step_from_binding(step)
+
+    def resume_insert(self) -> InsertStepResult:
+        return _insert_step_from_binding(self._binding.resume_insert())
+
+    def has_ongoing_insert(self) -> bool:
+        return self._binding.has_ongoing_insert()
+
+    def end_insert(self) -> list[CacheAction | ComponentAction]:
+        return _cache_actions_from_tagged(self._binding.end_insert())
+
+    def drive_host_eviction(
+        self, component_type: ComponentType, num_tokens: int
+    ) -> DriveHostEvictionResult:
+        binding_result = self._binding.drive_host_eviction(
+            int(component_type), num_tokens
+        )
+        return _fill_evict_result(binding_result, DriveHostEvictionResult())
+
+    def evict_excess_path_states(
+        self,
+        tail_node_id: NodeId,
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
+    ) -> None:
+        binding_result = self._binding.evict_excess_path_states(tail_node_id)
+        for component, tensors in binding_result.new_device_frees.items():
+            device_frees[ComponentType(component)].extend(tensors)
+        for component, tensors in binding_result.new_host_frees.items():
+            host_frees[ComponentType(component)].extend(tensors)
+
+    # ==== HiCache ====
+
+    def set_hicache_enabled(self) -> None:
+        self._binding.set_hicache_enabled()
+
+    @property
+    def page_size(self) -> int:
+        # Read-only: the Rust core freezes it at construction.
+        return self._page_size
+
+    @property
+    def enable_hicache(self) -> bool:
+        return self._binding.enable_hicache()
+
+    @property
+    def has_swa_host_pool(self) -> bool:
+        return self._binding.has_swa_host_pool()
+
+    @has_swa_host_pool.setter
+    def has_swa_host_pool(self, value: bool) -> None:
+        # The Rust core has no unset path; reject a True -> False transition.
+        assert value or not self.has_swa_host_pool
+        if value:
+            self._binding.set_has_swa_host_pool()
+
+    @property
+    def write_through_threshold(self) -> int:
+        return self._binding.write_through_threshold()
+
+    @write_through_threshold.setter
+    def write_through_threshold(self, value: int) -> None:
+        # The cache assigns tree_core.write_through_threshold at HiCache init.
+        self._binding.set_write_through_threshold(value)
+
+    @property
+    def is_write_back(self) -> bool:
+        return self._binding.is_write_back()
+
+    @is_write_back.setter
+    def is_write_back(self, value: bool) -> None:
+        # The cache assigns tree_core.is_write_back at HiCache init; forward it.
+        self._binding.set_is_write_back(value)
+
+    @property
+    def enable_storage(self) -> bool:
+        return self._binding.enable_storage()
+
+    @enable_storage.setter
+    def enable_storage(self, value: bool) -> None:
+        # The cache assigns tree_core.enable_storage at storage init; forward it.
+        self._binding.set_enable_storage(value)
+
+    @property
+    def enable_external_cache_linker(self) -> bool:
+        return False
+
+    @enable_external_cache_linker.setter
+    def enable_external_cache_linker(self, value: bool) -> None:
+        # TODO(Jialin): Port external cache linker support from #37091 and #37151.
+        if value:
+            raise ValueError(
+                "External cache linker is not supported by the Rust TreeCore"
+            )
+
+    def insert_host(
+        self,
+        node_id: NodeId,
+        key: RadixKey,
+        host_value: torch.Tensor,
+        hash_value: list[str],
+    ) -> InsertResult:
+        result = self._binding.insert_host(
+            node_id,
+            key.extra_key,
+            _radix_key_buffer(key),
+            host_value,
+            list(hash_value),
+            key.cache_salt,
+        )
+        return InsertResult(
+            prefix_len=result.prefix_len,
+            total_len=result.total_len,
+            last_device_node=result.last_device_node,
+            inserted_host_node=result.inserted_host_node,
+            host_insert_dropped=result.host_insert_dropped,
+            mamba_exist=result.mamba_exist,
+            cache_actions=_cache_actions_from_tagged(result.cache_actions),
+        )
+
+    def build_backup_spec(
+        self, node_id: NodeId
+    ) -> tuple[torch.Tensor, dict[ComponentType, list[PoolTransfer]]]:
+        device_value, comp_xfers = self._binding.build_backup_spec(node_id)
+        return device_value, _comp_xfers_from_binding(comp_xfers)
+
+    def build_storage_backup_spec(
+        self, node_id: NodeId, pass_prefix_keys: bool
+    ) -> Optional[StorageBackupSpec]:
+        spec = self._binding.build_storage_backup_spec(node_id, pass_prefix_keys)
+        if spec is None:
+            return None
+        # Token ids cross the boundary as raw int64 bytes, not per-token ints.
+        token_ids = array("q")
+        token_ids.frombytes(spec.token_ids)
+        return StorageBackupSpec(
+            host_value=spec.host_value,
+            token_ids=token_ids,
+            hash_value=spec.hash_value,
+            prefix_keys=spec.prefix_keys,
+            comp_xfers=_comp_xfers_from_binding(spec.comp_xfers),
+        )
+
+    def build_hicache_transfers(
+        self,
+        component_type: ComponentType,
+        node_id: NodeId,
+        phase: CacheTransferPhase,
+        *,
+        host_indices: Optional[torch.Tensor] = None,
+        token_ids: Optional[Sequence[int]] = None,
+        prefetch_tokens: int = 0,
+        last_hash: Optional[str] = None,
+    ) -> Optional[list[PoolTransfer]]:
+        transfers = self._binding.build_hicache_transfers(
+            int(component_type),
+            node_id,
+            phase.value,
+            host_indices,
+            # TODO: Forward token ids when Rust Mamba prefetch consumes them.
+            None,
+            prefetch_tokens,
+            last_hash,
+        )
+        if transfers is None:
+            return None
+        return [_transfer_from_binding(transfer) for transfer in transfers]
+
+    def build_load_back_spec(
+        self, node_id: NodeId, req: Optional[Req] = None
+    ) -> tuple[PoolTransfer, dict[ComponentType, list[PoolTransfer]]]:
+        # Component hooks take primitives, not Req: extract its fields here.
+        mamba_pool_idx = req.kv.mamba_pool_idx if req is not None else None
+        kv_xfer, comp_xfers = self._binding.build_load_back_spec(
+            node_id, mamba_pool_idx
+        )
+        return _transfer_from_binding(kv_xfer), _comp_xfers_from_binding(comp_xfers)
+
+    def prefetch_anchor_info(
+        self, node_id: NodeId
+    ) -> tuple[Optional[str], Optional[str]]:
+        return self._binding.prefetch_anchor_info(node_id)
+
+    def is_backuped(self, node_id: NodeId) -> bool:
+        return self._binding.node_backuped(node_id)
+
+    def is_root(self, node_id: NodeId) -> bool:
+        return self._binding.is_root(node_id)
+
+    def get_last_hash_value(self, node_id: NodeId) -> Optional[str]:
+        return self._binding.get_last_hash_value(node_id)
+
+    def get_prefix_hash_values(self, node_id: NodeId) -> list[str]:
+        return self._binding.get_prefix_hash_values(node_id)
+
+    def get_hash_values(self, node_id: NodeId) -> list[str]:
+        return self._binding.get_hash_values(node_id)
+
+    def snapshot_buffer_backup(
+        self, node_id: NodeId, pass_prefix_keys: bool
+    ) -> Optional[BufferBackupSnapshot]:
+        snapshot = self._binding.snapshot_buffer_backup(node_id, pass_prefix_keys)
+        if snapshot is None:
+            return None
+        token_ids = array("q")
+        token_ids.frombytes(snapshot.key_token_ids)
+        return BufferBackupSnapshot(
+            node_id=snapshot.node_id,
+            parent_node_id=snapshot.parent_node_id,
+            parent_is_root=snapshot.parent_is_root,
+            parent_last_hash=snapshot.parent_last_hash,
+            hash_values=snapshot.hash_values,
+            key=RadixKey(
+                token_ids,
+                extra_key=snapshot.extra_key,
+                is_bigram=snapshot.is_bigram,
+                cache_salt=snapshot.cache_salt,
+            ),
+            prefix_keys=snapshot.prefix_keys,
+        )
+
+    def validate_buffer_backup(
+        self, node_id: NodeId, expected_key_length: int
+    ) -> Optional[BufferBackupState]:
+        state = self._binding.validate_buffer_backup(node_id, expected_key_length)
+        if state is None:
+            return None
+        return BufferBackupState(
+            parent_node_id=state.parent_node_id,
+            parent_is_root=state.parent_is_root,
+            parent_last_hash=state.parent_last_hash,
+        )
+
+    def backfill_missing_hash_values(self) -> int:
+        return self._binding.backfill_missing_hash_values()
+
+    def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId:
+        return self._binding.root_node_handle(extra_key)
+
+    def dfs_weight_order(self, node_ids: Sequence[NodeId]) -> list[int]:
+        return self._binding.dfs_weight_order(list(node_ids))
+
+    def commit_hicache_transfers(
+        self,
+        node_id: NodeId,
+        phase: CacheTransferPhase,
+        comp_xfers: dict[ComponentType, list[PoolTransfer]],
+        *,
+        cache_actions: list[CacheAction | ComponentAction],
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        actions, mamba_exist = self._binding.commit_hicache_transfers(
+            node_id,
+            phase.value,
+            _comp_xfers_to_binding(comp_xfers),
+            (
+                None
+                if insert_result is None
+                else (
+                    insert_result.total_len,
+                    insert_result.inserted_host_node,
+                    insert_result.mamba_exist,
+                )
+            ),
+            (
+                None
+                if pool_storage_result is None
+                else (
+                    pool_storage_result.kv_hit_pages,
+                    dict(pool_storage_result.extra_pool_hit_pages),
+                )
+            ),
+        )
+        if insert_result is not None and mamba_exist is not None:
+            insert_result.mamba_exist = mamba_exist
+        cache_actions.extend(_cache_actions_from_tagged(actions))
+
+    def commit_backup(
+        self,
+        node_id: NodeId,
+        host_indices: torch.Tensor,
+        comp_xfers: dict[ComponentType, list[PoolTransfer]],
+    ) -> None:
+        self._binding.commit_backup(
+            node_id, host_indices, _comp_xfers_to_binding(comp_xfers)
+        )
+
+    def commit_load_back(
+        self,
+        node_id: NodeId,
+        device_indices: torch.Tensor,
+        kv_xfer: PoolTransfer,
+        comp_xfers: dict[ComponentType, list[PoolTransfer]],
+    ) -> list[CacheAction | ComponentAction]:
+        actions = self._binding.commit_load_back(
+            node_id,
+            device_indices,
+            _transfer_to_binding(kv_xfer),
+            _comp_xfers_to_binding(comp_xfers),
+        )
+        return _cache_actions_from_tagged(actions)
+
+    def drop_subtree_no_host(self, node_id: NodeId) -> DropSubtreeNoHostResult:
+        binding_result = self._binding.drop_subtree_no_host(node_id)
+        result = DropSubtreeNoHostResult(is_dropped=binding_result.dropped)
+        return _fill_evict_result(binding_result, result)
+
+    def mark_write_through_pending(self, node_id: NodeId) -> None:
+        self._binding.mark_write_through_pending(node_id)
+
+    def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None:
+        self._binding.finish_write_through(list(node_ids), ack_id)
+
+    def finish_load_back(self, anchor_node_id: NodeId) -> None:
+        self._binding.finish_load_back(anchor_node_id)
+
+    @property
+    def write_back_duplicate_reclaim_digest(self) -> int:
+        return self._binding.write_back_duplicate_reclaim_digest()
+
+    def set_component_device_value(
+        self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor
+    ) -> None:
+        self._binding.set_component_device_value(
+            node_id, int(component_type), value.to(torch.int64)
+        )
+
+    def get_component_device_value(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> Optional[torch.Tensor]:
+        return self._binding.get_component_device_value(node_id, int(component_type))
+
+    def component_has_host_value_only(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> bool:
+        return self._binding.component_has_host_value_only(node_id, int(component_type))
+
+    # ==== Others ====
+
+    def sanity_check(
+        self,
+        ongoing_write_through: list[tuple[int, NodeId]],
+        ongoing_load_back: list[tuple[int, NodeId]],
+    ) -> None:
+        self._binding.sanity_check(ongoing_write_through, ongoing_load_back)
+
+    def pretty_print(self) -> None:
+        self._binding.pretty_print()

python/sglang/srt/mem_cache/rust_tree_core/extension.py

@@ -0,0 +1,40 @@
+"""Load the bundled Rust TreeCore extension or a fingerprinted local build."""
+
+from pathlib import Path
+
+# Loading torch first makes its libtorch dependencies resident before dlopen.
+import torch
+
+from sglang.srt.rust_extensions import load_rust_extension
+from sglang.srt.rust_extensions.torch_build import torch_build_configuration
+
+_PYTHON_MODULE = "sglang.srt.mem_cache.rust_tree_core.mem_cache"
+_INSPECTION_MODULE = "sglang.srt.mem_cache.rust_tree_core.mem_cache_inspection"
+_CRATE_DIR = Path(__file__).resolve().parents[5] / "rust" / "mem-cache"
+_TORCH_COMPAT_HEADER = _CRATE_DIR / "torch_2_13_compat.h"
+
+
+def load_tree_core_extension(*, inspection: bool = False):
+    """Load the production binding or the test-only inspection variant."""
+    build = torch_build_configuration(
+        compat_header=_TORCH_COMPAT_HEADER,
+        python_module=_PYTHON_MODULE,
+        torch_module=torch,
+    )
+    return load_rust_extension(
+        _PYTHON_MODULE,
+        additional_features=("inspection",) if inspection else (),
+        extension_module=_INSPECTION_MODULE if inspection else None,
+        build_environment=build.environment,
+        build_fingerprint=build.fingerprint,
+    )
+
+
+bindings = load_tree_core_extension()
+
+DecLockRefParamsBinding = bindings.DecLockRefParamsBinding
+InsertParamsBinding = bindings.InsertParamsBinding
+MatchParamsBinding = bindings.MatchParamsBinding
+RustBigramUnifiedTreeCoreBinding = bindings.RustBigramUnifiedTreeCoreBinding
+RustUnifiedTreeCoreBinding = bindings.RustUnifiedTreeCoreBinding
+TreeCoreInitParamsBinding = bindings.TreeCoreInitParamsBinding

python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py

@@ -54,7 +54,17 @@ def _python_tree_core_factory(
     return UnifiedTreeCore(params, components)
 
 
+def _rust_tree_core_factory(
+    params: CacheInitParams, components: dict[ComponentType, TreeComponent]
+) -> UnifiedTreeCoreInterface:
+    """Load and construct the in-tree Rust TreeCore only when selected."""
+    from sglang.srt.mem_cache.rust_tree_core.adapter import RustUnifiedTreeCore
+
+    return RustUnifiedTreeCore(params)
+
+
 register_tree_core_backend("python", _python_tree_core_factory)
+register_tree_core_backend("rust", _rust_tree_core_factory)
 
 
 def create_tree_core(

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -35,6 +35,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     InsertResult,
     MatchPrefixParams,
     MatchResult,
+    _dfs_weight_order,
 )
 from sglang.srt.mem_cache.events import KVCacheEventRecorder
 from sglang.srt.mem_cache.hicache_storage import (
@@ -63,6 +64,8 @@ from sglang.srt.mem_cache.unified_cache.components import (
     get_and_increase_time_counter,
 )
 from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
+    BufferBackupSnapshot,
+    BufferBackupState,
     DecSwaLockOnlyResult,
     DemoteResult,
     DriveHostEvictionResult,
@@ -518,6 +521,55 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         """The hash values owned by this node, excluding its ancestors."""
         return self.node_by_id(node_id).hash_value or []
 
+    def snapshot_buffer_backup(
+        self, node_id: NodeId, pass_prefix_keys: bool
+    ) -> Optional[BufferBackupSnapshot]:
+        node = self._node_arena.get(node_id)
+        if (
+            node is None
+            or node is self.root_node
+            or not node.hash_value
+            or node.component_data[BASE_COMPONENT_TYPE].value is None
+        ):
+            return None
+        parent = node.parent
+        assert parent is not None and node.key is not None
+        return BufferBackupSnapshot(
+            node_id=node.id,
+            parent_node_id=parent.id,
+            parent_is_root=parent is self.root_node,
+            parent_last_hash=parent.get_last_hash_value(),
+            hash_values=list(node.hash_value),
+            key=RadixKey(
+                array("q", node.key.raw_token_ids()),
+                extra_key=node.key.extra_key,
+                is_bigram=node.key.is_bigram,
+                cache_salt=node.key.cache_salt,
+            ),
+            prefix_keys=(
+                node.get_prefix_hash_values(parent) if pass_prefix_keys else None
+            ),
+        )
+
+    def validate_buffer_backup(
+        self, node_id: NodeId, expected_key_length: int
+    ) -> Optional[BufferBackupState]:
+        node = self._node_arena.get(node_id)
+        if (
+            node is None
+            or node.component_data[BASE_COMPONENT_TYPE].value is None
+            or len(node.key) != expected_key_length
+        ):
+            return None
+        parent = node.parent
+        if parent is None:
+            return None
+        return BufferBackupState(
+            parent_node_id=parent.id,
+            parent_is_root=parent is self.root_node,
+            parent_last_hash=parent.get_last_hash_value(),
+        )
+
     def backfill_missing_hash_values(self) -> int:
         """Hash every node that was built while storage was disabled.
 
@@ -543,6 +595,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         """The NodeId anchoring matches; the single root serves every namespace."""
         return self.root_node.id
 
+    def dfs_weight_order(self, node_ids: Sequence[NodeId]) -> list[int]:
+        return _dfs_weight_order(self.root_node, node_ids, self.node_by_id)
+
     def _new_node(self, priority: int = 0) -> UnifiedTreeNode:
         """Create and register a tree node in the arena."""
         node = UnifiedTreeNode(self.component_types, priority=priority)
@@ -872,7 +927,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
     def begin_insert(self, params: InsertParams) -> InsertStepResult:
         """Start the insert, running to its first barrier or completion."""
         # Insert walks are single-flight; a live walk means re-entrancy.
-        assert self._ongoing_insert_walk_state is None, "concurrent insert walks"
+        if self._ongoing_insert_walk_state is not None:
+            raise RuntimeError("concurrent insert walks")
         key = params.key
         value = params.value
         key, value = key.maybe_to_bigram_view(self.is_eagle, value)
@@ -913,7 +969,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
 
     def resume_insert(self) -> InsertStepResult:
         """Continue the suspended insert after its step actions were executed."""
-        assert self._ongoing_insert_walk_state is not None, "no in-flight insert"
+        if self._ongoing_insert_walk_state is None:
+            raise RuntimeError("no in-flight insert")
         return self._advance_insert()
 
     def has_ongoing_insert(self) -> bool:

python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py

@@ -80,6 +80,22 @@ class RadixCacheWalkResult(msgspec.Struct, frozen=True, kw_only=True):
     prev_slot_indices: torch.Tensor
 
 
+class BufferBackupSnapshot(msgspec.Struct, frozen=True):
+    node_id: NodeId
+    parent_node_id: NodeId
+    parent_is_root: bool
+    parent_last_hash: Optional[str]
+    hash_values: list[str]
+    key: RadixKey
+    prefix_keys: Optional[list[str]]
+
+
+class BufferBackupState(msgspec.Struct, frozen=True):
+    parent_node_id: NodeId
+    parent_is_root: bool
+    parent_last_hash: Optional[str]
+
+
 class InsertStepResult(msgspec.Struct, frozen=True):
     """One step of a resumable insert: the Controller executes ``actions``, then
     resumes while ``result`` is None; ``result`` is set on the final step."""
@@ -181,6 +197,20 @@ class UnifiedTreeCoreInterface(ABC):
         """The hash values owned by this node, excluding its ancestors."""
         ...
 
+    @abstractmethod
+    def snapshot_buffer_backup(
+        self, node_id: NodeId, pass_prefix_keys: bool
+    ) -> Optional[BufferBackupSnapshot]:
+        """Snapshot an eligible buffer-only backup node."""
+        ...
+
+    @abstractmethod
+    def validate_buffer_backup(
+        self, node_id: NodeId, expected_key_length: int
+    ) -> Optional[BufferBackupState]:
+        """Validate a queued backup and return its current parent state."""
+        ...
+
     @abstractmethod
     def backfill_missing_hash_values(self) -> int:
         """Hash every node built while storage was disabled; return how many.
@@ -196,6 +226,11 @@ class UnifiedTreeCoreInterface(ABC):
         """The NodeId anchoring matches for the namespace."""
         ...
 
+    @abstractmethod
+    def dfs_weight_order(self, node_ids: Sequence[NodeId]) -> list[int]:
+        """Return input indices in depth-first, subtree-weight order."""
+        ...
+
     @abstractmethod
     def inc_lock_ref(
         self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -199,8 +199,9 @@ class UnifiedRadixCache(BasePrefixCache):
         )
         # The TreeCore owns the tree member-var state (structure, LRUs, sizes,
         # evictable leaves) and drives the components' tree-level hooks.
+        self._tree_core_backend = envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get()
         self.tree_core = create_tree_core(
-            name=envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get(),
+            name=self._tree_core_backend,
             params=params,
             components=self.components,
         )
@@ -386,6 +387,8 @@ class UnifiedRadixCache(BasePrefixCache):
         """Initialize HiCache infrastructure."""
         self.host_memory_mode = get_memory().hicache_host_memory_mode
         if self.host_memory_mode == "buffer_only":
+            # TODO(Jialin): Extend buffer-only state handoff to Mamba in a
+            # follow-up to #34798 and #35769.
             # FULL and FULL+SWA only: Mamba has no state-handoff channel on
             # the admission-time load-back read path and is not layer-gated.
             # Lifting the fence also needs the admission charge: a staged
@@ -1340,9 +1343,7 @@ class UnifiedRadixCache(BasePrefixCache):
             # FIFO ordering instead (BackupKV chains are parent-before-child
             # and every pipeline stage drains in order).
             for node_id in action.node_ids:
-                self.buffer_pipeline.enqueue_backup_intent(
-                    self.tree_core.node_by_id(node_id)
-                )
+                self.buffer_pipeline.enqueue_backup_intent(node_id)
             return 0
         written = 0
         for node_id in action.node_ids:
@@ -3120,3 +3121,6 @@ class UnifiedRadixCache(BasePrefixCache):
     def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId:
         """The root's NodeId -- URC match results carry NodeIds."""
         return self.tree_core.root_node_handle(extra_key)
+
+    def dfs_weight_order(self, node_handles: Sequence[NodeId]) -> list[int]:
+        return self.tree_core.dfs_weight_order(node_handles)

python/sglang/srt/rust_extensions/loader.py

@@ -19,7 +19,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
 from types import ModuleType
-from typing import Iterator, Literal
+from typing import Iterator, Literal, Mapping
 
 try:
     import tomllib
@@ -52,6 +52,7 @@ class _CrateSpec:
     package: str
     library: str
     python_module: str
+    manifest: Path
     workspace: Path
     features: tuple[str, ...]
 
@@ -69,6 +70,10 @@ def load_rust_extension(
     mode: RustBuildMode | None = None,
     cache_dir: Path | None = None,
     workspace: Path | None = None,
+    additional_features: tuple[str, ...] = (),
+    extension_module: str | None = None,
+    build_environment: Mapping[str, str] | None = None,
+    build_fingerprint: Mapping[str, object] | None = None,
 ) -> ModuleType:
     """Import a PyO3 extension, compiling it locally when permitted and needed.
 
@@ -77,9 +82,13 @@ def load_rust_extension(
     to ``python_module`` (the same metadata setup.py uses for wheel builds), so
     new crates need no registration here.
 
-    ``auto`` prefers a module bundled in the installed wheel, then a cached
-    local build, and finally Cargo. ``never`` permits the first two but never
-    invokes Cargo. ``force`` rebuilds from source and replaces the cache entry.
+    ``auto`` prefers a module bundled in an installed wheel. In a source tree,
+    it ignores unverified in-package artifacts and uses the fingerprinted cache
+    before invoking Cargo. ``never`` explicitly trusts a bundled module, then
+    permits the cache but never invokes Cargo. ``force`` rebuilds from source.
+    A same-name feature variant is always sourced from the fingerprinted cache.
+    A distinctly named variant may be supplied by test infrastructure and is
+    otherwise built into that cache after its bundled import misses.
     ``mode`` defaults to ``SGLANG_RUST_BUILD_MODE``.
     """
     if mode is None:
@@ -89,29 +98,45 @@ def load_rust_extension(
             f"invalid Rust extension build mode {mode!r}; expected auto, never, or force"
         )
 
-    if mode != "force":
-        module = _import_bundled_extension(python_module)
-        if module is not None:
-            return module
-    elif python_module in sys.modules:
+    load_module = extension_module or python_module
+    same_name_feature_variant = (
+        bool(additional_features) and load_module == python_module
+    )
+    if loaded := sys.modules.get(load_module):
+        if mode != "force":
+            return loaded
         raise RuntimeError(
-            f"cannot force-build {python_module} after it has been imported; "
+            f"cannot force-build {load_module} after it has been imported; "
             "start a new Python process"
         )
-
     if workspace is None:
         workspace = _RUST_WORKSPACE
+    source_checkout = (Path(workspace) / "Cargo.toml").is_file()
+    trust_bundled = mode == "never" or not source_checkout
+    if mode != "force" and trust_bundled and not same_name_feature_variant:
+        module = _import_bundled_extension(load_module)
+        if module is not None:
+            return module
+
     crate = _discover_crate(workspace, python_module)
-    context = _build_context(crate)
+    features = tuple(dict.fromkeys((*crate.features, *additional_features)))
+    context = _build_context(
+        crate,
+        features=features,
+        build_fingerprint=build_fingerprint,
+        extension_module=load_module,
+    )
     cache_root = _cache_root(cache_dir)
-    extension_path = _cached_extension_path(cache_root, crate, context.fingerprint)
+    extension_path = _cached_extension_path(
+        cache_root, crate, context.fingerprint, load_module
+    )
     lock_path = (
         cache_root / "locks" / f"{crate.package}-{context.target_fingerprint}.lock"
     )
 
     with _filesystem_lock(lock_path):
         if mode != "force" and extension_path.is_file():
-            return _load_extension_from_path(crate.python_module, extension_path)
+            return _load_extension_from_path(load_module, extension_path)
 
         if mode == "never":
             raise ModuleNotFoundError(
@@ -121,14 +146,19 @@ def load_rust_extension(
             )
 
         target_dir = cache_root / "targets" / context.target_fingerprint
-        artifact = _cargo_build(crate, target_dir)
+        artifact = _cargo_build(
+            crate,
+            target_dir,
+            features=features,
+            build_environment=build_environment,
+        )
         if _source_digest(crate.workspace) != context.source_digest:
             raise RuntimeError(
                 f"Rust sources under {crate.workspace} changed during the build; "
                 "the result was not cached"
             )
         _stage_atomically(artifact, extension_path)
-        return _load_extension_from_path(crate.python_module, extension_path)
+        return _load_extension_from_path(load_module, extension_path)
 
 
 def _import_bundled_extension(module_name: str) -> ModuleType | None:
@@ -143,16 +173,10 @@ def _import_bundled_extension(module_name: str) -> ModuleType | None:
 def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec:
     workspace = Path(workspace).resolve()
     workspace_manifest = workspace / "Cargo.toml"
-    lockfile = workspace / "Cargo.lock"
     if not workspace_manifest.is_file():
         raise FileNotFoundError(
             f"Rust workspace for {python_module} was not found at {workspace}"
         )
-    if not lockfile.is_file():
-        raise FileNotFoundError(
-            f"{lockfile} is required for reproducible `cargo build --locked` builds"
-        )
-
     matches: list[_CrateSpec] = []
     declared_modules: list[str] = []
     for manifest in _source_files(workspace):
@@ -178,12 +202,19 @@ def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec:
                 f"{manifest} declares python-module {python_module!r} but must "
                 "also set `package.name` and `lib.name`"
             )
+        crate_workspace = manifest.parent if "workspace" in document else workspace
+        lockfile = crate_workspace / "Cargo.lock"
+        if not lockfile.is_file():
+            raise FileNotFoundError(
+                f"{lockfile} is required for reproducible `cargo build --locked` builds"
+            )
         matches.append(
             _CrateSpec(
                 package=package_name,
                 library=library,
                 python_module=python_module,
-                workspace=workspace,
+                manifest=manifest,
+                workspace=crate_workspace,
                 features=tuple(sglang_metadata.get("features", ())),
             )
         )
@@ -203,7 +234,17 @@ def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec:
     return matches[0]
 
 
-def _build_context(crate: _CrateSpec) -> _BuildContext:
+def _build_context(
+    crate: _CrateSpec,
+    *,
+    features: tuple[str, ...] | None = None,
+    build_fingerprint: Mapping[str, object] | None = None,
+    extension_module: str | None = None,
+) -> _BuildContext:
+    if features is None:
+        features = crate.features
+    if extension_module is None:
+        extension_module = crate.python_module
     source_digest = _source_digest(crate.workspace)
     toolchain = {
         "cargo": _command_version(
@@ -224,6 +265,7 @@ def _build_context(crate: _CrateSpec) -> _BuildContext:
     }
     target_inputs = {
         "build_environment": build_environment,
+        "extension_build": dict(build_fingerprint or {}),
         "python_abi": python_abi,
         "toolchain": toolchain,
     }
@@ -234,6 +276,8 @@ def _build_context(crate: _CrateSpec) -> _BuildContext:
             "package": crate.package,
             "library": crate.library,
             "python_module": crate.python_module,
+            "extension_module": extension_module,
+            "features": features,
             "source_digest": source_digest,
             **target_inputs,
         }
@@ -301,12 +345,15 @@ def _cache_root(cache_dir: Path | None) -> Path:
 
 
 def _cached_extension_path(
-    cache_root: Path, crate: _CrateSpec, fingerprint: str
+    cache_root: Path,
+    crate: _CrateSpec,
+    fingerprint: str,
+    extension_module: str | None = None,
 ) -> Path:
     extension_suffix = sysconfig.get_config_var("EXT_SUFFIX")
     if not extension_suffix:
         raise RuntimeError("Python did not report an EXT_SUFFIX for native extensions")
-    module_leaf = crate.python_module.rsplit(".", 1)[-1]
+    module_leaf = (extension_module or crate.python_module).rsplit(".", 1)[-1]
     return (
         cache_root
         / "artifacts"
@@ -327,7 +374,15 @@ def _filesystem_lock(path: Path) -> Iterator[None]:
             fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
 
 
-def _cargo_build(crate: _CrateSpec, target_dir: Path) -> Path:
+def _cargo_build(
+    crate: _CrateSpec,
+    target_dir: Path,
+    *,
+    features: tuple[str, ...] | None = None,
+    build_environment: Mapping[str, str] | None = None,
+) -> Path:
+    if features is None:
+        features = crate.features
     target_dir.mkdir(parents=True, exist_ok=True)
     command = [
         "cargo",
@@ -337,10 +392,10 @@ def _cargo_build(crate: _CrateSpec, target_dir: Path) -> Path:
         "--package",
         crate.package,
     ]
-    if crate.features:
-        command.extend(("--features", ",".join(crate.features)))
+    if features:
+        command.extend(("--features", ",".join(features)))
 
-    environment = os.environ.copy()
+    environment = dict(os.environ if build_environment is None else build_environment)
     environment["CARGO_TARGET_DIR"] = os.fspath(target_dir)
     environment["PYO3_PYTHON"] = sys.executable
     logger.info("Building %s with `%s`", crate.python_module, " ".join(command))

python/sglang/srt/rust_extensions/torch_build.py

@@ -0,0 +1,115 @@
+"""Build settings for Rust extensions that link against the active PyTorch."""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import re
+import shlex
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from types import ModuleType
+from typing import Mapping
+
+_MIN_SUPPORTED_TORCH = (2, 11)
+_MAX_SUPPORTED_TORCH = (2, 13)
+
+
+@dataclass(frozen=True)
+class TorchBuildConfiguration:
+    """Environment overrides plus stable inputs for the artifact fingerprint."""
+
+    environment: dict[str, str]
+    fingerprint: dict[str, object]
+
+
+def torch_build_configuration(
+    *,
+    compat_header: Path,
+    python_module: str,
+    torch_module: ModuleType | None = None,
+    base_environment: Mapping[str, str] | None = None,
+    include_absolute_rpath: bool = True,
+) -> TorchBuildConfiguration:
+    """Describe a build against the torch package loaded by this interpreter."""
+    if sys.platform != "linux":
+        raise RuntimeError("the Rust TreeCore extension currently supports Linux only")
+
+    if torch_module is None:
+        try:
+            import torch as torch_module
+        except ModuleNotFoundError as exc:
+            raise RuntimeError(
+                "PyTorch must be installed before building the Rust TreeCore extension"
+            ) from exc
+
+    version = str(torch_module.__version__)
+    match = re.match(r"^(\d+)\.(\d+)", version)
+    if match is None:
+        raise RuntimeError(f"could not parse PyTorch version {version!r}")
+    major_minor = (int(match.group(1)), int(match.group(2)))
+    if not _MIN_SUPPORTED_TORCH <= major_minor <= _MAX_SUPPORTED_TORCH:
+        minimum = ".".join(map(str, _MIN_SUPPORTED_TORCH))
+        maximum = ".".join(map(str, _MAX_SUPPORTED_TORCH))
+        raise RuntimeError(
+            f"the Rust TreeCore supports PyTorch {minimum} through {maximum}; "
+            f"found {version}"
+        )
+
+    torch_file = getattr(torch_module, "__file__", None)
+    if torch_file is None:
+        raise RuntimeError("the active PyTorch package has no filesystem location")
+    torch_root = Path(torch_file).resolve().parent
+    torch_lib = torch_root / "lib"
+    if not torch_lib.is_dir():
+        raise RuntimeError(
+            f"the active PyTorch package has no library dir at {torch_lib}"
+        )
+
+    cxx11_abi_fn = getattr(torch_module, "compiled_with_cxx11_abi", None)
+    if cxx11_abi_fn is not None:
+        cxx11_abi = bool(cxx11_abi_fn())
+    else:
+        cxx11_abi = bool(torch_module._C._GLIBCXX_USE_CXX11_ABI)
+
+    environment = dict(os.environ if base_environment is None else base_environment)
+    environment["LIBTORCH_USE_PYTORCH"] = "1"
+    # tch 0.24 targets Torch 2.11. The compatibility header below covers the
+    # API removals in the supported 2.12/2.13 builds, after this explicit gate.
+    environment["LIBTORCH_BYPASS_VERSION_CHECK"] = "1"
+    environment["PYO3_PYTHON"] = sys.executable
+    environment["PATH"] = os.pathsep.join(
+        filter(None, (os.fspath(Path(sys.executable).parent), environment.get("PATH")))
+    )
+    environment["LD_LIBRARY_PATH"] = os.pathsep.join(
+        filter(None, (os.fspath(torch_lib), environment.get("LD_LIBRARY_PATH")))
+    )
+
+    cxxflags = environment.get("CXXFLAGS", "")
+    environment["CXXFLAGS"] = (
+        f"{cxxflags} -include {shlex.quote(os.fspath(compat_header.resolve()))}"
+    ).strip()
+
+    package_depth = len(python_module.split(".")) - 1
+    bundled_torch_lib = "$ORIGIN/" + "../" * package_depth + "torch/lib"
+    rustflags = environment.get("RUSTFLAGS", "")
+    rpath_flags = [f"-C link-arg=-Wl,-rpath,{bundled_torch_lib}"]
+    if include_absolute_rpath:
+        rpath_flags.append(f"-C link-arg=-Wl,-rpath,{torch_lib}")
+    environment["RUSTFLAGS"] = " ".join(filter(None, (rustflags, *rpath_flags)))
+
+    fingerprint = {
+        "torch_version": version,
+        "torch_root": os.fspath(torch_root),
+        "torch_cxx11_abi": cxx11_abi,
+        "torch_cuda": getattr(torch_module.version, "cuda", None),
+        "torch_hip": getattr(torch_module.version, "hip", None),
+        "include_absolute_rpath": include_absolute_rpath,
+        "compat_header_sha256": (
+            hashlib.sha256(compat_header.read_bytes()).hexdigest()
+            if compat_header.is_file()
+            else None
+        ),
+    }
+    return TorchBuildConfiguration(environment=environment, fingerprint=fingerprint)

python/sglang/test/test_utils.py

@@ -658,6 +658,17 @@ def _wait_for_server_health(
     return False, "Server failed to start within the timeout period"
 
 
+def unified_radix_tree_server_env(
+    tree_core_backend: str, **extra_env: str
+) -> dict[str, str]:
+    return {
+        **os.environ,
+        **extra_env,
+        "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+        "SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND": tree_core_backend,
+    }
+
+
 def popen_launch_server(
     model: str,
     base_url: str,