[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>

.github/actions/download-rust-ext/action.yml

@@ -1,8 +1,9 @@
 name: 'Download prebuilt Rust extensions'
 description: >
   Put rust-ext-build's PyO3 extension modules in the checkout and set
-  SGLANG_BUILD_RUST_EXTS=none for the job, so install skips the cargo build.
-  Sets nothing when neither source has them, leaving install to compile.
+  SGLANG_BUILD_RUST_EXTS=none and SGLANG_RUST_BUILD_MODE=never for the job, so
+  install skips the cargo build and runtime trusts these fingerprinted modules.
+  Sets nothing when neither source has them, leaving source builds enabled.
 
 inputs:
   artifact_name:
@@ -49,8 +50,10 @@ runs:
       if: steps.artifact.outcome != 'success'
       uses: actions/cache/restore@v4
       with:
-        path: python/sglang/srt/rust_extensions/_*.so
-        key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }}
+        path: |
+          python/sglang/srt/rust_extensions/_*.so
+          python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
+        key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
 
     # Job-wide, but only setup.py reads it, and only while building.
     # Whether the modules suit this interpreter is not decided here:
@@ -65,6 +68,9 @@ runs:
            || [ "${{ steps.cache.outputs.cache-hit }}" = "true" ]; then
             echo "hit=true" >> "$GITHUB_OUTPUT"
             echo "SGLANG_BUILD_RUST_EXTS=none" >> "$GITHUB_ENV"
+            # A source checkout normally ignores in-package native artifacts
+            # because they may be stale. These are fingerprint-keyed CI bytes.
+            echo "SGLANG_RUST_BUILD_MODE=never" >> "$GITHUB_ENV"
         else
             echo "hit=false" >> "$GITHUB_OUTPUT"
         fi

.github/workflows/_pr-test-rust-ext-build.yml

@@ -82,6 +82,8 @@ jobs:
           sparse-checkout: |
             rust
             python/setup.py
+            python/pyproject.toml
+            python/sglang/srt/rust_extensions/torch_build.py
             .github
             scripts/ci/utils
           sparse-checkout-cone-mode: false
@@ -93,14 +95,16 @@ jobs:
       - name: Ensure zstd so the saved entry is readable
         run: bash scripts/ci/utils/ensure_zstd.sh
 
-      # setup.py counts because it selects which crates get built. pyproject.toml
-      # is left out - it churns on bumps that cannot affect these modules.
+      # The setup hook and torch helper select and configure what gets built;
+      # pyproject.toml pins the libtorch ABI used by mem-cache.
       - name: Restore built modules
         id: cache
         uses: actions/cache/restore@v4
         with:
-          path: python/sglang/srt/rust_extensions/_*.so
-          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }}
+          path: |
+            python/sglang/srt/rust_extensions/_*.so
+            python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
+          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
 
       # On a miss: different hash = rust/setup.py moved; no entries = evicted.
       - name: Report cache lookup
@@ -112,7 +116,8 @@ jobs:
         run: |
           if [ -n "${MATCHED_KEY}" ]; then
               echo "hit: ${MATCHED_KEY}"
-              ls -l python/sglang/srt/rust_extensions/_*.so
+              ls -l python/sglang/srt/rust_extensions/_*.so \
+                python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
           else
               echo "miss: ${PRIMARY_KEY}"
               echo "entries under ${KEY_PREFIX}- (created / ref / size / key):"
@@ -134,7 +139,7 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ inputs.artifact_name }}
-          # Archive holds rust_extensions/_*.so, so it unpacks into python/sglang/srt/.
+          # Archive holds package-relative paths, so it unpacks into python/sglang/srt/.
           path: rust-ext-staging/
           if-no-files-found: error
           retention-days: 1
@@ -182,23 +187,21 @@ jobs:
           export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
           # Per-interpreter subdirs (set in the loop): PyO3's fingerprint tracks
           # the interpreter, so a shared dir rebuilds on every ABI switch.
-          cargo_target_root="${HOME}/.cache/sglang-cargo-target"
+          # ci_install_dependency.sh drops ${HOME}/.cache/sglang-cargo-target at 85%
+          # disk and unlocks before its own build, so a CUDA job sharing this host
+          # can delete the tree mid compile. Build in a per run dir nothing else
+          # touches. The .so cache above still carries results across runs.
+          cargo_target_root="${RUNNER_TEMP:-/tmp}/sglang-cargo-target-${GITHUB_RUN_ID:-norun}-$$"
           mkdir -p "${cargo_target_root}"
-          used="$(df --output=pcent "${cargo_target_root}" 2>/dev/null | tr -dc '0-9')"
-          if [ "${used:-0}" -ge 85 ]; then
-              echo "cargo target dir filesystem at ${used}%; dropping ${cargo_target_root}"
-              rm -rf "${cargo_target_root}"
-              mkdir -p "${cargo_target_root}"
-          fi
           python3 -m pip install --upgrade pip
           command -v uv >/dev/null 2>&1 || pip install uv
-          # build_rust needs only the build backend, not sglang's ~294 runtime deps.
+          # build_rust needs the build backend and torch, not sglang's ~294 other runtime deps.
           # Per-job path: these runners are persistent and shared, so a fixed one
           # both inherits the previous job's venv and races a concurrent build.
           venv_root="${RUNNER_TEMP:-/tmp}/sglang-ci-rust-ext-${GITHUB_RUN_ID:-norun}-$$"
           # Best-effort, like ci_cleanup_venv.sh: under set -e a failing EXIT trap
           # would fail the step, and nothing here is worth keeping for a postmortem.
-          trap 'rm -rf "${venv_root}" || true' EXIT
+          trap 'rm -rf "${venv_root}" "${cargo_target_root}" || true' EXIT
           for python_bin in "${{ steps.py310.outputs.python-path }}" "${{ steps.py312.outputs.python-path }}"; do
               minor="$("${python_bin}" -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')"
               export CARGO_TARGET_DIR="${cargo_target_root}/py${minor}"
@@ -206,8 +209,87 @@ jobs:
               uv venv "${venv}" --python "${python_bin}" --seed
               # shellcheck disable=SC1091
               source "${venv}/bin/activate"
-              uv pip install "setuptools>=61.0" "setuptools-rust>=1.10" "setuptools-scm>=8.0" wheel
+              uv pip install "setuptools>=61.0" "setuptools-rust>=1.11" "setuptools-scm>=8.0" "torch==2.13.0" wheel
+              # torch-sys resolves libtorch from the active interpreter and bakes that
+              # path into the persistent cargo cache. venv_root is per-run and deleted on
+              # exit, so a later run reuses a "Fresh torch-sys" whose -L points at a gone
+              # directory and fails with "unable to find library -ltorch". Pin LIBTORCH and
+              # drop the torch-shim units so they rebuild against this run's venv.
+              LIBTORCH="$(python -c 'import pathlib, torch; print(pathlib.Path(torch.__file__).parent)')"
+              export LIBTORCH
+              export LD_LIBRARY_PATH="${LIBTORCH}/lib:${LD_LIBRARY_PATH:-}"
+              cargo clean --release --manifest-path rust/mem-cache/Cargo.toml \
+                  -p torch-sys -p mem_cache 2>/dev/null || true
               (cd python && SGLANG_BUILD_RUST_EXTS=all python setup.py build_rust --inplace)
+              python - <<'PY'
+          import importlib.util
+          import os
+          import pathlib
+          import runpy
+          import shutil
+          import subprocess
+          import sysconfig
+
+          import torch  # noqa: F401 - preload libtorch before the extension
+
+          root = pathlib.Path.cwd()
+          suffix = sysconfig.get_config_var("EXT_SUFFIX")
+          production_path = pathlib.Path(
+              "python/sglang/srt/mem_cache/rust_tree_core/mem_cache" + suffix
+          ).resolve()
+          inspection_path = production_path.with_name(
+              "mem_cache_inspection" + suffix
+          )
+
+          helper = runpy.run_path(
+              root / "python/sglang/srt/rust_extensions/torch_build.py"
+          )
+          build = helper["torch_build_configuration"](
+              compat_header=root / "rust/mem-cache/torch_2_13_compat.h",
+              python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
+              torch_module=torch,
+              include_absolute_rpath=False,
+          )
+          subprocess.run(
+              [
+                  "cargo",
+                  "build",
+                  "--release",
+                  "--locked",
+                  "--manifest-path",
+                  "rust/mem-cache/Cargo.toml",
+                  "--features",
+                  "python-extension,inspection",
+              ],
+              env=build.environment,
+              check=True,
+          )
+          release_dir = pathlib.Path(os.environ["CARGO_TARGET_DIR"]) / "release"
+          if target := os.environ.get("CARGO_BUILD_TARGET"):
+              release_dir = pathlib.Path(os.environ["CARGO_TARGET_DIR"]) / target / "release"
+          shutil.copy2(release_dir / "libmem_cache.so", inspection_path)
+
+          def load(name, path):
+              spec = importlib.util.spec_from_file_location(name, path)
+              assert spec is not None and spec.loader is not None
+              module = importlib.util.module_from_spec(spec)
+              spec.loader.exec_module(module)
+              return module
+
+          production = load(
+              "sglang.srt.mem_cache.rust_tree_core.mem_cache", production_path
+          )
+          inspection = load(
+              "sglang.srt.mem_cache.rust_tree_core.mem_cache_inspection",
+              inspection_path,
+          )
+          assert not hasattr(
+              production.RustUnifiedTreeCoreBinding, "inspect_contains_node"
+          )
+          assert hasattr(
+              inspection.RustUnifiedTreeCoreBinding, "inspect_contains_node"
+          )
+          PY
               deactivate
           done
 
@@ -226,8 +308,10 @@ jobs:
       - name: Save built modules
         uses: actions/cache/save@v4
         with:
-          path: python/sglang/srt/rust_extensions/_*.so
-          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }}
+          path: |
+            python/sglang/srt/rust_extensions/_*.so
+            python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
+          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
 
       - name: Upload extension modules
         uses: actions/upload-artifact@v4

.github/workflows/release-pypi-nightly.yml

@@ -24,7 +24,8 @@ jobs:
     if: github.repository == 'sgl-project/sglang'
     runs-on: ubuntu-latest
     outputs:
-      nightly_version: ${{ steps.build.outputs.nightly_version }}
+      nightly_version: ${{ steps.wheel.outputs.wheel_version }}
+      wheel_filename: ${{ steps.wheel.outputs.wheel_filename }}
       commit_hash: ${{ steps.build.outputs.commit_hash }}
       build_date: ${{ steps.build.outputs.build_date }}
     steps:
@@ -41,7 +42,9 @@ jobs:
 
       - name: Install build dependencies
         run: |
-          pip install build wheel setuptools setuptools-scm
+          pip install \
+            auditwheel build patchelf wheel "setuptools>=61.0" \
+            "setuptools-rust>=1.11" "setuptools-scm>=8.0" "torch==2.13.0"
 
       # Needed by setuptools-rust to build the bundled native gRPC extension
       # (rust/sglang-grpc) when `python -m build` builds the sglang wheel.
@@ -66,7 +69,7 @@ jobs:
           MINOR=$(echo "$VERSION" | cut -d. -f2)
           PATCH_RAW=$(echo "$VERSION" | cut -d. -f3)
           # Strip pre-release suffixes (rc0, post1, etc.) to get numeric patch
-          PATCH=$(echo "$PATCH_RAW" | sed 's/[^0-9].*//')
+          PATCH=${PATCH_RAW%%[^0-9]*}
           NEXT_PATCH=$((PATCH + 1))
           NEXT_VERSION="${MAJOR}.${MINOR}.${NEXT_PATCH}"
 
@@ -77,24 +80,25 @@ jobs:
           export SETUPTOOLS_SCM_PRETEND_VERSION="$FORCE_VERSION"
 
           # Build wheel
-          python3 -m build --wheel
-
-          # Extract version from built wheel filename
-          WHEEL_FILE=$(ls dist/*.whl)
-          NIGHTLY_VERSION=$(echo "$WHEEL_FILE" | sed 's/.*sglang-\(.*\)-py3.*/\1/')
+          python3 -m build --wheel --no-isolation
 
           # Get commit info
           COMMIT_HASH=$(git rev-parse --short HEAD)
           BUILD_DATE=$(date -u +%Y-%m-%d)
 
-          echo "Built wheel: $WHEEL_FILE"
-          echo "Nightly version: ${NIGHTLY_VERSION}"
           echo "Commit: ${COMMIT_HASH}"
           echo "Build date: ${BUILD_DATE}"
 
-          echo "nightly_version=${NIGHTLY_VERSION}" >> $GITHUB_OUTPUT
-          echo "commit_hash=${COMMIT_HASH}" >> $GITHUB_OUTPUT
-          echo "build_date=${BUILD_DATE}" >> $GITHUB_OUTPUT
+          {
+            echo "commit_hash=${COMMIT_HASH}"
+            echo "build_date=${BUILD_DATE}"
+          } >> "$GITHUB_OUTPUT"
+
+      - name: Repair and smoke-test wheel
+        id: wheel
+        run: |
+          python3 scripts/release/prepare_sglang_wheel.py python/dist \
+            --github-output "$GITHUB_OUTPUT"
 
       - name: Upload wheel artifact
         uses: actions/upload-artifact@v4
@@ -139,7 +143,7 @@ jobs:
           token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }}
           prerelease: true
           body: |
-            Nightly build from commit ${{ github.sha }}
+            Nightly build from commit ${{ needs.build-nightly-wheel.outputs.commit_hash }}
             Build date: ${{ needs.build-nightly-wheel.outputs.build_date }}
             Version: ${{ needs.build-nightly-wheel.outputs.nightly_version }}
           files: |
@@ -147,7 +151,7 @@ jobs:
 
       - name: Clone wheel index repository
         run: |
-          git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl
+          git clone "https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git" sgl-whl
           cd sgl-whl
           git config --local user.name "sglang-bot"
           git config --local user.email "sglangbot@gmail.com"
@@ -162,10 +166,10 @@ jobs:
       - name: Update wheel index
         run: |
           python3 scripts/update_nightly_whl_index.py \
-            --commit-hash ${{ needs.build-nightly-wheel.outputs.commit_hash }} \
-            --nightly-version ${{ needs.build-nightly-wheel.outputs.nightly_version }} \
-            --cuda-version ${{ matrix.cuda_version }} \
-            --build-date ${{ needs.build-nightly-wheel.outputs.build_date }}
+            --commit-hash "${{ needs.build-nightly-wheel.outputs.commit_hash }}" \
+            --nightly-version "${{ needs.build-nightly-wheel.outputs.nightly_version }}" \
+            --cuda-version "${{ matrix.cuda_version }}" \
+            --build-date "${{ needs.build-nightly-wheel.outputs.build_date }}"
 
       - name: Push wheel index
         run: |

.github/workflows/release-pypi-pr.yml

@@ -17,7 +17,8 @@ jobs:
     if: github.repository == 'sgl-project/sglang'
     runs-on: ubuntu-latest
     outputs:
-      wheel_version: ${{ steps.gen_version.outputs.wheel_version }}
+      wheel_version: ${{ steps.wheel.outputs.wheel_version }}
+      wheel_filename: ${{ steps.wheel.outputs.wheel_filename }}
       commit_hash: ${{ steps.gen_version.outputs.commit_hash }}
       build_date: ${{ steps.gen_version.outputs.build_date }}
     steps:
@@ -34,7 +35,7 @@ jobs:
       - name: Generate PR wheel version
         id: gen_version
         run: |
-          LATEST_TAG=$(python3 scripts/release/get_version_tag.py)
+          LATEST_TAG=$(python3 scripts/release/get_version_tag.py --tag-only)
           BASE_VERSION=${LATEST_TAG#v}
           echo "Latest release tag: ${LATEST_TAG}"
 
@@ -57,10 +58,12 @@ jobs:
           echo "Commit: ${COMMIT_HASH}"
           echo "Build date: ${BUILD_DATE}"
 
-          echo "wheel_version=${WHEEL_VERSION}" >> $GITHUB_OUTPUT
-          echo "commit_hash=${COMMIT_HASH}" >> $GITHUB_OUTPUT
-          echo "base_version=${BASE_VERSION}" >> $GITHUB_OUTPUT
-          echo "build_date=${BUILD_DATE}" >> $GITHUB_OUTPUT
+          {
+            echo "wheel_version=${WHEEL_VERSION}"
+            echo "commit_hash=${COMMIT_HASH}"
+            echo "base_version=${BASE_VERSION}"
+            echo "build_date=${BUILD_DATE}"
+          } >> "$GITHUB_OUTPUT"
 
       - name: Update pyproject.toml with PR wheel version
         run: |
@@ -79,19 +82,26 @@ jobs:
 
       - name: Install build dependencies
         run: |
-          cd python
-          pip install build wheel setuptools
+          pip install \
+            auditwheel build patchelf wheel "setuptools>=61.0" \
+            "setuptools-rust>=1.11" "setuptools-scm>=8.0" "torch==2.13.0"
 
       - name: Build wheel
         run: |
           cd python
           cp ../README.md ../LICENSE .
-          python3 -m build --wheel
+          python3 -m build --wheel --no-isolation
 
           # List built wheels
           echo "Built wheel:"
           ls -lh dist/
 
+      - name: Repair and smoke-test wheel
+        id: wheel
+        run: |
+          python3 scripts/release/prepare_sglang_wheel.py python/dist \
+            --github-output "$GITHUB_OUTPUT"
+
       - name: Upload wheel artifact
         uses: actions/upload-artifact@v4
         with:
@@ -127,7 +137,7 @@ jobs:
           prerelease: true
           body: |
             PR wheel build from PR #${{ inputs.pr_number }}
-            Commit: ${{ github.sha }}
+            Commit: ${{ needs.build-pr-wheel.outputs.commit_hash }}
             Build date: ${{ needs.build-pr-wheel.outputs.build_date }}
             Version: ${{ needs.build-pr-wheel.outputs.wheel_version }}
 
@@ -143,14 +153,14 @@ jobs:
 
             **Direct installation:**
             ```bash
-            pip install https://github.com/sgl-project/whl/releases/download/pr-${{ inputs.pr_number }}-${{ needs.build-pr-wheel.outputs.build_date }}-${{ needs.build-pr-wheel.outputs.commit_hash }}/sglang-${{ needs.build-pr-wheel.outputs.wheel_version }}-py3-none-any.whl
+            pip install https://github.com/sgl-project/whl/releases/download/pr-${{ inputs.pr_number }}-${{ needs.build-pr-wheel.outputs.build_date }}-${{ needs.build-pr-wheel.outputs.commit_hash }}/${{ needs.build-pr-wheel.outputs.wheel_filename }}
             ```
           files: |
             dist/*.whl
 
       - name: Clone wheel index repository
         run: |
-          git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl
+          git clone "https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git" sgl-whl
           cd sgl-whl
           git config --local user.name "sglang-bot"
           git config --local user.email "sglangbot@gmail.com"
@@ -165,10 +175,10 @@ jobs:
       - name: Update wheel index
         run: |
           python3 scripts/update_pr_whl_index.py \
-            --pr-number ${{ inputs.pr_number }} \
-            --commit-hash ${{ needs.build-pr-wheel.outputs.commit_hash }} \
-            --wheel-version ${{ needs.build-pr-wheel.outputs.wheel_version }} \
-            --build-date ${{ needs.build-pr-wheel.outputs.build_date }}
+            --pr-number "${{ inputs.pr_number }}" \
+            --commit-hash "${{ needs.build-pr-wheel.outputs.commit_hash }}" \
+            --wheel-version "${{ needs.build-pr-wheel.outputs.wheel_version }}" \
+            --build-date "${{ needs.build-pr-wheel.outputs.build_date }}"
 
       - name: Push wheel index
         run: |

.github/workflows/release-pypi.yml

@@ -61,25 +61,23 @@ jobs:
         run: |
           cd python
           cp ../README.md ../LICENSE .
-          pip install build wheel setuptools setuptools-scm setuptools-rust
+          pip install \
+            build wheel "setuptools>=61.0" "setuptools-rust>=1.11" \
+            "setuptools-scm>=8.0" "torch==2.13.0"
           if [ -n "$RELEASE_VERSION" ]; then
             export SETUPTOOLS_SCM_PRETEND_VERSION="${RELEASE_VERSION#v}"
             echo "Pinning wheel version to $SETUPTOOLS_SCM_PRETEND_VERSION"
           fi
-          python3 -m build --wheel
+          python3 -m build --wheel --no-isolation
 
       # PyPI rejects plain `linux_x86_64` / `linux_aarch64` platform tags;
       # auditwheel rewrites the wheel's platform tag to a `manylinux_*` tag
       # and bundles any external native deps. The runner's glibc determines
       # the lowest acceptable manylinux policy.
-      - name: Repair wheel for manylinux
+      - name: Repair and smoke-test wheel
         run: |
-          cd python
           pip install auditwheel patchelf
-          mkdir -p dist-repaired
-          python3 -m auditwheel repair dist/*.whl -w dist-repaired/
-          rm dist/*.whl
-          mv dist-repaired/*.whl dist/
+          python3 scripts/release/prepare_sglang_wheel.py python/dist
 
       - name: Upload artifacts
         uses: actions/upload-artifact@v4

.github/workflows/seed-rust-ext-cache.yml

@@ -10,8 +10,11 @@ on:
     paths:
       - 'rust/**'
       - 'python/setup.py'
+      - 'python/pyproject.toml'
+      - 'python/sglang/srt/rust_extensions/torch_build.py'
       # The key's prefix lives in this file, so a bump there also moves the key.
       - '.github/workflows/_pr-test-rust-ext-build.yml'
+      - 'scripts/ci/utils/stage_rust_ext_modules.sh'
   workflow_dispatch:
 
 # Only the newest merge needs to seed; earlier ones are already stale.

.gitignore

@@ -231,6 +231,7 @@ work_dirs/
 # Rust lib
 Cargo.lock
 !rust/Cargo.lock
+!rust/mem-cache/Cargo.lock
 
 # Generated vision test fixtures (regenerate with: python scripts/generate_vision_golden.py)
 sgl-model-gateway/tests/fixtures/golden/

.pre-commit-config.yaml

@@ -151,13 +151,13 @@ repos:
         # cover sglang-mm's PyO3 bindings + rayon fan-out — both sit behind
         # non-default features, so `--workspace` alone never compiles them. protoc
         # is not required: sglang-grpc's build.rs falls back to a vendored binary.
-        entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings && cargo clippy -p sglang-mm --features python,parallel --lib -- -D warnings'
+        entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings && cargo clippy -p sglang-mm --features python,parallel --lib -- -D warnings && cargo clippy --manifest-path mem-cache/Cargo.toml --all-targets --no-default-features --features tch/doc-only -- -D warnings'
         language: system
         files: ^rust/.*\.rs$
         pass_filenames: false
       - id: rustfmt-rust-workspace
         name: rustfmt rust/ workspace
-        entry: bash -c 'cd rust && cargo fmt'
+        entry: bash -c 'cd rust && cargo fmt && cargo fmt --manifest-path mem-cache/Cargo.toml'
         language: system
         files: ^rust/.*\.rs$
         pass_filenames: false

docker/Dockerfile

@@ -671,6 +671,8 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     && ( if [ -f python/kernels.lock ]; then mv python/kernels.lock /root/.cache/sglang/; fi ) \
     && ( find /opt/sglang/lib/python3.12/site-packages -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true )
 
+ENV SGLANG_RUST_BUILD_MODE=never
+
 
 # Install pre-built gateway artifacts from parallel builder
 COPY --from=gateway_builder /build/sgl-model-gateway-bin /opt/sglang/bin/sgl-model-gateway
@@ -723,7 +725,8 @@ ARG GDRCOPY_VERSION=2.5.1
 
 ENV DEBIAN_FRONTEND=noninteractive \
     CUDA_HOME=/usr/local/cuda \
-    GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/
+    GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \
+    SGLANG_RUST_BUILD_MODE=never
 
 # Add GKE default lib and bin locations + CUDA compiler paths for FlashInfer JIT
 ENV PATH="${PATH}:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/cuda/nvvm/bin" \

docker/Dockerfile.cu134

@@ -944,6 +944,9 @@ ARG CUDA_VERSION
 ARG SGL_VERSION
 ARG USE_LATEST_SGLANG
 
+# TODO(Jialin): Set SGLANG_RUST_BUILD_MODE=never after Rust TreeCore supports
+# the PyTorch 2.15 nightly used by this preview image.
+
 WORKDIR /sgl-workspace
 
 COPY --from=local_src /src /tmp/local_src

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,

rust/Cargo.toml

@@ -5,6 +5,12 @@ members = [
     "sglang-mm",
     "sglang-server"
 ]
+exclude = ["mem-cache"]
+
+[workspace.metadata.sglang]
+# Extension crates that intentionally cannot share this workspace's dependency
+# graph (mem-cache currently uses the tch-compatible PyO3 0.22 API).
+extension-manifests = ["mem-cache/Cargo.toml"]
 
 [workspace.package]
 version = "0.1.0"

rust/mem-cache/.gitignore

@@ -0,0 +1 @@
+target/

rust/mem-cache/Cargo.lock

@@ -0,0 +1,905 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "adler2"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
+
+[[package]]
+name = "aes"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
+dependencies = [
+ "cfg-if",
+ "cipher",
+ "cpufeatures",
+]
+
+[[package]]
+name = "allocator-api2"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
+
+[[package]]
+name = "anyhow"
+version = "1.0.104"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
+
+[[package]]
+name = "autocfg"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
+
+[[package]]
+name = "base64ct"
+version = "1.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
+
+[[package]]
+name = "block-buffer"
+version = "0.10.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
+dependencies = [
+ "generic-array",
+]
+
+[[package]]
+name = "byteorder"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
+
+[[package]]
+name = "bzip2"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8"
+dependencies = [
+ "bzip2-sys",
+ "libc",
+]
+
+[[package]]
+name = "bzip2-sys"
+version = "0.1.13+1.0.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14"
+dependencies = [
+ "cc",
+ "pkg-config",
+]
+
+[[package]]
+name = "cc"
+version = "1.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273"
+dependencies = [
+ "find-msvc-tools",
+ "jobserver",
+ "libc",
+ "shlex",
+]
+
+[[package]]
+name = "cfg-if"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
+
+[[package]]
+name = "cipher"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
+dependencies = [
+ "crypto-common",
+ "inout",
+]
+
+[[package]]
+name = "constant_time_eq"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc"
+
+[[package]]
+name = "cpufeatures"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "crc32fast"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550"
+dependencies = [
+ "cfg-if",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17"
+
+[[package]]
+name = "crunchy"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
+
+[[package]]
+name = "crypto-common"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a"
+dependencies = [
+ "generic-array",
+ "typenum",
+]
+
+[[package]]
+name = "deranged"
+version = "0.5.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
+
+[[package]]
+name = "digest"
+version = "0.10.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
+dependencies = [
+ "block-buffer",
+ "crypto-common",
+ "subtle",
+]
+
+[[package]]
+name = "equivalent"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
+
+[[package]]
+name = "find-msvc-tools"
+version = "0.1.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890"
+
+[[package]]
+name = "flate2"
+version = "1.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide",
+]
+
+[[package]]
+name = "foldhash"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
+
+[[package]]
+name = "generic-array"
+version = "0.14.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
+dependencies = [
+ "typenum",
+ "version_check",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "r-efi",
+]
+
+[[package]]
+name = "half"
+version = "2.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
+dependencies = [
+ "cfg-if",
+ "crunchy",
+ "zerocopy",
+]
+
+[[package]]
+name = "hashbrown"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
+dependencies = [
+ "allocator-api2",
+ "equivalent",
+ "foldhash",
+]
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "hmac"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
+dependencies = [
+ "digest",
+]
+
+[[package]]
+name = "indoc"
+version = "2.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
+dependencies = [
+ "rustversion",
+]
+
+[[package]]
+name = "inout"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
+dependencies = [
+ "generic-array",
+]
+
+[[package]]
+name = "itoa"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
+
+[[package]]
+name = "jobserver"
+version = "0.1.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3"
+dependencies = [
+ "getrandom 0.4.3",
+ "libc",
+]
+
+[[package]]
+name = "lazy_static"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+
+[[package]]
+name = "libc"
+version = "0.2.189"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
+
+[[package]]
+name = "matrixmultiply"
+version = "0.3.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7"
+dependencies = [
+ "autocfg",
+ "rawpointer",
+]
+
+[[package]]
+name = "mem_cache"
+version = "0.1.0"
+dependencies = [
+ "hashbrown",
+ "pyo3",
+ "sha2",
+ "tch",
+ "thiserror",
+]
+
+[[package]]
+name = "memchr"
+version = "2.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
+
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "miniz_oxide"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c"
+dependencies = [
+ "adler2",
+ "simd-adler32",
+]
+
+[[package]]
+name = "ndarray"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-conv"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441"
+
+[[package]]
+name = "num-integer"
+version = "0.1.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "password-hash"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7676374caaee8a325c9e7a2ae557f216c5563a171d6997b0ef8a65af35147700"
+dependencies = [
+ "base64ct",
+ "rand_core",
+ "subtle",
+]
+
+[[package]]
+name = "pbkdf2"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917"
+dependencies = [
+ "digest",
+ "hmac",
+ "password-hash",
+ "sha2",
+]
+
+[[package]]
+name = "pkg-config"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548"
+
+[[package]]
+name = "portable-atomic"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "powerfmt"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
+
+[[package]]
+name = "ppv-lite86"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
+dependencies = [
+ "zerocopy",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.107"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "r-efi"
+version = "6.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
+
+[[package]]
+name = "rand"
+version = "0.8.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c"
+dependencies = [
+ "libc",
+ "rand_chacha",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+dependencies = [
+ "getrandom 0.2.17",
+]
+
+[[package]]
+name = "rawpointer"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
+
+[[package]]
+name = "rustversion"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f"
+
+[[package]]
+name = "safetensors"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d93279b86b3de76f820a8854dd06cbc33cfa57a417b19c47f6a25280112fb1df"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "serde"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
+dependencies = [
+ "serde_core",
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_core"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.4",
+]
+
+[[package]]
+name = "serde_json"
+version = "1.0.151"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14"
+dependencies = [
+ "itoa",
+ "memchr",
+ "serde",
+ "serde_core",
+ "zmij",
+]
+
+[[package]]
+name = "sha1"
+version = "0.10.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+]
+
+[[package]]
+name = "sha2"
+version = "0.10.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+ "sha2-asm",
+]
+
+[[package]]
+name = "sha2-asm"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b845214d6175804686b2bd482bcffe96651bb2d1200742b712003504a2dac1ab"
+dependencies = [
+ "cc",
+]
+
+[[package]]
+name = "shlex"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba"
+
+[[package]]
+name = "simd-adler32"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea"
+
+[[package]]
+name = "subtle"
+version = "2.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
+
+[[package]]
+name = "syn"
+version = "2.0.119"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn"
+version = "3.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
+[[package]]
+name = "tch"
+version = "0.24.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0d3f84a069d8ba16dbf720b61e8bf131d90ffb8e958a664eae8e4993c5c2fa6f"
+dependencies = [
+ "half",
+ "lazy_static",
+ "libc",
+ "ndarray",
+ "rand",
+ "safetensors",
+ "thiserror",
+ "torch-sys",
+ "zip",
+]
+
+[[package]]
+name = "thiserror"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
+dependencies = [
+ "thiserror-impl",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "time"
+version = "0.3.55"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134"
+dependencies = [
+ "deranged",
+ "num-conv",
+ "powerfmt",
+ "serde_core",
+ "time-core",
+]
+
+[[package]]
+name = "time-core"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109"
+
+[[package]]
+name = "torch-sys"
+version = "0.24.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4ba78777379cf09aaa79708c63e477cf0f95e021d04360c6821f1a9f56173f7"
+dependencies = [
+ "anyhow",
+ "cc",
+ "libc",
+ "zip",
+]
+
+[[package]]
+name = "typenum"
+version = "1.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
+
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
+[[package]]
+name = "version_check"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
+
+[[package]]
+name = "wasi"
+version = "0.11.1+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
+
+[[package]]
+name = "zerocopy"
+version = "0.8.56"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb"
+dependencies = [
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.8.56"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "zip"
+version = "0.6.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261"
+dependencies = [
+ "aes",
+ "byteorder",
+ "bzip2",
+ "constant_time_eq",
+ "crc32fast",
+ "crossbeam-utils",
+ "flate2",
+ "hmac",
+ "pbkdf2",
+ "sha1",
+ "time",
+ "zstd",
+]
+
+[[package]]
+name = "zmij"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
+
+[[package]]
+name = "zstd"
+version = "0.11.2+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4"
+dependencies = [
+ "zstd-safe",
+]
+
+[[package]]
+name = "zstd-safe"
+version = "5.0.2+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d2a5585e04f9eea4b2a3d1eca508c4dee9592a89ef6f450c11719da0726f4db"
+dependencies = [
+ "libc",
+ "zstd-sys",
+]
+
+[[package]]
+name = "zstd-sys"
+version = "2.0.16+zstd.1.5.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748"
+dependencies = [
+ "cc",
+ "pkg-config",
+]

rust/mem-cache/Cargo.toml

@@ -0,0 +1,40 @@
+[package]
+name = "mem_cache"
+version = "0.1.0"
+edition = "2024"
+description = "Rust tree core for the Unified Radix Cache"
+license = "Apache-2.0"
+repository = "https://github.com/sgl-project/sglang"
+
+[package.metadata.sglang]
+python-module = "sglang.srt.mem_cache.rust_tree_core.mem_cache"
+debug = false
+features = ["python-extension"]
+torch-compat-header = "torch_2_13_compat.h"
+
+[lib]
+name = "mem_cache"
+crate-type = ["cdylib"]
+
+# Keep this crate in its own workspace until its PyO3 version can be aligned
+# with the root workspace without conflicting native `links = "python"` crates.
+[workspace]
+
+[dependencies]
+hashbrown = "0.16"
+pyo3 = { version = "0.22", optional = true }
+sha2 = "0.10"
+tch = "=0.24.0"
+thiserror = "1"
+
+# Hardware SHA extensions; identical digests.
+[target.'cfg(all(target_arch = "aarch64", target_os = "linux"))'.dependencies]
+sha2 = { version = "0.10", features = ["asm"] }
+
+[features]
+# Keep the native core as the default so workspace tests do not link a Python
+# extension. Wheel and source builds select python-extension through package
+# metadata; shared white-box tests additionally select inspection.
+default = []
+python-extension = ["dep:pyo3", "pyo3/extension-module", "tch/python-extension"]
+inspection = []

rust/mem-cache/README.md

@@ -0,0 +1,42 @@
+# mem-cache
+
+Rust tree core for the Unified Radix Cache, covering Full attention, sliding window attention, and Mamba components. It implements the tree side of the `UnifiedTreeCoreInterface` split — match/insert walks, node arena, locks, eviction walks, HiCache backup/load-back specs, and KV events — behind a PyO3 binding, while the cache orchestration stays in Python.
+
+## Usage
+
+Select the backend with:
+
+```bash
+SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND=rust
+```
+
+SGLang wheels bundle the production extension. A source checkout falls back to
+the shared fingerprinted Rust-extension cache; it never writes a shared object
+into the Python package. LibTorch and the Python headers come from the running
+interpreter's PyTorch install. PyTorch 2.11 through 2.13 are accepted explicitly,
+and `torch_2_13_compat.h` covers two alignment APIs removed in PyTorch 2.13.
+
+## Development
+
+```bash
+# Build (libtorch from the installed torch package):
+cd rust/mem-cache
+LIBTORCH_USE_PYTORCH=1 \
+  LIBTORCH_BYPASS_VERSION_CHECK=1 \
+  CXXFLAGS="-include $PWD/torch_2_13_compat.h" \
+  cargo build --release --locked --features python-extension
+
+# Native tests do not enable pyo3's extension-module feature:
+TORCH_ROOT=$(python3 -c 'import pathlib, torch; print(pathlib.Path(torch.__file__).parent)')
+LIBTORCH_USE_PYTORCH=1 LIBTORCH_BYPASS_VERSION_CHECK=1 \
+  CXXFLAGS="-include $PWD/torch_2_13_compat.h" \
+  LD_LIBRARY_PATH="$TORCH_ROOT/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" \
+  cargo test --locked
+```
+
+The `inspection` Cargo feature adds white-box methods for the shared Python/Rust
+cache suite. Production wheels do not enable it.
+
+Unit tests live in `src/tests/`, mirroring the source layout one file per module (wired via `#[cfg(test)] #[path = ...]`), so implementation files stay free of inline test blocks.
+
+Supported component sets are `[Full]`, `[Full, SWA]`, `[Full, Mamba]`, and `[Full, SWA, Mamba]`.

rust/mem-cache/src/components/full.rs

@@ -0,0 +1,493 @@
+//! FULL attention component driver: overrides the methods FULL customizes and inherits
+//! the rest from the `TreeComponent` defaults.
+
+use std::cmp::Reverse;
+use std::collections::{BinaryHeap, HashMap, HashSet};
+
+use tch::{Kind, Tensor};
+
+use crate::components::TreeComponent;
+use crate::components::{ComponentType, FULL};
+use crate::node::ChildKeyType;
+use crate::node::Node;
+use crate::node::{NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx};
+use crate::unified_lru_list::PriorityKey;
+use crate::unified_tree_core::{
+    CacheAction, CacheTransferPhase, DecLockRefParams, EvictLayer, IncLockRefResult, InsertResult,
+    MatchPrefixParams, MatchResult, PoolName, PoolTransfer, PoolTransferResult, UnifiedTreeCore,
+};
+
+/// FULL attention component driver; owns the FULL device/host value slots.
+pub struct FullComponent;
+
+impl FullComponent {
+    /// The component's device value slot.
+    pub const DEVICE: ValueSlotIdx = ValueSlotIdx::device(FULL);
+    /// The component's host value slot.
+    pub const HOST: ValueSlotIdx = ValueSlotIdx::host(FULL);
+}
+
+impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
+    fn component_type(&self) -> ComponentType {
+        FULL
+    }
+
+    fn create_match_validator(
+        &self,
+        _tree_core: &UnifiedTreeCore<K>,
+        match_device_only: bool,
+    ) -> Box<dyn FnMut(&UnifiedTreeCore<K>, NodeIdx_) -> bool> {
+        // Device value present -> always a boundary; otherwise a backuped (host-resident)
+        // node also matches, unless the match is restricted to device.
+        Box::new(move |tree_core: &UnifiedTreeCore<K>, node_id: NodeIdx_| {
+            let node = tree_core.arena.node(node_id);
+            node.has_device_value(FULL) || (!match_device_only && node.has_host_value(FULL))
+        })
+    }
+
+    fn finalize_match_result_in_tree_core(
+        &self,
+        tree_core: &UnifiedTreeCore<K>,
+        mut result: MatchResult,
+        params: &MatchPrefixParams<'_, K>,
+        value_chunks: &[Tensor],
+        best_value_len: usize,
+    ) -> MatchResult {
+        // Compute Full KV host hit length: walk from last_host_node up to
+        // last_device_node, summing host_value lengths of evicted nodes.
+        let mut kv_host_hit = 0;
+        let mut node_idx = tree_core.arena.resolve(result.best_match_node_id);
+        let last_device_idx = tree_core.arena.resolve(result.last_device_node_id);
+        while node_idx != last_device_idx {
+            let node = tree_core.arena.node(node_idx);
+            let parent = node.try_parent().unwrap_or_else(|| {
+                panic!(
+                    "finalize walk from best_match_node {} hit root {} before \
+                     last_device_node {}",
+                    result.best_match_node_id, node.id, result.last_device_node_id
+                )
+            });
+            kv_host_hit += node.host_value_len(FULL);
+            node_idx = parent;
+        }
+        if kv_host_hit > 0 {
+            result.host_hit_length = result.host_hit_length.max(kv_host_hit);
+        }
+        result
+    }
+
+    fn redistribute_on_node_split(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        new_parent_id: NodeIdx_,
+        child_id: NodeIdx_,
+    ) {
+        let (new_parent, child) = tree_core.arena.node_pair_mut(new_parent_id, child_id);
+        let split_len = new_parent.key.atom_len() as i64;
+        new_parent.copy_device_lock_ref(FULL, child);
+        if child.has_device_value(FULL) {
+            Node::redistribute_child_device_value(new_parent, child, FULL, split_len);
+        }
+        if child.has_host_value(FULL) {
+            Node::redistribute_child_host_value(new_parent, child, FULL, split_len);
+        }
+    }
+
+    fn evict_component(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        target: EvictLayer,
+    ) -> (usize, usize) {
+        let node = tree_core.arena.node_mut(node_id);
+        let mut freed = 0;
+        let mut host_freed = 0;
+        if target.contains(EvictLayer::Device) && node.has_device_value(FULL) {
+            let value = node.device_value(FULL);
+            freed = node.device_value_len(FULL);
+            device_frees
+                .entry(FULL)
+                .or_default()
+                .push(value.shallow_clone());
+            // NOTE: cd.value = None is deferred to _cascade_evict (Full as trigger)
+            // because SWA's free_swa still needs to read Full.value.
+        }
+        if target.contains(EvictLayer::Host) && node.has_host_value(FULL) {
+            host_freed = node.host_value_len(FULL);
+            host_frees
+                .entry(FULL)
+                .or_default()
+                .push(node.take_host_value(FULL));
+        }
+        if freed > 0 {
+            tree_core.dec_evictable_size(FULL, freed);
+        }
+        (freed, host_freed)
+    }
+
+    fn eviction_priority(&self, is_leaf: bool) -> i64 {
+        if is_leaf { 0 } else { 2 }
+    }
+
+    fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<K>, request_cnt: usize) {
+        tree_core.set_evict_device_start(FULL, request_cnt);
+        tree_core.full_evict_device_heap.clear();
+        let arena = &tree_core.arena;
+        let strategy = &tree_core.eviction_strategy;
+        tree_core.full_evict_device_heap.extend(
+            tree_core
+                .evictable_device_leaves
+                .iter()
+                .map(|id| Reverse((strategy.get_priority(arena.node(id)), id))),
+        );
+    }
+
+    fn evict_device_next_node(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        tracker: &mut HashMap<ComponentType, usize>,
+        _device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        _host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) -> Option<NodeIdx_> {
+        let ct = FULL;
+        assert!(
+            tree_core.component_state(FULL).is_evict_device_ongoing,
+            "Full device eviction not started"
+        );
+        // Re-admit the previously returned leaf's parent once it became a
+        // D-leaf; the parent id was captured at return time because the leaf
+        // itself may have been freed by the eviction in between.
+        if let Some(last_node_parent) = tree_core.component_state(FULL).evict_device_cursor
+            && tree_core.evictable_device_leaves.contains(last_node_parent)
+        {
+            let key = tree_core
+                .eviction_strategy
+                .get_priority(tree_core.arena.node(last_node_parent));
+            tree_core
+                .full_evict_device_heap
+                .push(Reverse((key, last_node_parent)));
+        }
+        tree_core.component_state_mut(FULL).evict_device_cursor = None;
+        // The budget only advances between calls (the driver's evictions fill
+        // the tracker), so it gates the walk once up front.
+        if tracker[&ct] >= tree_core.component_state(FULL).evict_device_request_cnt {
+            return None;
+        }
+        while let Some(Reverse((_, x))) = tree_core.full_evict_device_heap.pop() {
+            if !tree_core.evictable_device_leaves.contains(x) {
+                continue;
+            }
+            let last_node_parent = tree_core.arena.node(x).try_parent();
+            tree_core.component_state_mut(FULL).evict_device_cursor = last_node_parent;
+            return Some(x);
+        }
+        None
+    }
+
+    fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<K>) {
+        tree_core.set_evict_device_end(FULL);
+        tree_core.full_evict_device_heap.clear();
+    }
+
+    fn reclaim_coexisting_host_values(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        num_tokens: usize,
+        tracker: &mut HashMap<ComponentType, usize>,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+        for spare_imminent_demotes in [true, false] {
+            if tracker[&FULL] >= num_tokens {
+                break;
+            }
+            let candidates: Vec<NodeIdx_> = tree_core.full_coexisting_host_nodes.iter().collect();
+            for node_id in candidates {
+                if tracker[&FULL] >= num_tokens {
+                    break;
+                }
+                let node = tree_core.arena.node(node_id);
+                if !node.has_device_value(FULL) || !node.has_host_value(FULL) {
+                    tree_core.full_coexisting_host_nodes.discard(node_id);
+                    continue;
+                }
+                if spare_imminent_demotes && tree_core.evictable_device_leaves.contains(node_id) {
+                    continue;
+                }
+                if !tree_core.can_reclaim_coexisting_host_value_(node_id, FULL) {
+                    continue;
+                }
+                tree_core.release_coexisting_host_value_(
+                    node_id,
+                    FULL,
+                    tracker,
+                    device_frees,
+                    host_frees,
+                );
+                tree_core.full_coexisting_host_nodes.discard(node_id);
+            }
+        }
+    }
+
+    /// Evict host leaves to free KV host pool space.
+    fn drive_host_eviction(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        num_tokens: usize,
+        tracker: &mut HashMap<ComponentType, usize>,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+        let ct = FULL;
+        let arena = &tree_core.arena;
+        let strategy = &tree_core.eviction_strategy;
+        let mut heap: BinaryHeap<Reverse<(PriorityKey, NodeIdx_)>> = tree_core
+            .evictable_host_leaves
+            .iter()
+            .map(|id| Reverse((strategy.get_priority(arena.node(id)), id)))
+            .collect();
+        while tracker[&ct] < num_tokens {
+            let Some(Reverse((_, x))) = heap.pop() else {
+                break;
+            };
+            if !tree_core.evictable_host_leaves.contains(x) {
+                continue;
+            }
+            // The parent id is captured before the eviction frees the leaf.
+            let parent = tree_core.arena.node(x).try_parent();
+            tree_core.evict_host_leaf_(x, tracker, device_frees, host_frees);
+            if let Some(parent) = parent
+                && tree_core.evictable_host_leaves.contains(parent)
+            {
+                let key = tree_core
+                    .eviction_strategy
+                    .get_priority(tree_core.arena.node(parent));
+                heap.push(Reverse((key, parent)));
+            }
+        }
+    }
+
+    fn acquire_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        mut result: IncLockRefResult,
+        lock_host: bool,
+    ) -> IncLockRefResult {
+        let ct = FULL;
+
+        // Only the last host node needs to be protected.
+        if lock_host {
+            let node = tree_core.arena.node_mut(node_id);
+            // write_back mode: the anchor may be device-only (no host_value); pin it anyway.
+            if !node.has_host_value(FULL) && !tree_core.is_write_back {
+                return result;
+            }
+            node.inc_host_lock_ref(FULL);
+            tree_core.update_evictable_leaf_sets_(node_id);
+            return result;
+        }
+
+        // Skip the bottom evicted segment, recording it for the matching release.
+        let on_boundary = |node: &Node<K>| node.is_root() || node.has_device_value(FULL);
+        let mut cur = node_id;
+        let mut node = tree_core.arena.node(cur);
+        if !on_boundary(node) {
+            let skip_lock_node_ids = result.skip_lock_node_ids.entry(ct).or_default();
+            loop {
+                skip_lock_node_ids.insert(node.id);
+                cur = node.parent();
+                node = tree_core.arena.node(cur);
+                if on_boundary(node) {
+                    break;
+                }
+            }
+        }
+
+        // Lock the device-on segment up to the root.
+        let mut delta = 0;
+        loop {
+            let node = tree_core.arena.node_mut(cur);
+            if node.is_root() {
+                break;
+            }
+            assert!(
+                node.has_device_value(FULL),
+                "FULL invariant broken: evicted ancestor {cur} above device-on segment"
+            );
+            let parent = node.parent();
+            let newly_locked_len = if node.device_lock_ref(FULL) == 0 {
+                Some(node.device_value_len(FULL))
+            } else {
+                None
+            };
+            node.inc_device_lock_ref(FULL);
+            if let Some(key_len) = newly_locked_len {
+                tree_core.dec_evictable_size(FULL, key_len);
+                tree_core.inc_protected_size(FULL, key_len);
+                delta += key_len;
+            }
+            tree_core.evictable_device_leaves.discard(cur);
+            cur = parent;
+        }
+        result.delta = Some(delta);
+        result
+    }
+
+    fn release_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        params: Option<&DecLockRefParams>,
+        lock_host: bool,
+    ) {
+        let ct = FULL;
+
+        if lock_host {
+            let node = tree_core.arena.node_mut(node_id);
+            if node.host_lock_ref(FULL) == 0 {
+                return;
+            }
+            // Mirror of `acquire`. write_back uses a pure counter.
+            if !node.has_host_value(FULL) && !tree_core.is_write_back {
+                return;
+            }
+            node.dec_host_lock_ref(FULL);
+            tree_core.update_evictable_leaf_sets_(node_id);
+            return;
+        }
+
+        let empty = HashSet::new();
+        let skip_lock_node_ids = params
+            .and_then(|p| p.skip_lock_node_ids.get(&ct))
+            .unwrap_or(&empty);
+        let mut cur = node_id;
+        loop {
+            let node = tree_core.arena.node_mut(cur);
+            if node.is_root() {
+                break;
+            }
+            let parent = node.parent();
+            if skip_lock_node_ids.contains(&node.id) {
+                cur = parent;
+                continue;
+            }
+            assert!(
+                node.has_device_value(FULL),
+                "release_component_lock: node {cur} has no FULL device value"
+            );
+            let old_lock_ref = node.device_lock_ref(FULL);
+            assert!(
+                old_lock_ref > 0,
+                "release_component_lock: node {cur} is not locked"
+            );
+            let newly_unlocked_len = if old_lock_ref == 1 {
+                Some(node.device_value_len(FULL))
+            } else {
+                None
+            };
+            node.dec_device_lock_ref(FULL);
+            if let Some(key_len) = newly_unlocked_len {
+                tree_core.dec_protected_size(FULL, key_len);
+                tree_core.inc_evictable_size(FULL, key_len);
+                tree_core.update_evictable_leaf_sets_(cur);
+            }
+            cur = parent;
+        }
+    }
+
+    fn build_hicache_transfers(
+        &self,
+        tree_core: &UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        phase: CacheTransferPhase,
+        _mamba_pool_idx: Option<Tensor>,
+        _host_indices: Option<Tensor>,
+        _token_ids: Option<&[i64]>,
+        _prefetch_tokens: usize,
+        _last_hash: Option<&str>,
+    ) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
+        Ok(match phase {
+            // Full KV backup is handled by the main flow
+            // (cache_controller.write on host_value directly).
+            // No extra PoolTransfer needed.
+            CacheTransferPhase::BackupHost => None,
+            CacheTransferPhase::LoadBack => {
+                // `node` is best_match_node. FULL device evict only from leaves,
+                // so once we hit a device-on node, everything above is also device-on.
+                let mut backed_up: Vec<Tensor> = Vec::new();
+                let mut nodes_to_load: Vec<NodeId> = Vec::new();
+                let mut cur = tree_core.arena.node(node_id);
+                while cur.evicted() {
+                    backed_up.push(cur.host_value(FULL).shallow_clone());
+                    nodes_to_load.push(cur.id);
+                    cur = tree_core.arena.node(cur.parent());
+                }
+                backed_up.reverse();
+                nodes_to_load.reverse();
+                let host_indices = if backed_up.is_empty() {
+                    Tensor::empty([0], (Kind::Int64, tch::Device::Cpu))
+                } else {
+                    Tensor::cat(&backed_up, 0)
+                };
+                Some(vec![PoolTransfer {
+                    name: PoolName::Kv,
+                    host_indices: Some(host_indices),
+                    nodes_to_load: Some(nodes_to_load),
+                    ..Default::default()
+                }])
+            }
+            CacheTransferPhase::BackupStorage | CacheTransferPhase::Prefetch => None,
+        })
+    }
+
+    fn commit_hicache_transfer(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        phase: CacheTransferPhase,
+        transfers: Vec<PoolTransfer>,
+        cache_actions: &mut Vec<CacheAction>,
+        insert_result: Option<&mut InsertResult>,
+        pool_storage_result: Option<&PoolTransferResult>,
+    ) {
+        match phase {
+            CacheTransferPhase::BackupHost => {
+                if let Some(transfer) = transfers.first()
+                    && let Some(host_indices) = &transfer.host_indices
+                {
+                    tree_core
+                        .arena
+                        .set_host_value(node_id, FULL, host_indices.copy());
+                }
+            }
+            CacheTransferPhase::LoadBack => {
+                if let Some(transfer) = transfers.first()
+                    && let Some(device_indices) = &transfer.device_indices
+                {
+                    let mut offset = 0i64;
+                    for &loaded_id in transfer.nodes_to_load.iter().flatten() {
+                        let loaded_idx = tree_core.arena.resolve(loaded_id);
+                        let loaded = tree_core.arena.node_mut(loaded_idx);
+                        let n_len = loaded.host_value_len(FULL) as i64;
+                        loaded
+                            .set_device_value(FULL, device_indices.narrow(0, offset, n_len).copy());
+                        offset += n_len;
+                        // Full uses leaf sets, not LRU.
+                        tree_core.inc_evictable_size(FULL, n_len as usize);
+                        tree_core.update_evictable_leaf_sets_(loaded_idx);
+                    }
+                }
+                tree_core.update_evictable_leaf_sets_(node_id);
+            }
+            // The Full component has no storage-phase commits.
+            CacheTransferPhase::BackupStorage | CacheTransferPhase::Prefetch => {}
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "../tests/components/full.rs"]
+mod tests;

rust/mem-cache/src/components/mamba.rs

@@ -0,0 +1,751 @@
+//! Mamba (SSM state) component driver: overrides the methods Mamba customizes
+//! and inherits the rest from the `TreeComponent` defaults.
+//! Mamba data is per-leaf single-slot state; sizes count slots, not tokens.
+
+use std::collections::HashMap;
+
+use tch::Tensor;
+
+use crate::components::TreeComponent;
+use crate::components::{ComponentType, MAMBA};
+use crate::node::ChildKeyType;
+use crate::node::Node;
+use crate::node::{NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx};
+use crate::unified_tree_core::{
+    CacheAction, CacheInitParams, CacheTransferPhase, DecLockRefParams, EvictLayer,
+    IncLockRefResult, InsertParams, InsertResult, LRURefreshPhase, MatchPrefixParams, MatchResult,
+    PoolHitPolicy, PoolName, PoolTransfer, PoolTransferResult, UnifiedTreeCore,
+};
+
+/// Mamba component driver; owns the Mamba device/host value slots.
+pub struct MambaComponent {
+    /// Joint chunk/tree-page alignment for the mamba branching seqlen.
+    mamba_checkpoint_grid: usize,
+    /// Per-root-path cap on cached Mamba states; None means unlimited.
+    mamba_max_states_per_path: Option<usize>,
+}
+
+impl MambaComponent {
+    /// The component's device value slot.
+    pub const DEVICE: ValueSlotIdx = ValueSlotIdx::device(MAMBA);
+    /// The component's host value slot.
+    pub const HOST: ValueSlotIdx = ValueSlotIdx::host(MAMBA);
+}
+
+impl MambaComponent {
+    /// Build the driver from the tree's init params.
+    pub fn new(params: &CacheInitParams) -> Self {
+        let mamba_cache_chunk_size = params
+            .mamba_cache_chunk_size
+            .expect("the Mamba component requires mamba_cache_chunk_size");
+        MambaComponent {
+            // A donated checkpoint must land on both the model's chunk grid and
+            // a radix-node boundary. `params.page_size` is already widened by DCP.
+            mamba_checkpoint_grid: least_common_multiple(mamba_cache_chunk_size, params.page_size),
+            mamba_max_states_per_path: params.mamba_max_states_per_path,
+        }
+    }
+}
+
+fn least_common_multiple(lhs: usize, rhs: usize) -> usize {
+    let mut a = lhs;
+    let mut b = rhs;
+    while b != 0 {
+        (a, b) = (b, a % b);
+    }
+    lhs / a * rhs
+}
+
+impl MambaComponent {
+    // Tier-selected mamba slot read for the lock paths; `host` picks the host slot.
+    fn has_value<K: ChildKeyType>(node: &Node<K>, host: bool) -> bool {
+        if host {
+            node.has_host_value(MAMBA)
+        } else {
+            node.has_device_value(MAMBA)
+        }
+    }
+
+    /// Defer the path-cap eviction so it runs after the insert's BackupKV.
+    fn emit_excess_path_states_eviction_(
+        &self,
+        tail_node_id: NodeId,
+        cache_actions: &mut Vec<CacheAction>,
+    ) {
+        if self.mamba_max_states_per_path.is_none() {
+            return;
+        }
+        cache_actions.push(CacheAction::MambaEvictExcessPathStates { tail_node_id });
+    }
+}
+
+impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
+    fn component_type(&self) -> ComponentType {
+        MAMBA
+    }
+
+    fn needs_incremental_backup(&self, tree_core: &UnifiedTreeCore<K>, node_id: NodeIdx_) -> bool {
+        let node = tree_core.arena.node(node_id);
+        node.has_device_value(MAMBA) && !node.has_host_value(MAMBA)
+    }
+
+    /// A match consumes only the best-match node's mamba state, so MATCH_END
+    /// touches just that node; new-leaf states enter the LRU at insert commit,
+    /// so WALKDOWN and INSERT_END are no-ops.
+    fn refresh_lru(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        phase: LRURefreshPhase,
+        node_id: NodeIdx_,
+    ) {
+        match phase {
+            LRURefreshPhase::Walkdown => {}
+            LRURefreshPhase::MatchEnd => {
+                if tree_core.arena.has_device_value(node_id, MAMBA) {
+                    tree_core.device_lru_list_mut(MAMBA).reset_node_mru(node_id);
+                }
+            }
+            LRURefreshPhase::InsertEnd => {}
+        }
+    }
+
+    /// A per-match predicate accepting nodes that hold mamba data.
+    fn create_match_validator(
+        &self,
+        _tree_core: &UnifiedTreeCore<K>,
+        match_device_only: bool,
+    ) -> Box<dyn FnMut(&UnifiedTreeCore<K>, NodeIdx_) -> bool> {
+        // HiCache: evicted + backuped (host_value present) is also a valid match.
+        Box::new(move |tree_core: &UnifiedTreeCore<K>, node_id: NodeIdx_| {
+            let node = tree_core.arena.node(node_id);
+            node.has_device_value(MAMBA) || (!match_device_only && node.has_host_value(MAMBA))
+        })
+    }
+
+    /// The mamba branching seqlen and the host-only hit bump.
+    fn finalize_match_result_in_tree_core(
+        &self,
+        tree_core: &UnifiedTreeCore<K>,
+        mut result: MatchResult,
+        _params: &MatchPrefixParams<'_, K>,
+        _value_chunks: &[Tensor],
+        _best_value_len: usize,
+    ) -> MatchResult {
+        let mamba_boundary_len = result.device_indices.size()[0] as usize + result.host_hit_length;
+
+        // Full KV may extend beyond the latest reusable Mamba state. The branching
+        // point is the last checkpoint-grid-aligned position within the Full-KV hit
+        // that lies beyond the current Mamba boundary.
+        let aligned_seqlen =
+            result.full_kv_hit_length / self.mamba_checkpoint_grid * self.mamba_checkpoint_grid;
+        result.mamba_branching_seqlen =
+            (aligned_seqlen > mamba_boundary_len).then_some(aligned_seqlen);
+
+        // HiCache: if mamba was evicted from device but has host backup,
+        // ensure mamba_host_hit_length >= 1 so load_back is triggered.
+        let last_node = tree_core
+            .arena
+            .node(tree_core.arena.resolve(result.best_match_node_id));
+        if !last_node.has_device_value(MAMBA) && last_node.has_host_value(MAMBA) {
+            result.mamba_host_hit_length = result.mamba_host_hit_length.max(1);
+        }
+        result
+    }
+
+    /// Attach the donated mamba slot to the insert target leaf.
+    fn commit_insert_component_data(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        is_new_leaf: bool,
+        params: &InsertParams<'_, K>,
+        result: &mut InsertResult,
+        cache_actions: &mut Vec<CacheAction>,
+    ) {
+        let mamba_value = params
+            .mamba_value
+            .as_ref()
+            .expect("mamba insert requires a donated mamba_value");
+        let slot_len = mamba_value.size()[0] as usize;
+
+        if is_new_leaf {
+            tree_core
+                .arena
+                .set_device_value(node_id, MAMBA, mamba_value.shallow_clone());
+            tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
+            tree_core.inc_evictable_size(MAMBA, slot_len);
+            self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions);
+            return;
+        }
+        if !tree_core.arena.has_device_value(node_id, MAMBA) {
+            // Tombstone refill: the node moves from the host LRU to the device LRU.
+            tree_core
+                .arena
+                .set_device_value(node_id, MAMBA, mamba_value.shallow_clone());
+            let host_lru = tree_core.host_lru_list_mut(MAMBA);
+            if host_lru.in_list(Some(node_id)) {
+                host_lru.remove_node(node_id);
+            }
+            tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
+            tree_core.inc_evictable_size(MAMBA, slot_len);
+            let tick = tree_core.arena.get_and_bump_access_counter();
+            tree_core.arena.node_mut(node_id).last_access_counter = tick;
+            self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions);
+            return;
+        }
+        tree_core.device_lru_list_mut(MAMBA).reset_node_mru(node_id);
+        let tick = tree_core.arena.get_and_bump_access_counter();
+        tree_core.arena.node_mut(node_id).last_access_counter = tick;
+        result.mamba_exist = true;
+    }
+
+    /// Mamba data stays on the original leaf; the new prefix node gets none.
+    /// Evict shallow Mamba device checkpoints beyond the per-path cap on the
+    /// tail's root path; Full KV, host backups, the tail, forks, locked nodes,
+    /// and device leaves are preserved (a best-effort soft cap).
+    fn evict_excess_path_states(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        tail_node_id: NodeIdx_,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+        let Some(cap) = self.mamba_max_states_per_path else {
+            return;
+        };
+        // Mamba-value holders on the root path, tail-first.
+        let mut holders: Vec<NodeIdx_> = Vec::new();
+        let mut cursor = Some(tail_node_id);
+        while let Some(node_id) = cursor {
+            let node = tree_core.arena.node(node_id);
+            if node.is_root() {
+                break;
+            }
+            if node.has_device_value(MAMBA) {
+                holders.push(node_id);
+            }
+            cursor = node.parent;
+        }
+        let mut excess = holders.len().saturating_sub(cap);
+        if excess == 0 {
+            return;
+        }
+        // Cache-level apply: the counts are not reported, only the frees.
+        let mut tracker: HashMap<ComponentType, usize> = HashMap::new();
+        for &node_id in holders.iter().rev() {
+            if excess == 0 || node_id == tail_node_id {
+                break;
+            }
+            let node = tree_core.arena.node(node_id);
+            if node.device_lock_ref(MAMBA) > 0 || node.children.len() != 1 {
+                continue;
+            }
+            if tree_core.evictable_device_leaves.contains(node_id) {
+                continue;
+            }
+            tree_core.evict_component_and_detach_lru_(
+                node_id,
+                MAMBA,
+                device_frees,
+                host_frees,
+                EvictLayer::Device,
+                Some(&mut tracker),
+            );
+            tree_core.cascade_evict_(
+                node_id,
+                MAMBA,
+                &mut tracker,
+                device_frees,
+                host_frees,
+                EvictLayer::Device,
+            );
+            excess -= 1;
+        }
+    }
+
+    fn redistribute_on_node_split(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        new_parent_id: NodeIdx_,
+        _child_id: NodeIdx_,
+    ) {
+        let new_parent = tree_core.arena.node_mut(new_parent_id);
+        if new_parent.has_device_value(MAMBA) {
+            let _ = new_parent.take_device_value(MAMBA);
+        }
+        new_parent.set_lock_ref_(ValueSlotIdx::device(MAMBA), 0);
+        // HiCache: mamba host_value stays on child (mamba = leaf-only data).
+        if new_parent.has_host_value(MAMBA) {
+            let _ = new_parent.take_host_value(MAMBA);
+        }
+        new_parent.set_lock_ref_(ValueSlotIdx::host(MAMBA), 0);
+    }
+
+    /// Free the node's mamba slot on the targeted layer(s).
+    fn evict_component(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        target: EvictLayer,
+    ) -> (usize, usize) {
+        let ct = MAMBA;
+        let node = tree_core.arena.node_mut(node_id);
+        let mut freed = 0;
+        let mut host_freed = 0;
+
+        // Device layer
+        if target.contains(EvictLayer::Device) && node.has_device_value(MAMBA) {
+            freed = node.device_value_len(MAMBA);
+            device_frees
+                .entry(ct)
+                .or_default()
+                .push(node.take_device_value(MAMBA));
+            tree_core.dec_evictable_size(MAMBA, freed);
+        }
+
+        // Host layer
+        let node = tree_core.arena.node_mut(node_id);
+        if target.contains(EvictLayer::Host) && node.has_host_value(MAMBA) {
+            host_freed = node.host_value_len(MAMBA);
+            host_frees
+                .entry(ct)
+                .or_default()
+                .push(node.take_host_value(MAMBA));
+            let host_lru = tree_core.host_lru_list_mut(MAMBA);
+            if host_lru.in_list(Some(node_id)) {
+                host_lru.remove_node(node_id);
+            }
+        }
+
+        // After device tombstone: if only host_value remains, insert into host LRU
+        let node = tree_core.arena.node(node_id);
+        if target == EvictLayer::Device
+            && !node.has_device_value(MAMBA)
+            && node.has_host_value(MAMBA)
+        {
+            let host_lru = tree_core.host_lru_list_mut(MAMBA);
+            if !host_lru.in_list(Some(node_id)) {
+                host_lru.insert_mru(node_id);
+            }
+        }
+
+        (freed, host_freed)
+    }
+
+    /// Begin the device-eviction walk from this component's LRU cursor.
+    fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<K>, request_cnt: usize) {
+        tree_core.set_evict_device_start(MAMBA, request_cnt);
+        let cursor = tree_core
+            .device_lru_list(MAMBA)
+            .get_lru_no_lock(&tree_core.arena);
+        tree_core.component_state_mut(MAMBA).evict_device_cursor = cursor;
+    }
+
+    /// Advance one device-eviction step and return a leaf, if selected.
+    ///
+    /// An internal tombstone is one complete step so the caller can apply its
+    /// pending frees and recheck allocator capacity before the next mutation.
+    fn evict_device_next_node(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        tracker: &mut HashMap<ComponentType, usize>,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) -> Option<NodeIdx_> {
+        let ct = MAMBA;
+        assert!(
+            tree_core.component_state(MAMBA).is_evict_device_ongoing,
+            "Mamba device eviction not started"
+        );
+        let mut cursor = tree_core.component_state(MAMBA).evict_device_cursor;
+        // The cursor is re-validated (reset to LRU head) if the previous
+        // node's eviction removed it.
+        if cursor.is_some_and(|c| !tree_core.device_lru_list(MAMBA).in_list(Some(c))) {
+            cursor = tree_core
+                .device_lru_list(MAMBA)
+                .get_lru_no_lock(&tree_core.arena);
+        }
+        let next = loop {
+            if tracker[&ct] >= tree_core.component_state(MAMBA).evict_device_request_cnt {
+                break None;
+            }
+            let Some(x) = cursor else {
+                break None;
+            };
+            if !tree_core.device_lru_list(MAMBA).in_list(Some(x)) {
+                break None;
+            }
+            assert!(
+                tree_core.arena.has_device_value(x, MAMBA),
+                "Mamba eviction cursor on a valueless node {x}"
+            );
+            cursor = tree_core
+                .device_lru_list(MAMBA)
+                .get_prev_no_lock(x, &tree_core.arena);
+            // A load-back pin means an in-flight DMA targets this node's slices.
+            if tree_core.arena.node(x).is_load_back_pending() {
+                continue;
+            }
+            if tree_core.evictable_device_leaves.contains(x) {
+                break Some(x);
+            }
+            // Internal nodes are tombstoned inline (no IO).
+            tree_core.evict_component_and_detach_lru_(
+                x,
+                ct,
+                device_frees,
+                host_frees,
+                EvictLayer::Device,
+                Some(tracker),
+            );
+            tree_core.cascade_evict_(x, ct, tracker, device_frees, host_frees, EvictLayer::Device);
+            break None;
+        };
+        tree_core.component_state_mut(MAMBA).evict_device_cursor = cursor;
+        next
+    }
+
+    /// Clear the device-eviction walk cursor state.
+    fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<K>) {
+        tree_core.set_evict_device_end(MAMBA);
+    }
+
+    /// Single-node mamba lock; host locks also detach from the host LRU.
+    fn acquire_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        mut result: IncLockRefResult,
+        lock_host: bool,
+    ) -> IncLockRefResult {
+        let node = tree_core.arena.node(node_id);
+        if node.is_root() {
+            return result;
+        }
+        // A node in skip_lock_node_ids was a tombstone when this lock was acquired.
+        if !Self::has_value(node, lock_host) {
+            result
+                .skip_lock_node_ids
+                .entry(MAMBA)
+                .or_default()
+                .insert(node.id);
+            return result;
+        }
+        if lock_host {
+            if node.host_lock_ref(MAMBA) == 0 {
+                let host_lru = tree_core.host_lru_list_mut(MAMBA);
+                if host_lru.in_list(Some(node_id)) {
+                    host_lru.remove_node(node_id);
+                }
+            }
+            tree_core.arena.inc_host_lock_ref(node_id, MAMBA);
+        } else {
+            let value_len = node.device_value_len(MAMBA);
+            if node.device_lock_ref(MAMBA) == 0 {
+                tree_core.dec_evictable_size(MAMBA, value_len);
+                tree_core.inc_protected_size(MAMBA, value_len);
+            }
+            tree_core.arena.inc_device_lock_ref(node_id, MAMBA);
+        }
+        result
+    }
+
+    /// Single-node mamba unlock; host unlocks reinsert into the host LRU.
+    fn release_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        params: Option<&DecLockRefParams>,
+        lock_host: bool,
+    ) {
+        if tree_core.arena.node(node_id).is_root() {
+            return;
+        }
+        if let Some(params) = params
+            && params
+                .skip_lock_node_ids
+                .get(&MAMBA)
+                .is_some_and(|ids| ids.contains(&tree_core.arena.node(node_id).id))
+        {
+            return;
+        }
+        if lock_host {
+            let node = tree_core.arena.node_mut(node_id);
+            node.dec_host_lock_ref(MAMBA);
+            if node.host_lock_ref(MAMBA) == 0
+                && !node.has_device_value(MAMBA)
+                && node.has_host_value(MAMBA)
+            {
+                let host_lru = tree_core.host_lru_list_mut(MAMBA);
+                if !host_lru.in_list(Some(node_id)) {
+                    host_lru.insert_mru(node_id);
+                }
+            }
+            return;
+        }
+        let node = tree_core.arena.node(node_id);
+        let device_lock_ref = node.device_lock_ref(MAMBA);
+        if device_lock_ref > 0 {
+            if device_lock_ref == 1 {
+                let value_len = node.device_value_len(MAMBA);
+                tree_core.inc_evictable_size(MAMBA, value_len);
+                tree_core.dec_protected_size(MAMBA, value_len);
+            }
+            tree_core.arena.dec_device_lock_ref(node_id, MAMBA);
+        }
+    }
+
+    /// Build the mamba transfer descriptors for the given phase.
+    fn build_hicache_transfers(
+        &self,
+        tree_core: &UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        phase: CacheTransferPhase,
+        mamba_pool_idx: Option<Tensor>,
+        host_indices: Option<Tensor>,
+        _token_ids: Option<&[i64]>,
+        _prefetch_tokens: usize,
+        _last_hash: Option<&str>,
+    ) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
+        Ok(match phase {
+            CacheTransferPhase::BackupHost => {
+                let node = tree_core.arena.node(node_id);
+                if node.has_host_value(MAMBA) {
+                    return Ok(None);
+                }
+                node.try_device_value(MAMBA).map(|value| {
+                    vec![PoolTransfer {
+                        name: PoolName::Mamba,
+                        device_indices: Some(value.shallow_clone()),
+                        ..Default::default()
+                    }]
+                })
+            }
+            CacheTransferPhase::LoadBack => {
+                let node = tree_core.arena.node(node_id);
+                if node.has_device_value(MAMBA) {
+                    return Ok(None);
+                }
+                let mut transfers = Vec::new();
+                // restore single node if host_value exists
+                if let Some(host_value) = node.try_host_value(MAMBA) {
+                    transfers.push(PoolTransfer {
+                        name: PoolName::Mamba,
+                        host_indices: Some(host_value.shallow_clone()),
+                        nodes_to_load: Some(vec![node.id]),
+                        ..Default::default()
+                    });
+                }
+                // Per-request mamba CoW (H->D copy into the request's device slot,
+                // pre-allocated on caller side).
+                if let (Some(mamba_pool_idx), Some(host_value)) =
+                    (mamba_pool_idx, node.try_host_value(MAMBA))
+                {
+                    transfers.push(PoolTransfer {
+                        name: PoolName::Mamba,
+                        host_indices: Some(host_value.shallow_clone()),
+                        device_indices: Some(mamba_pool_idx.unsqueeze(0)),
+                        ..Default::default()
+                    });
+                }
+                if transfers.is_empty() {
+                    None
+                } else {
+                    Some(transfers)
+                }
+            }
+            CacheTransferPhase::BackupStorage => {
+                let node = tree_core.arena.node(node_id);
+                let Some(host_value) = node.try_host_value(MAMBA) else {
+                    return Ok(None);
+                };
+                let Some(hash_value) = node.hash_value.as_ref().filter(|h| !h.is_empty()) else {
+                    return Ok(None);
+                };
+                Some(vec![PoolTransfer {
+                    name: PoolName::Mamba,
+                    host_indices: Some(host_value.shallow_clone()),
+                    keys: Some(vec![hash_value[hash_value.len() - 1].clone()]),
+                    hit_policy: PoolHitPolicy::TrailingPages,
+                    ..Default::default()
+                }])
+            }
+            CacheTransferPhase::Prefetch => {
+                let host_indices =
+                    host_indices.expect("Mamba PREFETCH build requires host indices");
+                Some(vec![PoolTransfer {
+                    name: PoolName::Mamba,
+                    host_indices: Some(host_indices),
+                    keys: Some(vec!["__placeholder__".to_string()]),
+                    hit_policy: PoolHitPolicy::TrailingPages,
+                    ..Default::default()
+                }])
+            }
+        })
+    }
+
+    /// Post-transfer mamba bookkeeping for the given phase.
+    fn commit_hicache_transfer(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        phase: CacheTransferPhase,
+        transfers: Vec<PoolTransfer>,
+        cache_actions: &mut Vec<CacheAction>,
+        insert_result: Option<&mut InsertResult>,
+        pool_storage_result: Option<&PoolTransferResult>,
+    ) {
+        match phase {
+            CacheTransferPhase::BackupHost => {
+                if let Some(transfer) = transfers.first()
+                    && let Some(host_indices) = &transfer.host_indices
+                {
+                    let node = tree_core.arena.node_mut(node_id);
+                    if !node.has_host_value(MAMBA) {
+                        node.set_host_value(MAMBA, host_indices.copy());
+                    }
+                }
+            }
+            CacheTransferPhase::LoadBack => {
+                let Some(transfer) = transfers.first() else {
+                    return;
+                };
+                if let Some(device_indices) = &transfer.device_indices {
+                    let node = tree_core.arena.node_mut(node_id);
+                    node.set_device_value(MAMBA, device_indices.copy());
+                    let count = node.device_value_len(MAMBA);
+                    // Move from host LRU to device LRU
+                    let host_lru = tree_core.host_lru_list_mut(MAMBA);
+                    if host_lru.in_list(Some(node_id)) {
+                        host_lru.remove_node(node_id);
+                    }
+                    tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
+                    tree_core.inc_evictable_size(MAMBA, count);
+                }
+            }
+            // The python elif chain has no BACKUP_STORAGE arm.
+            CacheTransferPhase::BackupStorage => {}
+            CacheTransferPhase::Prefetch => {
+                let Some(transfer) = transfers.first() else {
+                    return;
+                };
+                let host_indices = transfer.host_indices.as_ref();
+                let loaded = pool_storage_result.is_some_and(|result| {
+                    result
+                        .extra_pool_hit_pages
+                        .get(&PoolName::Mamba)
+                        .copied()
+                        .unwrap_or(0)
+                        >= 1
+                });
+                let target_node_id = insert_result
+                    .as_deref()
+                    .and_then(|result| result.inserted_host_node)
+                    .map(|id| tree_core.arena.resolve(id));
+                let attach_target = match (host_indices, target_node_id) {
+                    (Some(_), Some(target))
+                        if loaded && !tree_core.arena.has_host_value(target, MAMBA) =>
+                    {
+                        Some(target)
+                    }
+                    _ => None,
+                };
+                let Some(target) = attach_target else {
+                    // The buffer cannot attach: free it and let the caller keep
+                    // its own donated slot bookkeeping.
+                    cache_actions.push(CacheAction::FreeComponentHostSlot {
+                        component_type: MAMBA,
+                        host_indices: host_indices
+                            .map(|host| vec![host.shallow_clone()])
+                            .unwrap_or_default(),
+                    });
+                    if let Some(insert_result) = insert_result {
+                        insert_result.mamba_exist = true;
+                    }
+                    return;
+                };
+                let host_indices = host_indices.expect("an attach target implies host indices");
+                tree_core
+                    .arena
+                    .set_host_value(target, MAMBA, host_indices.copy());
+                if !tree_core.arena.has_device_value(target, MAMBA) {
+                    let host_lru = tree_core.host_lru_list_mut(MAMBA);
+                    if !host_lru.in_list(Some(target)) {
+                        host_lru.insert_mru(target);
+                    }
+                }
+                if let Some(insert_result) = insert_result {
+                    insert_result.mamba_exist = false;
+                }
+            }
+        }
+    }
+
+    /// Evict mamba host resources: internal nodes tombstone privately, host
+    /// leaves evict atomically.
+    fn drive_host_eviction(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        num_tokens: usize,
+        tracker: &mut HashMap<ComponentType, usize>,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+        let ct = MAMBA;
+        let mut x = tree_core
+            .host_lru_list(MAMBA)
+            .get_lru_no_lock(&tree_core.arena);
+        loop {
+            if tracker[&ct] >= num_tokens {
+                break;
+            }
+            let Some(cur) = x else {
+                break;
+            };
+            if !tree_core.host_lru_list(MAMBA).in_list(Some(cur)) {
+                break;
+            }
+            let x_next = tree_core
+                .host_lru_list(MAMBA)
+                .get_prev_no_lock(cur, &tree_core.arena);
+            // A load-back pin means an in-flight DMA reads this node's host slices.
+            if tree_core.arena.node(cur).is_load_back_pending() {
+                x = x_next;
+                continue;
+            }
+            if tree_core.evictable_host_leaves.contains(cur) {
+                // Host leaf: atomic eviction (all components host + delete)
+                tree_core.evict_host_leaf_(cur, tracker, device_frees, host_frees);
+            } else {
+                // Internal: tombstone Mamba + cascade
+                assert!(
+                    tree_core.arena.has_host_value(cur, MAMBA),
+                    "Mamba host LRU member {cur} has no host value"
+                );
+                tree_core.evict_component_and_detach_lru_(
+                    cur,
+                    ct,
+                    device_frees,
+                    host_frees,
+                    EvictLayer::Host,
+                    Some(tracker),
+                );
+                tree_core.cascade_evict_(
+                    cur,
+                    ct,
+                    tracker,
+                    device_frees,
+                    host_frees,
+                    EvictLayer::Host,
+                );
+            }
+            x = x_next;
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "../tests/components/mamba.rs"]
+mod tests;

rust/mem-cache/src/components/mod.rs

@@ -0,0 +1,491 @@
+//! Per-component drivers; each receives the whole `UnifiedTreeCore` for backward access.
+#![allow(unused_variables)]
+
+use std::collections::HashMap;
+
+use tch::Tensor;
+
+use crate::node::{ChildKeyType, NodeArena, NodeIdx_, TreeCoreRuntimeError};
+use crate::unified_tree_core::{
+    CacheAction, CacheTransferPhase, DecLockRefParams, EvictLayer, IncLockRefResult, InsertParams,
+    InsertResult, LRURefreshPhase, MatchPrefixParams, MatchResult, PoolTransfer,
+    PoolTransferResult, UnifiedTreeCore,
+};
+
+mod full;
+mod mamba;
+mod swa;
+
+pub use full::FullComponent;
+pub use mamba::MambaComponent;
+pub use swa::SwaComponent;
+
+/// Whether `node_id` holds the component's data on `target`, checking its
+/// device or host slot.
+pub(crate) fn node_has_component_data<K: ChildKeyType>(
+    arena: &NodeArena<K>,
+    node_id: NodeIdx_,
+    component_type: ComponentType,
+    target: EvictLayer,
+) -> bool {
+    match target {
+        EvictLayer::Device => arena.has_device_value(node_id, component_type),
+        EvictLayer::Host => arena.has_host_value(node_id, component_type),
+        EvictLayer::All => panic!("node_has_component_data: EvictLayer::All is not a single layer"),
+    }
+}
+
+/// Every device value of the component across all roots, concatenated.
+pub(crate) fn all_values_flatten<K: ChildKeyType>(
+    tree_core: &UnifiedTreeCore<K>,
+    component_type: ComponentType,
+) -> Tensor {
+    let mut values: Vec<Tensor> = Vec::new();
+    let mut stack: Vec<NodeIdx_> = vec![tree_core.arena.root()];
+    while let Some(node_id) = stack.pop() {
+        let node = tree_core.arena.node(node_id);
+        if let Some(value) = node.try_device_value(component_type) {
+            values.push(value.shallow_clone());
+        }
+        stack.extend(node.children.values().copied());
+    }
+    if values.is_empty() {
+        return tree_core.empty_device_indices.shallow_clone();
+    }
+    Tensor::cat(&values, 0)
+}
+
+/// A per-component lock/value/eviction driver over the shared `UnifiedTreeCore`.
+pub trait TreeComponent<K: ChildKeyType> {
+    /// The component this driver serves.
+    fn component_type(&self) -> ComponentType;
+
+    /// Whether this component has device data that still needs a host backup.
+    fn needs_incremental_backup(
+        &self,
+        _tree_core: &UnifiedTreeCore<K>,
+        _node_id: NodeIdx_,
+    ) -> bool {
+        false
+    }
+
+    /// Refresh this component's LRU position for `node_id` at the given walk phase.
+    fn refresh_lru(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        phase: LRURefreshPhase,
+        node_id: NodeIdx_,
+    ) {
+        // Python reference — tree_component.py::TreeComponent.refresh_lru:
+        //     def refresh_lru(
+        //         self,
+        //         phase: LRURefreshPhase,
+        //         node: UnifiedTreeNode,
+        //         root_node: UnifiedTreeNode,
+        //     ) -> None:
+        //         ct = self.component_type
+        //         match phase:
+        //             case LRURefreshPhase.WALKDOWN:
+        //                 if node.component_data[ct].value is None:
+        //                     return
+        //                 self.tree_core.lru_lists[ct].reset_node_mru(node)
+        //             case LRURefreshPhase.MATCH_END:
+        //                 self.tree_core.lru_lists[ct].reset_node_and_parents_mru(
+        //                     node, root_node, self.node_has_component_data
+        //                 )
+        //             case LRURefreshPhase.INSERT_END:
+        //                 # WALKDOWN already refreshed every node on the insert path
+        //                 # (including the new leaf), so there is nothing more to do.
+        //                 return
+        //             case _:
+        //                 raise ValueError(f"Unknown LRURefreshPhase: {phase}")
+        unimplemented!("TreeComponent.refresh_lru")
+    }
+
+    /// Return a per-match stateful predicate deciding whether a node is a valid
+    /// match boundary for this component.
+    // Python reference — tree_component.py::TreeComponent.create_match_validator:
+    //     @abstractmethod
+    //     def create_match_validator(
+    //         self, match_device_only: bool = False
+    //     ) -> Callable[[UnifiedTreeNode], bool]:
+    //         """Return a per-match stateful predicate that decides whether a node
+    //         is a valid match boundary for this component.
+    //         Called once per match_prefix; the returned closure may carry state.
+    //         When match_device_only is true, host-backed nodes must not be accepted
+    //         as valid match boundaries.
+    //         - Full: returns True if the node has full component data.
+    //         - SWA: tracks accumulated length since last gap; returns True only
+    //           when the contiguous window reaches swa_sliding_window_size.
+    //         - Mamba: returns True iff the node has mamba component data."""
+    //         ...
+    fn create_match_validator(
+        &self,
+        tree_core: &UnifiedTreeCore<K>,
+        match_device_only: bool,
+    ) -> Box<dyn FnMut(&UnifiedTreeCore<K>, NodeIdx_) -> bool>;
+
+    /// Tree-side post-processing inside the match walk (no cache access).
+    fn finalize_match_result_in_tree_core(
+        &self,
+        tree_core: &UnifiedTreeCore<K>,
+        result: MatchResult,
+        params: &MatchPrefixParams<'_, K>,
+        value_chunks: &[Tensor],
+        best_value_len: usize,
+    ) -> MatchResult {
+        result
+    }
+
+    /// Called per-node when an insert's key overlaps an existing node.
+    /// Returns the index within `value_slice` from which this component
+    /// consumed (took ownership of) the underlying KV pool slots.
+    /// Returns `prefix_len` if nothing was consumed (default).
+    /// The insert walk uses this to free only the non-consumed duplicate
+    /// portion: `value_slice[dup_start..consumed_from]`.
+    fn update_component_on_insert_overlap(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        prefix_len: usize,
+        total_prefix_len: usize,
+        value_slice: Tensor,
+        params: &InsertParams<'_, K>,
+        result: &mut InsertResult,
+        cache_actions: &mut Vec<CacheAction>,
+    ) -> usize {
+        prefix_len
+    }
+
+    /// Called after `unevict_node_on_insert_` restores the base (Full) value
+    /// on an evicted node. Aux components (e.g. SWA) override this to rebuild
+    /// their own data from the freshly assigned base value when their entry
+    /// is still tombstoned. Default no-op.
+    fn recover_after_unevict(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        prefix_len: usize,
+        total_prefix_len: usize,
+        params: &InsertParams<'_, K>,
+        result: &mut InsertResult,
+        cache_actions: &mut Vec<CacheAction>,
+    ) {
+    }
+
+    /// Finalize component data on the target (leaf) node after the insert
+    /// walk completes. Called once per insert.
+    /// - Full: no-op (full data is handled by `add_new_node_`).
+    /// - SWA: for new leaves, checks whether the node straddles the SWA
+    ///   eviction boundary (`swa_evicted_seqlen`). If so, splits the node
+    ///   via `split_node_` — the parent becomes a tombstone (no SWA) and the
+    ///   child (the deeper portion) receives SWA data. If the entire node
+    ///   is within the window, sets SWA directly. If entirely outside,
+    ///   leaves SWA as None (tombstone).
+    /// - Mamba: sets the mamba component value from params, inserts into the
+    ///   mamba LRU list, and increments evictable size. If the node already
+    ///   has mamba data, resets its LRU position instead.
+    fn commit_insert_component_data(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        is_new_leaf: bool,
+        params: &InsertParams<'_, K>,
+        result: &mut InsertResult,
+        cache_actions: &mut Vec<CacheAction>,
+    ) {
+    }
+
+    /// Evict shallow device checkpoints beyond the per-path state cap on the
+    /// tail's root path; only the Mamba component caps its states.
+    fn evict_excess_path_states(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        tail_node_id: NodeIdx_,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+    }
+
+    /// Redistribute component data between `new_parent` and `child` when a node is
+    /// split; `new_parent` is the newly created prefix node.
+    // Python reference — tree_component.py::TreeComponent.redistribute_on_node_split:
+    //     @abstractmethod
+    //     def redistribute_on_node_split(
+    //         self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
+    //     ):
+    //         """Redistribute component data between new_parent and child when a
+    //         node is split. new_parent is the newly created prefix node.
+    //         - Full: copies child's lock_ref to new_parent.
+    //         - SWA: slices (or clones) the swa value for new_parent, copies
+    //           lock_ref and component_uuid metadata, then syncs child's swa
+    //           value with its (now-trimmed) full_value.
+    //         - Mamba: sets new_parent's mamba value to None and lock_ref to 0
+    //           (mamba data stays on the original leaf, not on prefix nodes)."""
+    //         ...
+    fn redistribute_on_node_split(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        new_parent_id: NodeIdx_,
+        child_id: NodeIdx_,
+    );
+
+    /// Free this component's KV resources on a node being evicted; returns
+    /// (device_freed, host_freed) token counts.
+    // Python reference — tree_component.py::TreeComponent.evict_component:
+    //     @abstractmethod
+    //     def evict_component(
+    //         self,
+    //         node: UnifiedTreeNode,
+    //         device_frees: dict[ComponentType, list[torch.Tensor]],
+    //         host_frees: dict[ComponentType, list[torch.Tensor]],
+    //         target: EvictLayer = EvictLayer.DEVICE,
+    //     ) -> tuple[int, int]:
+    //         """Free this component's KV resources on a node being evicted.
+    //
+    //         *target* controls which layer(s) to evict:
+    //           - DEVICE: free device memory and tombstone (value = None).
+    //                     Host data is untouched.
+    //           - HOST:   free host memory (host_value = None).
+    //                     Device data is untouched.
+    //           - ALL:    free both device and host memory.
+    //                     No tombstone — caller will delete the node.
+    //
+    //         Returns (device_freed, host_freed) token counts."""
+    //         ...
+    fn evict_component(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        target: EvictLayer,
+    ) -> (usize, usize);
+
+    /// Eviction priority on this node type; higher = evicted later, and evicting a
+    /// component cascade-evicts every component of equal or lower priority.
+    fn eviction_priority(&self, is_leaf: bool) -> i64 {
+        0
+    }
+
+    /// Begin this component's device-eviction walk (build its cursor/heap).
+    fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<K>, request_cnt: usize);
+
+    /// Advance one eviction step and return a device leaf, if selected.
+    ///
+    /// Implementations must return after one allocator-relevant internal
+    /// mutation so the caller can drain pending frees before continuing.
+    fn evict_device_next_node(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        tracker: &mut HashMap<ComponentType, usize>,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) -> Option<NodeIdx_>;
+
+    /// Clear this component's device-eviction walk state.
+    fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<K>);
+
+    /// Increment component lock refs, protecting nodes from eviction.
+    // Python reference — tree_component.py::TreeComponent.acquire_component_lock:
+    //     @abstractmethod
+    //     def acquire_component_lock(
+    //         self,
+    //         node: UnifiedTreeNode,
+    //         result: IncLockRefResult,
+    //         lock_host: bool = False,
+    //     ) -> IncLockRefResult:
+    //         """Increment component lock refs, protecting nodes from
+    //         eviction. Updates evictable → protected size on first lock.
+    //         - Full: path-lock — walks from node up to root, incrementing
+    //           lock_ref on every ancestor.
+    //         - SWA: path-lock — walks upward collecting swa values until the
+    //           sliding window is filled; records a component_uuid at the
+    //           boundary for release_component_lock to know where to stop.
+    //         - Mamba: single-node lock — only increments lock_ref on the
+    //           node itself (mamba state is per-leaf, not per-path).
+    //
+    //         When ``lock_host`` is True, the lock applies to host-side state:
+    //         - Full: single-node host lock.
+    //         - SWA: host window-lock with a dedicated host UUID boundary.
+    //         - Mamba: single-node host lock with host LRU detach."""
+    //         ...
+    fn acquire_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        result: IncLockRefResult,
+        lock_host: bool,
+    ) -> IncLockRefResult;
+
+    /// Decrement component lock refs, un-protecting nodes.
+    // Python reference — tree_component.py::TreeComponent.release_component_lock:
+    //     @abstractmethod
+    //     def release_component_lock(
+    //         self,
+    //         node: UnifiedTreeNode,
+    //         params: Optional[DecLockRefParams],
+    //         lock_host: bool = False,
+    //     ) -> None:
+    //         """Decrement component lock refs, un-protecting nodes.
+    //         Updates protected → evictable size when lock_ref drops to 0.
+    //         - Full: path-unlock — walks from node up to root, decrementing
+    //           lock_ref on every ancestor.
+    //         - SWA: path-unlock — walks upward, stopping at the node whose
+    //           component_uuid matches the one recorded during acquire.
+    //         - Mamba: single-node unlock — only decrements lock_ref on the
+    //           node itself.
+    //
+    //         When ``lock_host`` is True, the inverse host-side semantics apply."""
+    //         ...
+    fn release_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        params: Option<&DecLockRefParams>,
+        lock_host: bool,
+    );
+
+    /// Early-release the SWA lock along [node, swa_uuid_for_lock] while leaving
+    /// the other components' locks intact; only the SWA component supports it.
+    fn release_window_lock(
+        &self,
+        _tree_core: &mut UnifiedTreeCore<K>,
+        _node_id: NodeIdx_,
+        _swa_uuid_for_lock: Option<i64>,
+        _device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        _host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+        unimplemented!("release_window_lock is SWA-only")
+    }
+
+    /// Build transfer descriptors for this component in the given phase; None when
+    /// the component has nothing to transfer.
+    fn build_hicache_transfers(
+        &self,
+        tree_core: &UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        phase: CacheTransferPhase,
+        mamba_pool_idx: Option<Tensor>,
+        host_indices: Option<Tensor>,
+        token_ids: Option<&[i64]>,
+        prefetch_tokens: usize,
+        last_hash: Option<&str>,
+    ) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
+        // Python reference — tree_component.py::TreeComponent.build_hicache_transfers:
+        //     def build_hicache_transfers(
+        //         self,
+        //         node: UnifiedTreeNode,
+        //         phase: CacheTransferPhase,
+        //         *,
+        //         mamba_pool_idx: Optional[torch.Tensor] = None,
+        //         host_indices: Optional[torch.Tensor] = None,
+        //         token_ids: Optional[Sequence[int]] = None,
+        //         prefetch_tokens: int = 0,
+        //         last_hash: Optional[str] = None,
+        //     ) -> Optional[list[PoolTransfer]]:
+        //         """Build transfer descriptors for this component in the given phase.
+        //         Returns None if the component has nothing to transfer."""
+        //         return None
+        unimplemented!("TreeComponent.build_hicache_transfers")
+    }
+
+    /// Post-transfer bookkeeping: store host indices, update LRU, etc.
+    fn commit_hicache_transfer(
+        &self,
+        tree_core: &mut UnifiedTreeCore<K>,
+        node_id: NodeIdx_,
+        phase: CacheTransferPhase,
+        transfers: Vec<PoolTransfer>,
+        cache_actions: &mut Vec<CacheAction>,
+        insert_result: Option<&mut InsertResult>,
+        pool_storage_result: Option<&PoolTransferResult>,
+    ) {
+        // Python reference — tree_component.py::TreeComponent.commit_hicache_transfer:
+        //     def commit_hicache_transfer(
+        //         self,
+        //         node: UnifiedTreeNode,
+        //         phase: CacheTransferPhase,
+        //         transfers: list[PoolTransfer] = (),
+        //         *,
+        //         cache_actions: list[CacheAction | ComponentAction],
+        //         insert_result: Optional[InsertResult] = None,
+        //         pool_storage_result: Optional[PoolTransferResult] = None,
+        //     ) -> None:
+        //         """Post-transfer bookkeeping: store host indices, update LRU, etc."""
+        //         pass
+        unimplemented!("TreeComponent.commit_hicache_transfer")
+    }
+
+    /// Reclaim host values that coexist with device values before ordinary
+    /// host eviction. Called only under the write-back policy.
+    fn reclaim_coexisting_host_values(
+        &self,
+        _tree_core: &mut UnifiedTreeCore<K>,
+        _num_tokens: usize,
+        _tracker: &mut HashMap<ComponentType, usize>,
+        _device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        _host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+    }
+
+    /// Evict from this component's host-side resources.
+    /// Called by HostPoolGroup when the host pool is full.
+    /// Default no-op for components without host storage.
+    fn drive_host_eviction(
+        &self,
+        _tree_core: &mut UnifiedTreeCore<K>,
+        _num_tokens: usize,
+        _tracker: &mut HashMap<ComponentType, usize>,
+        _device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        _host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) {
+    }
+}
+
+// Tree component types.
+
+/// The tree components; discriminants define the per-component array indexes.
+#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]
+pub enum ComponentType {
+    Full = 0,
+    Swa = 1,
+    Mamba = 2,
+}
+
+/// Short call-site aliases for the component types.
+pub const FULL: ComponentType = ComponentType::Full;
+pub const SWA: ComponentType = ComponentType::Swa;
+pub const MAMBA: ComponentType = ComponentType::Mamba;
+
+/// The base component every tree runs; the others are auxiliary.
+pub const BASE_COMPONENT_TYPE: ComponentType = ComponentType::Full;
+
+/// Slots per tier — the arrays are sized to this, not the enabled subset.
+pub const NUM_COMPONENT_TYPES: usize = ComponentType::Mamba as usize + 1;
+
+impl ComponentType {
+    /// Index into a per-component array.
+    pub const fn idx(self) -> usize {
+        self as usize
+    }
+
+    /// Whether the component stores one state slot per node (Mamba) instead of
+    /// one row per key atom.
+    pub fn single_value_per_node(self) -> bool {
+        matches!(self, ComponentType::Mamba)
+    }
+
+    /// The component at a per-component array index; panics out of range.
+    pub fn from_idx(idx: usize) -> ComponentType {
+        match idx {
+            0 => ComponentType::Full,
+            1 => ComponentType::Swa,
+            2 => ComponentType::Mamba,
+            _ => panic!("from_idx: {idx} is not a component index"),
+        }
+    }
+}
+#[cfg(test)]
+#[path = "../tests/components/base.rs"]
+mod tests;

rust/mem-cache/src/lib.rs

@@ -0,0 +1,27 @@
+//! Rust radix tree core for SGLang's KV cache, exposed to Python as `mem_cache`.
+// TODO(Jialin): Replace recoverable panics with explicit Rust errors and map
+// them to PyErr at the Python boundary.
+#![allow(
+    dead_code,
+    unsafe_op_in_unsafe_fn,
+    clippy::unwrap_used,
+    clippy::expect_used,
+    clippy::owned_cow,
+    clippy::panic,
+    clippy::print_stdout,
+    clippy::too_many_arguments,
+    clippy::type_complexity,
+    clippy::unimplemented,
+    clippy::unreachable,
+    clippy::useless_conversion
+)]
+
+mod components;
+mod node;
+#[cfg(feature = "python-extension")]
+mod python_bindings;
+#[cfg(test)]
+#[path = "tests/test_utils.rs"]
+pub(crate) mod test_utils;
+mod unified_lru_list;
+mod unified_tree_core;

rust/mem-cache/src/tests/components/base.rs

@@ -0,0 +1,160 @@
+use super::*;
+use crate::components::FULL;
+use crate::unified_tree_core::CacheInitParams;
+
+// Test-only component exercising the trait defaults; abstract hooks stay unimplemented.
+struct DefaultComponentForTest;
+
+impl TreeComponent<Vec<i64>> for DefaultComponentForTest {
+    fn component_type(&self) -> ComponentType {
+        FULL
+    }
+
+    fn create_match_validator(
+        &self,
+        _tree_core: &UnifiedTreeCore<Vec<i64>>,
+        match_device_only: bool,
+    ) -> Box<dyn FnMut(&UnifiedTreeCore<Vec<i64>>, NodeIdx_) -> bool> {
+        unimplemented!()
+    }
+
+    fn redistribute_on_node_split(
+        &self,
+        tree_core: &mut UnifiedTreeCore<Vec<i64>>,
+        new_parent_id: NodeIdx_,
+        child_id: NodeIdx_,
+    ) {
+        unimplemented!()
+    }
+
+    fn evict_component(
+        &self,
+        tree_core: &mut UnifiedTreeCore<Vec<i64>>,
+        node_id: NodeIdx_,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        target: EvictLayer,
+    ) -> (usize, usize) {
+        unimplemented!()
+    }
+
+    fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<Vec<i64>>, request_cnt: usize) {
+        unimplemented!()
+    }
+
+    fn evict_device_next_node(
+        &self,
+        tree_core: &mut UnifiedTreeCore<Vec<i64>>,
+        tracker: &mut HashMap<ComponentType, usize>,
+        device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+        host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    ) -> Option<NodeIdx_> {
+        unimplemented!()
+    }
+
+    fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<Vec<i64>>) {
+        unimplemented!()
+    }
+
+    fn acquire_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<Vec<i64>>,
+        node_id: NodeIdx_,
+        result: IncLockRefResult,
+        lock_host: bool,
+    ) -> IncLockRefResult {
+        unimplemented!()
+    }
+
+    fn release_component_lock(
+        &self,
+        tree_core: &mut UnifiedTreeCore<Vec<i64>>,
+        node_id: NodeIdx_,
+        params: Option<&DecLockRefParams>,
+        lock_host: bool,
+    ) {
+        unimplemented!()
+    }
+}
+
+#[test]
+fn insert_overlap_default_consumes_nothing() {
+    let mut tc: UnifiedTreeCore<Vec<i64>> =
+        UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]);
+    let root = tc.arena.root();
+    let consumed_from = DefaultComponentForTest.update_component_on_insert_overlap(
+        &mut tc,
+        root,
+        /* prefix_len = */ 3,
+        /* total_prefix_len = */ 0,
+        Tensor::from_slice(&[0i64, 1, 2]),
+        &InsertParams {
+            key: &vec![0, 1, 2],
+            namespace: Default::default(),
+            value: Tensor::from_slice(&[0i64, 1, 2]),
+            mamba_value: None,
+            prev_prefix_len: 0,
+            swa_evicted_seqlen: 0,
+            chunked: false,
+            priority: 0,
+            track_adopted_ranges: false,
+        },
+        &mut InsertResult::default(),
+        &mut Vec::new(),
+    );
+    // Nothing consumed: the whole overlap stays freeable as duplicates.
+    assert_eq!(consumed_from, 3);
+}
+
+#[test]
+fn finalize_match_result_default_returns_result_unchanged() {
+    let tc: UnifiedTreeCore<Vec<i64>> =
+        UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]);
+    let result = MatchResult {
+        last_device_node_id: 3,
+        best_match_node_id: 7,
+        host_hit_length: 11,
+        ..tc.empty_match_result()
+    };
+    let out = DefaultComponentForTest.finalize_match_result_in_tree_core(
+        &tc,
+        result,
+        &MatchPrefixParams {
+            key: &Vec::new(),
+            namespace: Default::default(),
+        },
+        &[],
+        0,
+    );
+    assert_eq!(out.last_device_node_id, 3);
+    assert_eq!(out.best_match_node_id, 7);
+    assert_eq!(out.host_hit_length, 11);
+}
+
+#[test]
+fn drive_host_eviction_default_is_a_noop() {
+    let mut tc: UnifiedTreeCore<Vec<i64>> =
+        UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]);
+    let mut tracker = HashMap::from([(FULL, 5usize)]);
+    let mut device_frees = HashMap::new();
+    let mut host_frees = HashMap::new();
+    DefaultComponentForTest.drive_host_eviction(
+        &mut tc,
+        /* num_tokens = */ 100,
+        &mut tracker,
+        &mut device_frees,
+        &mut host_frees,
+    );
+    assert_eq!(tracker[&FULL], 5);
+    assert!(device_frees.is_empty());
+    assert!(host_frees.is_empty());
+}
+
+// Component types.
+
+#[test]
+fn idx_matches_discriminants() {
+    assert_eq!(ComponentType::Full.idx(), 0);
+    assert_eq!(ComponentType::Swa.idx(), 1);
+    assert_eq!(ComponentType::Mamba.idx(), 2);
+}

rust/mem-cache/src/tests/test_utils.rs

@@ -0,0 +1,46 @@
+//! Shared helpers for the crate's unit tests.
+
+use std::collections::HashMap;
+
+use tch::Tensor;
+
+use crate::components::ComponentType;
+use crate::unified_tree_core::{CacheAction, EvictionStepResult};
+
+/// Fold an eviction step into a caller's running accumulators (the Controller
+/// consumption contract: deltas add, freed tensors append).
+pub(crate) fn accumulate_step(
+    step: EvictionStepResult,
+    tracker: &mut HashMap<ComponentType, usize>,
+    device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+    host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
+) {
+    for (ct, delta) in step.tracker {
+        *tracker.entry(ct).or_insert(0) += delta;
+    }
+    for (ct, tensors) in step.device_frees {
+        device_frees.entry(ct).or_default().extend(tensors);
+    }
+    for (ct, tensors) in step.host_frees {
+        host_frees.entry(ct).or_default().extend(tensors);
+    }
+}
+
+/// Short variant names for diagnosing an action sequence's shape.
+pub(crate) fn action_kinds(actions: &[CacheAction]) -> Vec<&'static str> {
+    actions
+        .iter()
+        .map(|action| match action {
+            CacheAction::FreeDeviceKV(_) => "FreeDeviceKV",
+            CacheAction::FreeDeviceKVFullOnly(_) => "FreeDeviceKVFullOnly",
+            CacheAction::BackupKV(_) => "BackupKV",
+            CacheAction::ReplaceWriteThroughOnNodeSplit { .. } => "ReplaceWriteThroughOnNodeSplit",
+            CacheAction::MambaEvictExcessPathStates { .. } => "MambaEvictExcessPathStates",
+            CacheAction::FreeComponentDeviceSlot { .. } => "FreeComponentDeviceSlot",
+            CacheAction::FreeComponentHostSlot { .. } => "FreeComponentHostSlot",
+            CacheAction::RebuildFullToSwaMapping { .. } => "RebuildFullToSwaMapping",
+            CacheAction::RecoverSwaWithLockedFull { .. } => "RecoverSwaWithLockedFull",
+            CacheAction::SwaRebuild { .. } => "SwaRebuild",
+        })
+        .collect()
+}

rust/mem-cache/src/tests/unified_lru_list.rs

@@ -0,0 +1,708 @@
+use super::*;
+use crate::components::FULL;
+use crate::node::{NodeArena, NodeIdx_, ValueSlotIdx};
+
+fn order(list: &UnifiedLRUList) -> Vec<NodeIdx_> {
+    list.iter().collect()
+}
+
+#[test]
+fn fresh_list_reads_are_empty() {
+    let list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    assert_eq!(list.get_lru_where(|_| true), None);
+    assert_eq!(list.iter().count(), 0);
+    assert_eq!(list.len(), 0);
+    list.validate();
+}
+
+#[test]
+fn insert_mru_orders_most_recent_first() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.insert_mru(NodeIdx_(30));
+    assert_eq!(order(&list), vec![NodeIdx_(30), NodeIdx_(20), NodeIdx_(10)]);
+    assert_eq!(list.len(), 3);
+    assert!(list.in_list(Some(NodeIdx_(10))));
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "already in the LRU list")]
+fn insert_mru_panics_when_already_a_member() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(1));
+    list.insert_mru(NodeIdx_(1));
+}
+
+#[test]
+fn remove_node_updates_membership_immediately() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.insert_mru(NodeIdx_(30));
+    list.remove_node(NodeIdx_(20));
+    assert_eq!(order(&list), vec![NodeIdx_(30), NodeIdx_(10)]);
+    assert!(!list.in_list(Some(NodeIdx_(20))));
+    assert_eq!(list.len(), 2);
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "index out of bounds")]
+fn remove_node_panics_when_absent() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.remove_node(NodeIdx_(1));
+}
+
+#[test]
+#[should_panic(expected = "not in the LRU list")]
+fn remove_node_panics_on_a_node_removed_earlier() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.remove_node(NodeIdx_(10));
+    // The cell is still allocated but reset; membership must gate the removal.
+    list.remove_node(NodeIdx_(10));
+}
+
+#[test]
+#[should_panic(expected = "not in the LRU list")]
+fn remove_node_panics_on_an_unlisted_cell() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.remove_node(NodeIdx_(10));
+    list.remove_node_(UnifiedLRUList::cell_of_(NodeIdx_(10)));
+}
+
+#[test]
+#[should_panic(expected = "already in the LRU list")]
+fn add_node_panics_on_a_linked_cell() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.add_node_(UnifiedLRUList::cell_of_(NodeIdx_(10)));
+}
+
+#[test]
+fn removed_nodes_can_be_reinserted() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.remove_node(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(10));
+    assert_eq!(order(&list), vec![NodeIdx_(10), NodeIdx_(20)]);
+    list.validate();
+}
+
+#[test]
+fn reset_node_mru_moves_a_member_to_the_front() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.insert_mru(NodeIdx_(30));
+    list.reset_node_mru(NodeIdx_(10));
+    assert_eq!(order(&list), vec![NodeIdx_(10), NodeIdx_(30), NodeIdx_(20)]);
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "index out of bounds")]
+fn reset_node_mru_panics_on_a_non_member() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.reset_node_mru(NodeIdx_(7));
+}
+
+// Arena chain root -> a -> b -> c plus a sibling, two atoms per key.
+fn arena_chain() -> (
+    NodeArena<Vec<i64>>,
+    NodeIdx_,
+    NodeIdx_,
+    NodeIdx_,
+    NodeIdx_,
+    NodeIdx_,
+) {
+    let mut arena = NodeArena::new(vec![crate::components::FULL], /* page_size = */ 1);
+    let root = arena.root();
+    let a = arena
+        .alloc_child(
+            root,
+            /* key = */ vec![1, 11],
+            /* priority = */ 0,
+            /* extra_key = */ None,
+        )
+        .unwrap();
+    let b = arena
+        .alloc_child(
+            a,
+            /* key = */ vec![2, 22],
+            /* priority = */ 0,
+            /* extra_key = */ None,
+        )
+        .unwrap();
+    let c = arena
+        .alloc_child(
+            b,
+            /* key = */ vec![3, 33],
+            /* priority = */ 0,
+            /* extra_key = */ None,
+        )
+        .unwrap();
+    let other = arena
+        .alloc_child(
+            root,
+            /* key = */ vec![9, 99],
+            /* priority = */ 0,
+            /* extra_key = */ None,
+        )
+        .unwrap();
+    (arena, root, a, b, c, other)
+}
+
+#[test]
+fn reset_parents_mru_reranks_included_nodes_deepest_first() {
+    let (arena, _root, a, b, c, other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.insert_mru(c);
+    list.insert_mru(other);
+    // b is excluded and skipped; c then a become the MRU run.
+    list.reset_node_and_parents_mru(c, &arena, |node| node.idx != b);
+    assert_eq!(order(&list), vec![c, a, other]);
+    list.validate();
+}
+
+#[test]
+fn reset_parents_mru_reranks_ancestors_when_the_deepest_is_excluded() {
+    let (arena, _root, a, _b, c, other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.insert_mru(other);
+    // c and b are excluded; a alone becomes the new MRU head.
+    list.reset_node_and_parents_mru(c, &arena, |node| node.idx == a);
+    assert_eq!(order(&list), vec![a, other]);
+    list.validate();
+}
+
+#[test]
+fn reset_walks_are_noops_when_node_is_the_root() {
+    let (arena, root, a, _b, _c, _other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.reset_node_and_parents_mru(root, &arena, |_| true);
+    list.reset_node_and_window_ancestors_mru(root, 4, &arena, |_| true);
+    assert_eq!(order(&list), vec![a]);
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "not in the LRU list")]
+fn reset_parents_mru_panics_on_an_unlisted_included_node() {
+    let (arena, _root, _a, _b, c, _other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(c);
+    list.remove_node(c);
+    list.reset_node_and_parents_mru(c, &arena, |_| true);
+}
+
+#[test]
+fn reset_window_ancestors_mru_stops_at_the_window() {
+    let (arena, _root, a, b, c, other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(b);
+    list.insert_mru(c);
+    list.insert_mru(a);
+    list.insert_mru(other);
+    // A window of 4 atoms covers c and b; a stays put beyond it.
+    list.reset_node_and_window_ancestors_mru(c, 4, &arena, |_| true);
+    assert_eq!(order(&list), vec![c, b, other, a]);
+    list.validate();
+}
+
+#[test]
+fn reset_window_ancestors_mru_includes_the_straddling_ancestor() {
+    let (arena, _root, a, b, c, other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.insert_mru(b);
+    list.insert_mru(c);
+    list.insert_mru(other);
+    // A window of 5 atoms ends mid-a: the straddling ancestor is still included.
+    list.reset_node_and_window_ancestors_mru(c, 5, &arena, |_| true);
+    assert_eq!(order(&list), vec![c, b, a, other]);
+    list.validate();
+}
+
+#[test]
+fn reset_walks_stop_at_the_salted_chains_root() {
+    let mut arena: NodeArena<Vec<i64>> =
+        NodeArena::new(vec![crate::components::FULL], /* page_size = */ 1);
+    let named = arena.root();
+    let a = arena
+        .alloc_child(
+            named,
+            /* key = */ vec![1, 11],
+            /* priority = */ 0,
+            Some("lora-1"),
+        )
+        .unwrap();
+    let b = arena
+        .alloc_child(
+            a,
+            /* key = */ vec![2, 22],
+            /* priority = */ 0,
+            /* extra_key = */ None,
+        )
+        .unwrap();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.insert_mru(b);
+    // Both walks terminate at the root without visiting it.
+    list.reset_node_and_parents_mru(b, &arena, |_| true);
+    assert_eq!(order(&list), vec![b, a]);
+    list.reset_node_and_window_ancestors_mru(b, 100, &arena, |_| true);
+    assert_eq!(order(&list), vec![b, a]);
+    list.validate();
+}
+
+#[test]
+fn get_lru_no_lock_returns_the_lru_most_unlocked_member() {
+    let (mut arena, _root, a, b, c, _other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.insert_mru(b);
+    list.insert_mru(c);
+    assert_eq!(list.get_lru_no_lock(&arena), Some(a));
+    // A lock on the list's own slot hides the LRU end from the walker.
+    arena
+        .node_mut(a)
+        .set_lock_ref_(ValueSlotIdx::device(FULL), 1);
+    assert_eq!(list.get_lru_no_lock(&arena), Some(b));
+    arena
+        .node_mut(b)
+        .set_lock_ref_(ValueSlotIdx::device(FULL), 1);
+    arena
+        .node_mut(c)
+        .set_lock_ref_(ValueSlotIdx::device(FULL), 1);
+    assert_eq!(list.get_lru_no_lock(&arena), None);
+}
+
+#[test]
+fn get_prev_no_lock_skips_locked_members_toward_the_mru_end() {
+    let (mut arena, _root, a, b, c, other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.insert_mru(b);
+    list.insert_mru(c);
+    list.insert_mru(other);
+    assert_eq!(list.get_prev_no_lock(a, &arena), Some(b));
+    // The locked b is skipped; from the MRU end there is no predecessor left.
+    arena
+        .node_mut(b)
+        .set_lock_ref_(ValueSlotIdx::device(FULL), 1);
+    assert_eq!(list.get_prev_no_lock(a, &arena), Some(c));
+    assert_eq!(list.get_prev_no_lock(other, &arena), None);
+    // A lock on a different slot does not gate this list's walker.
+    arena.node_mut(c).set_lock_ref_(ValueSlotIdx::host(FULL), 1);
+    assert_eq!(list.get_prev_no_lock(a, &arena), Some(c));
+}
+
+#[test]
+fn reset_window_accumulation_counts_excluded_nodes() {
+    let (arena, _root, a, b, c, other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(c);
+    list.insert_mru(a);
+    list.insert_mru(other);
+    // b is excluded but its atoms still consume the window, keeping a out of reach.
+    list.reset_node_and_window_ancestors_mru(c, 4, &arena, |node| node.idx != b);
+    assert_eq!(order(&list), vec![c, other, a]);
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "not in the LRU list")]
+fn reset_node_mru_panics_on_a_node_removed_earlier() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.remove_node(NodeIdx_(10));
+    // The cell is still allocated but unlisted; the gated read must reject it.
+    list.reset_node_mru(NodeIdx_(10));
+}
+
+#[test]
+fn in_list_is_false_for_none_and_non_members() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    assert!(!list.in_list(None));
+    assert!(!list.in_list(Some(NodeIdx_(5))));
+    list.insert_mru(NodeIdx_(5));
+    assert!(list.in_list(Some(NodeIdx_(5))));
+}
+
+#[test]
+fn get_lru_where_walks_from_the_tail() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.insert_mru(NodeIdx_(30));
+    assert_eq!(list.get_lru_where(|_| true), Some(NodeIdx_(10)));
+    assert_eq!(
+        list.get_lru_where(|id| id != NodeIdx_(10)),
+        Some(NodeIdx_(20))
+    );
+    assert_eq!(list.get_lru_where(|_| false), None);
+}
+
+#[test]
+fn get_prev_where_walks_toward_the_head_from_a_member() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.insert_mru(NodeIdx_(30));
+    // Order is [30, 20, 10]; 10's predecessors are 20 then 30.
+    assert_eq!(
+        list.get_prev_where(NodeIdx_(10), |_| true),
+        Some(NodeIdx_(20))
+    );
+    assert_eq!(
+        list.get_prev_where(NodeIdx_(10), |id| id != NodeIdx_(20)),
+        Some(NodeIdx_(30))
+    );
+    assert_eq!(list.get_prev_where(NodeIdx_(30), |_| true), None);
+}
+
+#[test]
+#[should_panic(expected = "index out of bounds")]
+fn get_prev_where_panics_on_a_non_member() {
+    let list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.get_prev_where(NodeIdx_(7), |_| true);
+}
+
+#[test]
+#[should_panic(expected = "not in the LRU list")]
+fn get_prev_where_panics_on_a_node_removed_earlier() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.remove_node(NodeIdx_(10));
+    // The cell is still allocated but unlisted; the gated read must reject it.
+    list.get_prev_where(NodeIdx_(10), |_| true);
+}
+
+#[test]
+fn get_prev_before_remove_keeps_the_walk_consistent() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.insert_mru(NodeIdx_(30));
+    // The eviction-cursor contract: compute the predecessor, then remove.
+    let next = list.get_prev_where(NodeIdx_(10), |_| true);
+    list.remove_node(NodeIdx_(10));
+    assert_eq!(next, Some(NodeIdx_(20)));
+    assert!(list.in_list(next));
+    assert_eq!(
+        list.get_prev_where(NodeIdx_(20), |_| true),
+        Some(NodeIdx_(30))
+    );
+    list.validate();
+}
+
+#[test]
+fn insert_mru_grows_the_cell_table_one_by_one() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(0));
+    list.insert_mru(NodeIdx_(1));
+    list.insert_mru(NodeIdx_(2));
+    assert_eq!(order(&list), vec![NodeIdx_(2), NodeIdx_(1), NodeIdx_(0)]);
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "broken prev link")]
+fn validate_panics_on_a_corrupted_prev_link() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    let cell = UnifiedLRUList::cell_of_(NodeIdx_(10));
+    list.cells[cell.0].prev = cell;
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "membership mismatch")]
+fn validate_panics_on_a_linked_cell_without_the_flag() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.cells[UnifiedLRUList::cell_of_(NodeIdx_(10)).0].in_list = false;
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "membership mismatch")]
+fn validate_panics_on_a_flagged_unlinked_cell() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.insert_mru(NodeIdx_(20));
+    list.remove_node(NodeIdx_(20));
+    // The cell is reset but a stray flag claims membership.
+    list.cells[UnifiedLRUList::cell_of_(NodeIdx_(20)).0].in_list = true;
+    list.validate();
+}
+
+#[test]
+fn reset_window_ancestors_mru_is_a_noop_on_a_zero_window() {
+    let (arena, _root, a, _b, c, _other) = arena_chain();
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(a);
+    list.insert_mru(c);
+    list.reset_node_and_window_ancestors_mru(c, 0, &arena, |_| true);
+    assert_eq!(order(&list), vec![c, a]);
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "out of bounds")]
+fn validate_panics_on_an_out_of_range_link() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.cells[UnifiedLRUList::cell_of_(NodeIdx_(10)).0].next = CellId(99);
+    list.validate();
+}
+
+#[test]
+#[should_panic(expected = "length mismatch")]
+fn validate_panics_on_a_desynced_member_counter() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(10));
+    list.len = 2;
+    list.validate();
+}
+
+#[test]
+fn len_drops_to_zero_after_all_members_removed() {
+    let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    list.insert_mru(NodeIdx_(1));
+    list.remove_node(NodeIdx_(1));
+    assert_eq!(list.len(), 0);
+    assert_eq!(list.iter().count(), 0);
+    list.validate();
+}
+
+#[test]
+fn check_linked_list_accepts_a_clean_list() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.insert_mru(NodeIdx_(1));
+    lru.insert_mru(NodeIdx_(2));
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert!(errors.is_empty());
+}
+
+#[test]
+fn check_linked_list_reports_a_broken_prev() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.insert_mru(NodeIdx_(1));
+    lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].prev = UnifiedLRUList::cell_of_(NodeIdx_(0));
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert_eq!(errors.len(), 1);
+    assert!(errors[0].contains("broken prev at node 0"));
+}
+
+#[test]
+fn check_linked_list_reports_an_unflagged_member() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.insert_mru(NodeIdx_(1));
+    lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].in_list = false;
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert_eq!(errors.len(), 1);
+    assert!(errors[0].contains("node 0 in list not flagged"));
+}
+
+#[test]
+fn check_linked_list_reports_a_cycle() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.insert_mru(NodeIdx_(1));
+    // 0's next loops back to 1 instead of reaching the tail.
+    lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].next = UnifiedLRUList::cell_of_(NodeIdx_(1));
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert!(errors.iter().any(|e| e.contains("cycle at node 1")));
+}
+
+#[test]
+fn check_linked_list_reports_a_count_mismatch() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.len = 2;
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert_eq!(errors.len(), 1);
+    assert!(errors[0].contains("list=1 != len=2"));
+}
+
+#[test]
+fn check_linked_list_reports_an_out_of_bounds_link() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].next = CellId(999);
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert!(errors.iter().any(|e| e.contains("cell 999 out of bounds")));
+}
+
+#[test]
+fn check_linked_list_reports_a_broken_tail_backlink() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.cells[TAIL.0].prev = HEAD;
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert_eq!(errors.len(), 1);
+    assert!(errors[0].contains("broken tail backlink"));
+}
+
+#[test]
+fn check_linked_list_reports_a_flagged_unreachable_cell() {
+    let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
+    lru.insert_mru(NodeIdx_(0));
+    lru.insert_mru(NodeIdx_(1));
+    lru.remove_node(NodeIdx_(0));
+    // Re-flag the unlinked cell without relinking it.
+    lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].in_list = true;
+    let mut errors = Vec::new();
+    lru.check_linked_list_("[t]", &mut errors);
+    assert!(
+        errors
+            .iter()
+            .any(|e| e.contains("node 0 flagged but unreachable"))
+    );
+}
+
+// Eviction priority keys.
+
+// A node with distinct field values: last_access 5, creation 7, hits 3, priority 9.
+fn arena_with_node() -> (NodeArena<Vec<i64>>, NodeIdx_) {
+    let mut arena: NodeArena<Vec<i64>> = NodeArena::new(vec![FULL], /* page_size = */ 1);
+    let root = arena.root();
+    let a = arena
+        .alloc_child(
+            root,
+            /* key = */ vec![1],
+            /* priority = */ 9,
+            /* extra_key = */ None,
+        )
+        .unwrap();
+    let node = arena.node_mut(a);
+    node.last_access_counter = 5;
+    node.creation_counter = 7;
+    node.hit_count = 3;
+    (arena, NodeIdx_(a.0))
+}
+
+#[test]
+fn each_strategy_maps_its_node_fields_into_the_key() {
+    let (arena, a) = arena_with_node();
+    let node = arena.node(NodeIdx_(a.0));
+    assert_eq!(LruStrategy.get_priority(node), PriorityKey(5, 0));
+    assert_eq!(LfuStrategy.get_priority(node), PriorityKey(3, 5));
+    assert_eq!(FifoStrategy.get_priority(node), PriorityKey(7, 0));
+    assert_eq!(MruStrategy.get_priority(node), PriorityKey(-5, 0));
+    assert_eq!(FiloStrategy.get_priority(node), PriorityKey(-7, 0));
+    assert_eq!(PriorityStrategy.get_priority(node), PriorityKey(9, 5));
+}
+
+#[test]
+fn slru_segments_on_the_protected_threshold() {
+    let (mut arena, a) = arena_with_node();
+    let slru = SlruStrategy {
+        protected_threshold: 2,
+    };
+    // 3 hits >= threshold 2: protected segment.
+    assert_eq!(
+        slru.get_priority(arena.node(NodeIdx_(a.0))),
+        PriorityKey(1, 5)
+    );
+    // Exactly at the threshold counts as protected.
+    arena.node_mut(NodeIdx_(a.0)).hit_count = 2;
+    assert_eq!(
+        slru.get_priority(arena.node(NodeIdx_(a.0))),
+        PriorityKey(1, 5)
+    );
+    arena.node_mut(NodeIdx_(a.0)).hit_count = 1;
+    assert_eq!(
+        slru.get_priority(arena.node(NodeIdx_(a.0))),
+        PriorityKey(0, 5)
+    );
+}
+
+#[test]
+fn get_eviction_strategy_resolves_each_policy_name() {
+    let (arena, a) = arena_with_node();
+    let node = arena.node(NodeIdx_(a.0));
+    // Distinct node fields make each policy's key identify its strategy.
+    let cases = [
+        ("lru", PriorityKey(5, 0)),
+        ("LFU", PriorityKey(3, 5)),
+        ("fifo", PriorityKey(7, 0)),
+        ("mru", PriorityKey(-5, 0)),
+        ("filo", PriorityKey(-7, 0)),
+        ("priority", PriorityKey(9, 5)),
+        ("slru", PriorityKey(1, 5)),
+    ];
+    for (policy, expected) in cases {
+        assert_eq!(
+            get_eviction_strategy::<Vec<i64>>(policy).get_priority(node),
+            expected,
+            "policy {policy}"
+        );
+    }
+}
+
+#[test]
+fn eviction_policy_names_are_case_insensitive() {
+    let (arena, a) = arena_with_node();
+    let node = arena.node(NodeIdx_(a.0));
+    // Mixed-case names resolve to the same strategies as their lowercase forms.
+    assert_eq!(
+        get_eviction_strategy::<Vec<i64>>("LRU").get_priority(node),
+        PriorityKey(5, 0)
+    );
+    assert_eq!(
+        get_eviction_strategy::<Vec<i64>>("Priority").get_priority(node),
+        PriorityKey(9, 5)
+    );
+}
+
+#[test]
+fn get_eviction_strategy_slru_default_threshold_is_two() {
+    let (mut arena, a) = arena_with_node();
+    let slru = get_eviction_strategy::<Vec<i64>>("slru");
+    // Exactly 2 hits is protected under the factory default; 1 is not.
+    arena.node_mut(NodeIdx_(a.0)).hit_count = 2;
+    assert_eq!(
+        slru.get_priority(arena.node(NodeIdx_(a.0))),
+        PriorityKey(1, 5)
+    );
+    arena.node_mut(NodeIdx_(a.0)).hit_count = 1;
+    assert_eq!(
+        slru.get_priority(arena.node(NodeIdx_(a.0))),
+        PriorityKey(0, 5)
+    );
+}
+
+#[test]
+#[should_panic(expected = "Unknown eviction policy: random. Supported policies:")]
+fn get_eviction_strategy_panics_on_an_unknown_policy() {
+    get_eviction_strategy::<Vec<i64>>("Random");
+}
+
+#[test]
+fn priority_keys_order_lexicographically() {
+    assert!(PriorityKey(0, 9) < PriorityKey(1, 0));
+    assert!(PriorityKey(1, 2) < PriorityKey(1, 3));
+}

rust/mem-cache/src/unified_lru_list.rs

@@ -0,0 +1,501 @@
+//! Self-contained LRU order over `NodeIdx_`s: MRU at the head side, LRU at the
+//! tail side. Node semantics stay with callers through predicates; the reset
+//! walks read parent links from the arena.
+
+use std::collections::HashSet;
+
+use crate::node::ChildKeyType;
+use crate::node::Node;
+use crate::node::NodeArena;
+use crate::node::{NodeIdx_, ValueSlotIdx};
+
+/// Index into the cell table; distinct from `NodeIdx_` so shifted and unshifted
+/// ids cannot be mixed.
+#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
+struct CellId(usize);
+
+/// Head sentinel cell.
+const HEAD: CellId = CellId(0);
+/// Tail sentinel cell.
+const TAIL: CellId = CellId(1);
+/// Table offset: node ids map to cells after the sentinels.
+const OFFSET: usize = 2;
+
+/// One doubly-linked cell; a node's cell lives at `NodeIdx_ + OFFSET`.
+#[derive(Clone, Copy, Default)]
+struct Cell {
+    prev: CellId,
+    next: CellId,
+    in_list: bool,
+}
+
+/// LRU list over `NodeIdx_`s, with head/tail sentinel cells keeping the link
+/// operations branchless. External APIs take `NodeIdx_`s; internal (`_`-suffixed)
+/// interfaces work on `CellId`s.
+pub struct UnifiedLRUList {
+    /// The (component × tier) value slot whose lock gates this list's walkers.
+    slot: ValueSlotIdx,
+    /// Cell table indexed by `NodeIdx_ + OFFSET`; cells 0/1 are the sentinels.
+    cells: Vec<Cell>,
+    /// Number of member cells, excluding the sentinels.
+    len: usize,
+}
+
+impl UnifiedLRUList {
+    pub fn new(slot: ValueSlotIdx) -> Self {
+        UnifiedLRUList {
+            slot,
+            cells: vec![
+                // Sentinels link to each other and stay permanently flagged so
+                // the gated cell accessors admit them.
+                Cell {
+                    prev: HEAD,
+                    next: TAIL,
+                    in_list: true,
+                },
+                Cell {
+                    prev: HEAD,
+                    next: TAIL,
+                    in_list: true,
+                },
+            ],
+            len: 0,
+        }
+    }
+
+    // ==== List operations ====
+
+    fn add_node_after_(&mut self, prev: CellId, cell: CellId) {
+        self.new_cell_(cell);
+        let next = self.cell_(prev).next;
+        self.connect_(cell, next);
+        self.connect_(prev, cell);
+    }
+
+    fn add_node_(&mut self, cell: CellId) {
+        self.add_node_after_(HEAD, cell);
+    }
+
+    fn remove_node_(&mut self, cell: CellId) {
+        let Cell { prev, next, .. } = *self.cell_(cell);
+        self.connect_(prev, next);
+        // Unflag the cell; the stale prev/next are never read while unlisted.
+        self.cell_mut_(cell).in_list = false;
+        self.len -= 1;
+    }
+
+    // ==== NodeIdx_ <-> CellId plumbing ====
+
+    /// The node's cell slot; the only `NodeIdx_` -> `CellId` crossing.
+    fn cell_of_(node_id: NodeIdx_) -> CellId {
+        CellId(node_id.0 + OFFSET)
+    }
+
+    /// The cell's node; the only `CellId` -> `NodeIdx_` crossing.
+    fn node_of_(cell: CellId) -> NodeIdx_ {
+        NodeIdx_(cell.0 - OFFSET)
+    }
+
+    /// The cell, asserting it is linked (sentinels always are).
+    #[track_caller]
+    fn cell_(&self, id: CellId) -> &Cell {
+        let cell = &self.cells[id.0];
+        assert!(
+            cell.in_list,
+            "node {} not in the LRU list",
+            Self::node_of_(id)
+        );
+        cell
+    }
+
+    #[track_caller]
+    fn cell_mut_(&mut self, id: CellId) -> &mut Cell {
+        let cell = &mut self.cells[id.0];
+        assert!(
+            cell.in_list,
+            "node {} not in the LRU list",
+            Self::node_of_(id)
+        );
+        cell
+    }
+
+    /// Admit an unlisted cell: grow the table to cover it, then flag and count
+    /// it before any connections.
+    fn new_cell_(&mut self, cell: CellId) {
+        if cell.0 >= self.cells.len() {
+            self.cells.resize(cell.0 + 1, Cell::default());
+        }
+        assert!(
+            !self.cells[cell.0].in_list,
+            "new_cell_: cell {cell:?} already in the LRU list"
+        );
+        self.cells[cell.0].in_list = true;
+        self.len += 1;
+    }
+
+    /// Whether the cell is linked into the list; safe on cells beyond the table.
+    fn in_list_(&self, cell: CellId) -> bool {
+        self.cells.get(cell.0).is_some_and(|cell| cell.in_list)
+    }
+
+    /// Link `a -> b`.
+    fn connect_(&mut self, a: CellId, b: CellId) {
+        self.cell_mut_(a).next = b;
+        self.cell_mut_(b).prev = a;
+    }
+
+    /// Insert a node as the most-recently-used; panics if already a member.
+    pub fn insert_mru(&mut self, node_id: NodeIdx_) {
+        self.add_node_(Self::cell_of_(node_id));
+    }
+
+    /// Remove a member node, resetting its cell; panics if not a member.
+    pub fn remove_node(&mut self, node_id: NodeIdx_) {
+        self.remove_node_(Self::cell_of_(node_id));
+    }
+
+    /// Move a member node back to the most-recently-used position.
+    pub fn reset_node_mru(&mut self, node_id: NodeIdx_) {
+        let cell = Self::cell_of_(node_id);
+        self.remove_node_(cell);
+        self.add_node_(cell);
+    }
+
+    /// Re-rank the `should_include` nodes from `node_id` up to its root
+    /// (exclusive) as the MRU run, deepest first.
+    pub fn reset_node_and_parents_mru<K: ChildKeyType>(
+        &mut self,
+        node_id: NodeIdx_,
+        arena: &NodeArena<K>,
+        mut should_include: impl FnMut(&Node<K>) -> bool,
+    ) {
+        let mut prev = HEAD;
+        let mut cur = node_id;
+        loop {
+            let node = arena.node(cur);
+            let Some(parent) = node.try_parent() else {
+                break;
+            };
+            if should_include(node) {
+                let cell = Self::cell_of_(cur);
+                self.remove_node_(cell);
+                self.add_node_after_(prev, cell);
+                prev = cell;
+            }
+            cur = parent;
+        }
+    }
+
+    /// Like `reset_node_and_parents_mru`, stopping once `window_size` atoms
+    /// are covered; excluded ancestors consume the window too.
+    pub fn reset_node_and_window_ancestors_mru<K: ChildKeyType>(
+        &mut self,
+        node_id: NodeIdx_,
+        window_size: usize,
+        arena: &NodeArena<K>,
+        mut should_include: impl FnMut(&Node<K>) -> bool,
+    ) {
+        let mut prev = HEAD;
+        let mut accumulated = 0;
+        let mut cur = node_id;
+        while accumulated < window_size {
+            let node = arena.node(cur);
+            let Some(parent) = node.try_parent() else {
+                break;
+            };
+            if should_include(node) {
+                let cell = Self::cell_of_(cur);
+                self.remove_node_(cell);
+                self.add_node_after_(prev, cell);
+                prev = cell;
+            }
+            accumulated += node.key.atom_len();
+            cur = parent;
+        }
+    }
+
+    /// Whether the node is a member (`None` is never a member).
+    pub fn in_list(&self, node_id: Option<NodeIdx_>) -> bool {
+        node_id.is_some_and(|id| self.in_list_(Self::cell_of_(id)))
+    }
+
+    /// The nearest predecessor of `cell` satisfying `pred`, walking toward the
+    /// head; `cell` itself is excluded.
+    fn get_prev_where_(
+        &self,
+        cell: CellId,
+        mut pred: impl FnMut(NodeIdx_) -> bool,
+    ) -> Option<NodeIdx_> {
+        let mut cell = self.cell_(cell).prev;
+        while cell != HEAD {
+            let node = Self::node_of_(cell);
+            if pred(node) {
+                return Some(node);
+            }
+            cell = self.cell_(cell).prev;
+        }
+        None
+    }
+
+    /// The nearest predecessor of a member satisfying `pred`; panics if
+    /// `node_id` is not a member.
+    pub fn get_prev_where(
+        &self,
+        node_id: NodeIdx_,
+        pred: impl FnMut(NodeIdx_) -> bool,
+    ) -> Option<NodeIdx_> {
+        self.get_prev_where_(Self::cell_of_(node_id), pred)
+    }
+
+    /// The least-recent member whose lock on the list's own slot is free.
+    pub fn get_lru_no_lock<K: ChildKeyType>(&self, arena: &NodeArena<K>) -> Option<NodeIdx_> {
+        self.get_lru_where(|id| arena.node(id).lock_ref_(self.slot) == 0)
+    }
+
+    /// The nearest more-recent member whose lock on the list's own slot is
+    /// free, from `node_id`.
+    pub fn get_prev_no_lock<K: ChildKeyType>(
+        &self,
+        node_id: NodeIdx_,
+        arena: &NodeArena<K>,
+    ) -> Option<NodeIdx_> {
+        self.get_prev_where(node_id, |id| arena.node(id).lock_ref_(self.slot) == 0)
+    }
+
+    /// The least-recently-used member satisfying `pred`.
+    pub fn get_lru_where(&self, pred: impl FnMut(NodeIdx_) -> bool) -> Option<NodeIdx_> {
+        self.get_prev_where_(TAIL, pred)
+    }
+
+    /// Number of member cells, excluding the sentinels.
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    /// Materialize the current members from most to least recent.
+    ///
+    /// Inspection callers need an owned snapshot across the Python boundary;
+    /// the linked-list iterator itself never escapes the Rust core.
+    pub(crate) fn snapshot_node_ids(&self) -> Vec<NodeIdx_> {
+        let mut node_ids = Vec::with_capacity(self.len);
+        let mut cell = self.cell_(HEAD).next;
+        while cell != TAIL {
+            node_ids.push(Self::node_of_(cell));
+            cell = self.cell_(cell).next;
+        }
+        node_ids
+    }
+
+    // ==== Test-only conveniences ====
+
+    /// The members, MRU to LRU.
+    #[cfg(test)]
+    pub fn iter(&self) -> impl Iterator<Item = NodeIdx_> + '_ {
+        let mut cell = self.cell_(HEAD).next;
+        std::iter::from_fn(move || {
+            if cell == TAIL {
+                return None;
+            }
+            let node = Self::node_of_(cell);
+            cell = self.cell_(cell).next;
+            Some(node)
+        })
+    }
+
+    /// Panics if the links, membership flags, or member counter are inconsistent.
+    /// Reads cells raw: it inspects possibly-inconsistent state that the gated
+    /// accessors would reject.
+    #[cfg(test)]
+    pub fn validate(&self) {
+        let mut count = 0;
+        let mut prev = HEAD;
+        let mut cell = self.cells[HEAD.0].next;
+        while cell != TAIL {
+            assert!(
+                cell.0 >= OFFSET && cell.0 < self.cells.len(),
+                "validate: cell {cell:?} out of bounds"
+            );
+            assert_eq!(
+                self.cells[cell.0].prev, prev,
+                "validate: broken prev link at cell {cell:?}"
+            );
+            assert!(
+                self.cells[cell.0].in_list,
+                "validate: membership mismatch at cell {cell:?}"
+            );
+            count += 1;
+            assert!(count <= self.len, "validate: cycle detected");
+            prev = cell;
+            cell = self.cells[cell.0].next;
+        }
+        assert_eq!(self.cells[TAIL.0].prev, prev, "validate: broken tail link");
+        assert_eq!(count, self.len, "validate: length mismatch");
+        let flagged = self
+            .cells
+            .iter()
+            .skip(OFFSET)
+            .filter(|cell| cell.in_list)
+            .count();
+        assert_eq!(flagged, self.len, "validate: membership mismatch");
+    }
+    /// Test-only: desynchronize `len` to force integrity errors.
+    #[cfg(test)]
+    pub(crate) fn bump_len_for_test(&mut self) {
+        self.len += 1;
+    }
+
+    /// Walk a LRU doubly-linked list, collect integrity errors.
+    pub(crate) fn check_linked_list_(&self, label: &str, errors: &mut Vec<String>) {
+        let mut visited: HashSet<usize> = HashSet::new();
+        let mut prev = HEAD;
+        let mut x = self.cells[HEAD.0].next;
+        while x != TAIL {
+            if x.0 < OFFSET {
+                errors.push(format!("{label} broken chain: link points at a sentinel"));
+                break;
+            }
+            let Some(cell) = self.cells.get(x.0) else {
+                errors.push(format!("{label} broken chain: cell {} out of bounds", x.0));
+                break;
+            };
+            if cell.prev != prev {
+                errors.push(format!("{label} broken prev at node {}", Self::node_of_(x)));
+            }
+            if !cell.in_list {
+                errors.push(format!(
+                    "{label} node {} in list not flagged",
+                    Self::node_of_(x)
+                ));
+            }
+            if !visited.insert(x.0) {
+                errors.push(format!("{label} cycle at node {}", Self::node_of_(x)));
+                break;
+            }
+            prev = x;
+            x = cell.next;
+        }
+        // The tail backlink closes the list onto the last visited member.
+        if x == TAIL && self.cells[TAIL.0].prev != prev {
+            errors.push(format!("{label} broken tail backlink"));
+        }
+        // Every flagged member cell must be reachable from the head.
+        for (idx, cell) in self.cells.iter().enumerate().skip(OFFSET) {
+            if cell.in_list && !visited.contains(&idx) {
+                errors.push(format!(
+                    "{label} node {} flagged but unreachable",
+                    idx - OFFSET
+                ));
+            }
+        }
+        if visited.len() != self.len {
+            errors.push(format!(
+                "{label} list={} != len={}",
+                visited.len(),
+                self.len
+            ));
+        }
+    }
+}
+
+// Eviction priority keys.
+
+/// Eviction-priority key, ordered lexicographically; lower evicts first.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
+pub struct PriorityKey(pub i64, pub i64);
+
+/// Ranks nodes for eviction; lower priority evicts first.
+pub trait EvictionStrategy<K: ChildKeyType> {
+    /// The node's eviction priority.
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey;
+}
+
+/// Least-recently-used.
+pub struct LruStrategy;
+
+impl<K: ChildKeyType> EvictionStrategy<K> for LruStrategy {
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey {
+        PriorityKey(node.last_access_counter, 0)
+    }
+}
+
+/// Least-frequently-used; LRU within a hit count.
+pub struct LfuStrategy;
+
+impl<K: ChildKeyType> EvictionStrategy<K> for LfuStrategy {
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey {
+        PriorityKey(node.hit_count, node.last_access_counter)
+    }
+}
+
+/// First-in-first-out over creation order.
+pub struct FifoStrategy;
+
+impl<K: ChildKeyType> EvictionStrategy<K> for FifoStrategy {
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey {
+        PriorityKey(node.creation_counter, 0)
+    }
+}
+
+/// Most-recently-used first.
+pub struct MruStrategy;
+
+impl<K: ChildKeyType> EvictionStrategy<K> for MruStrategy {
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey {
+        PriorityKey(-node.last_access_counter, 0)
+    }
+}
+
+/// First-in-last-out over creation order.
+pub struct FiloStrategy;
+
+impl<K: ChildKeyType> EvictionStrategy<K> for FiloStrategy {
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey {
+        PriorityKey(-node.creation_counter, 0)
+    }
+}
+
+/// Priority-aware: lower node priority evicts first, LRU within a priority.
+pub struct PriorityStrategy;
+
+impl<K: ChildKeyType> EvictionStrategy<K> for PriorityStrategy {
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey {
+        PriorityKey(node.priority, node.last_access_counter)
+    }
+}
+
+/// Segmented LRU: probationary nodes (hits below the threshold) evict before
+/// protected ones, LRU within a segment.
+pub struct SlruStrategy {
+    pub protected_threshold: i64,
+}
+
+impl<K: ChildKeyType> EvictionStrategy<K> for SlruStrategy {
+    fn get_priority(&self, node: &Node<K>) -> PriorityKey {
+        PriorityKey(
+            (node.hit_count >= self.protected_threshold) as i64,
+            node.last_access_counter,
+        )
+    }
+}
+
+/// The strategy for an eviction-policy name.
+pub fn get_eviction_strategy<K: ChildKeyType>(policy: &str) -> Box<dyn EvictionStrategy<K> + Send> {
+    match policy.to_lowercase().as_str() {
+        "lru" => Box::new(LruStrategy),
+        "lfu" => Box::new(LfuStrategy),
+        "fifo" => Box::new(FifoStrategy),
+        "mru" => Box::new(MruStrategy),
+        "filo" => Box::new(FiloStrategy),
+        "priority" => Box::new(PriorityStrategy),
+        "slru" => Box::new(SlruStrategy {
+            protected_threshold: 2,
+        }),
+        other => panic!(
+            "Unknown eviction policy: {other}. Supported policies: \
+             'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority', 'slru'."
+        ),
+    }
+}
+#[cfg(test)]
+#[path = "tests/unified_lru_list.rs"]
+mod tests;

rust/mem-cache/torch_2_13_compat.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <stdexcept>
+#include <torch/version.h>
+
+#if TORCH_VERSION_MAJOR > 2 ||                                                 \
+    (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 13)
+// Keep tch 0.24's removed alignment wrappers as explicit runtime errors.
+#define align_as(...)                                                          \
+  alias();                                                                     \
+  throw std::runtime_error("align_as is unavailable in PyTorch 2.13+")
+#define align_tensors(...)                                                     \
+  autograd::variable_list{};                                                   \
+  throw std::runtime_error("align_tensors is unavailable in PyTorch 2.13+")
+#endif

scripts/ci/cuda/ci_install_dependency.sh

@@ -277,11 +277,24 @@ clean_site_packages() {
 }
 
 setup_cargo_cache() {
+    if [ "${SGLANG_BUILD_RUST_EXTS:-}" = "none" ]; then
+        echo "Using prebuilt Rust extensions; skipping Cargo target setup"
+        mark_step_done "${FUNCNAME[0]}"
+        return
+    fi
+
     # actions/checkout's `git clean -ffdx` deletes the gitignored in-repo
     # rust/target, so every job recompiles the whole dependency graph. Move the
     # target dir out of the tree: setuptools-rust has no target-dir option of its
     # own and defers to CARGO_TARGET_DIR, which uv passes to the build backend.
     export CARGO_TARGET_DIR="${HOME}/.cache/sglang-cargo-target"
+    local cargo_target_lock="${HOME}/.cache/sglang-cargo-target.lock"
+    mkdir -p "${HOME}/.cache"
+    exec 9>"${cargo_target_lock}"
+    echo "Waiting for exclusive cargo target lock: ${cargo_target_lock}"
+    flock --exclusive 9
+    CARGO_TARGET_LOCK_HELD=1
+    echo "Acquired cargo target lock"
     mkdir -p "${CARGO_TARGET_DIR}"
 
     # Same disk-pressure guard as the uv cache in ci_cleanup_venv.sh (which
@@ -298,6 +311,15 @@ setup_cargo_cache() {
     mark_step_done "${FUNCNAME[0]}"
 }
 
+release_cargo_cache_lock() {
+    if [ "${CARGO_TARGET_LOCK_HELD:-0}" = "1" ]; then
+        flock --unlock 9
+        exec 9>&-
+        CARGO_TARGET_LOCK_HELD=0
+        echo "Released cargo target lock"
+    fi
+}
+
 setup_pip_toolchain() {
     if [ "$USE_VENV" = "1" ]; then
         # The bootstrap upgrade hit system pip; this upgrades the venv's own.
@@ -473,10 +495,19 @@ require_prebuilt_rust_exts() {
     for module in server grpc multimodal; do
         [ -f "python/sglang/srt/rust_extensions/_${module}${suffix}" ] || missing+=("${module}")
     done
+    [ -f "python/sglang/srt/mem_cache/rust_tree_core/mem_cache${suffix}" ] \
+        || missing+=("mem_cache")
+    [ -f "python/sglang/srt/mem_cache/rust_tree_core/mem_cache_inspection${suffix}" ] \
+        || missing+=("mem_cache_inspection")
     if [ ${#missing[@]} -gt 0 ]; then
         echo "::warning::no prebuilt Rust extension ${suffix} for: ${missing[*]}; building from source"
         ls -l python/sglang/srt/rust_extensions/_*.so 2>/dev/null || echo "(no extension modules at all)"
+        ls -l python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so 2>/dev/null || true
         export SGLANG_BUILD_RUST_EXTS=
+        export SGLANG_RUST_BUILD_MODE=auto
+        if [ -n "${GITHUB_ENV:-}" ]; then
+            echo "SGLANG_RUST_BUILD_MODE=auto" >> "${GITHUB_ENV}"
+        fi
         mark_step_done "${FUNCNAME[0]}"
         return
     fi
@@ -846,14 +877,15 @@ main() {
     install_apt_packages
     install_gdrcopy
     clean_site_packages
-    setup_cargo_cache
     require_prebuilt_rust_exts
     setup_pip_toolchain
     remove_stale_cuda12_nvidia_wheels
     uninstall_stale_flashinfer
     install_pytorch_stack
     install_cuda12_deepep_wheel
+    setup_cargo_cache
     install_sglang
+    release_cargo_cache_lock
     install_nccl
     # Diffusion B200 CI imports torch inside install_sglang_kernel after removing
     # stale CUDA 12 NVIDIA wheels, so opt into one early LD_LIBRARY_PATH refresh.

scripts/ci/utils/stage_rust_ext_modules.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
-# Copy the built PyO3 extension modules into rust-ext-staging/rust_extensions/ for
-# upload-artifact. Shared by both jobs of _pr-test-rust-ext-build.yml, so the
-# archive layout and the module-count check cannot drift between them.
+# Copy the built PyO3 extension modules into their package-relative paths under
+# rust-ext-staging/. Shared by both jobs of _pr-test-rust-ext-build.yml, so the
+# archive layout and module-count checks cannot drift between them.
 #
 # MAX_GLIBC (optional): also reject a module requiring a newer GLIBC symbol
 # version than the test runners have. Only set where the modules were just
@@ -32,6 +32,23 @@ for module in server grpc multimodal; do
     cp "${found[@]}" rust-ext-staging/rust_extensions/
     built+=("${found[@]}")
 done
+
+mkdir -p rust-ext-staging/mem_cache/rust_tree_core
+for module in mem_cache mem_cache_inspection; do
+    tree_core=(python/sglang/srt/mem_cache/rust_tree_core/"${module}".*.so)
+    if [ ${#tree_core[@]} -eq 0 ]; then
+        echo "::error::no Rust TreeCore ${module} extension module found"
+        exit 1
+    fi
+    tree_core_suffixes=$(printf '%s\n' "${tree_core[@]##*/${module}}" | sort)
+    if [ "${tree_core_suffixes}" != "${expected_suffixes}" ]; then
+        echo "::error::Rust TreeCore ${module} extension does not match the interpreter set"
+        printf 'have:\n%s\nwant:\n%s\n' "${tree_core_suffixes}" "${expected_suffixes}"
+        exit 1
+    fi
+    cp "${tree_core[@]}" rust-ext-staging/mem_cache/rust_tree_core/
+    built+=("${tree_core[@]}")
+done
 max_allowed="${MAX_GLIBC:-}"
 [ -n "${max_allowed}" ] || exit 0
 

scripts/release/prepare_sglang_wheel.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""Repair an SGLang wheel and smoke-test its production Rust TreeCore."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import textwrap
+import zipfile
+from email.parser import BytesParser
+from pathlib import Path, PurePosixPath
+
+_LIBTORCH_EXCLUDES = (
+    "libc10.so",
+    "libc10_cuda.so",
+    "libtorch.so",
+    "libtorch_cpu.so",
+    "libtorch_cuda.so",
+    "libtorch_python.so",
+)
+_TREE_CORE_DIR = PurePosixPath("sglang/srt/mem_cache/rust_tree_core")
+_BINDING_CLASSES = (
+    "RustUnifiedTreeCoreBinding",
+    "RustBigramUnifiedTreeCoreBinding",
+    "TreeCoreInitParamsBinding",
+)
+
+
+def _single_wheel(directory: Path) -> Path:
+    wheels = sorted(directory.glob("*.whl"))
+    if len(wheels) != 1:
+        raise RuntimeError(f"expected one wheel in {directory}, found {wheels}")
+    return wheels[0]
+
+
+def _metadata(wheel: Path) -> tuple[str, str]:
+    with zipfile.ZipFile(wheel) as archive:
+        metadata_files = [
+            name for name in archive.namelist() if name.endswith(".dist-info/METADATA")
+        ]
+        if len(metadata_files) != 1:
+            raise RuntimeError(
+                f"expected one METADATA file in {wheel}, found {metadata_files}"
+            )
+        metadata = BytesParser().parsebytes(archive.read(metadata_files[0]))
+    return str(metadata["Name"]), str(metadata["Version"])
+
+
+def _smoke_test_tree_core(wheel: Path) -> None:
+    with tempfile.TemporaryDirectory(prefix="sglang-wheel-smoke-") as temp_dir:
+        root = Path(temp_dir)
+        with zipfile.ZipFile(wheel) as archive:
+            names = archive.namelist()
+            inspection_modules = [
+                name
+                for name in names
+                if PurePosixPath(name).parent == _TREE_CORE_DIR
+                and PurePosixPath(name).name.startswith("mem_cache_inspection")
+                and name.endswith(".so")
+            ]
+            if inspection_modules:
+                raise RuntimeError(
+                    f"production wheel contains inspection modules: {inspection_modules}"
+                )
+            production_modules = [
+                name
+                for name in names
+                if PurePosixPath(name).parent == _TREE_CORE_DIR
+                and PurePosixPath(name).name.startswith("mem_cache.")
+                and name.endswith(".so")
+            ]
+            if len(production_modules) != 1:
+                raise RuntimeError(
+                    "expected one production Rust TreeCore module, found "
+                    f"{production_modules}"
+                )
+
+        install_dir = root / "installed"
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "--disable-pip-version-check",
+                "--no-compile",
+                "--no-deps",
+                "--no-index",
+                "--target",
+                os.fspath(install_dir),
+                os.fspath(wheel),
+            ],
+            check=True,
+        )
+
+        smoke_script = textwrap.dedent(f"""
+            import sys
+            import types
+            from pathlib import Path
+
+            site_packages = Path({os.fspath(install_dir)!r}).resolve()
+            sys.path.insert(0, str(site_packages))
+            package = types.ModuleType("sglang")
+            package.__package__ = "sglang"
+            package.__path__ = [str(site_packages / "sglang")]
+            sys.modules["sglang"] = package
+
+            from sglang.srt.mem_cache.rust_tree_core.extension import bindings
+
+            module_path = Path(bindings.__file__).resolve()
+            if site_packages not in module_path.parents:
+                raise RuntimeError(
+                    f"loaded TreeCore outside installed wheel: {{module_path}}"
+                )
+            if bindings.__name__ != "sglang.srt.mem_cache.rust_tree_core.mem_cache":
+                raise RuntimeError(
+                    f"loaded unexpected TreeCore module: {{bindings.__name__}}"
+                )
+            for class_name in {_BINDING_CLASSES!r}:
+                binding = getattr(bindings, class_name, None)
+                if binding is None:
+                    raise RuntimeError(
+                        f"production TreeCore is missing {{class_name}}"
+                    )
+                inspection_methods = [
+                    name for name in dir(binding) if name.startswith("inspect_")
+                ]
+                if inspection_methods:
+                    raise RuntimeError(
+                        f"production {{class_name}} exposes inspection methods: "
+                        f"{{inspection_methods}}"
+                    )
+
+            from array import array
+
+            hashes = bindings.get_hash_str(array("q", [1, 2]), None, 1)
+            if len(hashes) != 2 or any(len(value) != 64 for value in hashes):
+                raise RuntimeError(f"unexpected TreeCore hash result: {{hashes}}")
+            """)
+        environment = os.environ.copy()
+        environment["SGLANG_RUST_BUILD_MODE"] = "never"
+        environment.pop("PYTHONPATH", None)
+        subprocess.run(
+            [sys.executable, "-I", "-c", smoke_script],
+            cwd=root,
+            env=environment,
+            check=True,
+        )
+
+
+def _write_github_outputs(path: Path, *, wheel: Path, version: str) -> None:
+    with path.open("a", encoding="utf-8") as output:
+        output.write(f"wheel_filename={wheel.name}\n")
+        output.write(f"wheel_version={version}\n")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("wheel_dir", type=Path)
+    parser.add_argument("--github-output", type=Path)
+    args = parser.parse_args()
+
+    wheel_dir = args.wheel_dir.resolve()
+    source_wheel = _single_wheel(wheel_dir)
+    with tempfile.TemporaryDirectory(
+        prefix="sglang-wheel-repair-", dir=wheel_dir.parent
+    ) as repair_dir:
+        command = [
+            sys.executable,
+            "-m",
+            "auditwheel",
+            "repair",
+            os.fspath(source_wheel),
+            "--wheel-dir",
+            repair_dir,
+        ]
+        for library in _LIBTORCH_EXCLUDES:
+            command.extend(("--exclude", library))
+        subprocess.run(command, check=True)
+
+        repaired_wheel = _single_wheel(Path(repair_dir))
+        name, version = _metadata(repaired_wheel)
+        if name.casefold() != "sglang":
+            raise RuntimeError(f"expected sglang wheel, found {name!r}")
+        _smoke_test_tree_core(repaired_wheel)
+
+        destination = wheel_dir / repaired_wheel.name
+        source_wheel.unlink()
+        shutil.move(repaired_wheel, destination)
+
+    if args.github_output is not None:
+        _write_github_outputs(
+            args.github_output.resolve(), wheel=destination, version=version
+        )
+    print(f"Prepared {destination.name} (sglang {version})")
+
+
+if __name__ == "__main__":
+    main()

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py

@@ -18,13 +18,14 @@ from sglang.test.test_utils import (
     is_in_ci,
     popen_launch_server,
     terminate_and_kill_process_tree,
+    unified_radix_tree_server_env,
 )
 
 DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
 DSV4_DSPARK_MODEL = "deepseek-ai/DeepSeek-V4-Flash-DSpark"
 DSV4_FLASH_LAUNCH_TIMEOUT = 3600
 
-register_cuda_ci(est_time=2400, stage="extra-b", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=4800, stage="extra-b", runner_config="4-gpu-h100")
 
 
 def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label):
@@ -37,6 +38,7 @@ def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label):
 class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
     """DeepSeek V4 Flash FP8 + HiCache + UnifiedRadixCache."""
 
+    tree_core_backend = "python"
     tp_size = 4
     pp_size = 1
     hicache_io_backend = "direct"
@@ -100,10 +102,10 @@ class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCas
             cls.base_url,
             timeout=DSV4_FLASH_LAUNCH_TIMEOUT,
             other_args=cls._server_args(),
-            env={
-                "SGLANG_DSV4_FP4_EXPERTS": "0",
-                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
-            },
+            env=unified_radix_tree_server_env(
+                cls.tree_core_backend,
+                SGLANG_DSV4_FP4_EXPERTS="0",
+            ),
         )
         cls.input_ids = get_input_ids(cls.model, num_samples=18)
 
@@ -127,6 +129,7 @@ class TestUnifiedDeepSeekV4FlashHiCachePageFirstDirect(
 class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
     """DeepSeek V4 Flash FP8 + HiCache L3 (file backend) + UnifiedRadixCache."""
 
+    tree_core_backend = "python"
     l3_prefetch_page_size = 256
     l3_prefetch_prompt_pages = 4
     max_running_requests = 4
@@ -171,11 +174,11 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
                 "--max-running-requests",
                 str(cls.max_running_requests),
             ],
-            env={
-                "SGLANG_DSV4_FP4_EXPERTS": "0",
-                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
-                "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
-            },
+            env=unified_radix_tree_server_env(
+                cls.tree_core_backend,
+                SGLANG_DSV4_FP4_EXPERTS="0",
+                SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR=cls.hicache_dir,
+            ),
         )
 
     @classmethod
@@ -188,6 +191,7 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
 class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
     """DeepSeek V4 Flash EAGLE + HiCache L3 should load from storage."""
 
+    tree_core_backend = "python"
     page_size = 256
     l3_prefetch_page_size = 256
     l3_prefetch_prompt_pages = 4
@@ -248,11 +252,11 @@ class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestC
                 "--speculative-num-draft-tokens",
                 "4",
             ],
-            env={
-                "SGLANG_DSV4_FP4_EXPERTS": "0",
-                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
-                "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
-            },
+            env=unified_radix_tree_server_env(
+                cls.tree_core_backend,
+                SGLANG_DSV4_FP4_EXPERTS="0",
+                SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR=cls.hicache_dir,
+            ),
         )
 
     @classmethod
@@ -385,12 +389,38 @@ class TestUnifiedDeepSeekV4FlashDSparkHiCacheL3(
                 "--speculative-algorithm",
                 "DSPARK",
             ],
-            env={
-                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
-                "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
-            },
+            env=unified_radix_tree_server_env(
+                cls.tree_core_backend,
+                SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR=cls.hicache_dir,
+            ),
         )
 
 
+class TestRustUnifiedDeepSeekV4FlashHiCache(TestUnifiedDeepSeekV4FlashHiCache):
+    tree_core_backend = "rust"
+
+
+class TestRustUnifiedDeepSeekV4FlashHiCachePageFirstDirect(
+    TestUnifiedDeepSeekV4FlashHiCachePageFirstDirect
+):
+    tree_core_backend = "rust"
+
+
+class TestRustUnifiedDeepSeekV4FlashHiCacheL3(TestUnifiedDeepSeekV4FlashHiCacheL3):
+    tree_core_backend = "rust"
+
+
+class TestRustUnifiedDeepSeekV4FlashEagleHiCacheL3(
+    TestUnifiedDeepSeekV4FlashEagleHiCacheL3
+):
+    tree_core_backend = "rust"
+
+
+class TestRustUnifiedDeepSeekV4FlashDSparkHiCacheL3(
+    TestUnifiedDeepSeekV4FlashDSparkHiCacheL3
+):
+    tree_core_backend = "rust"
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py

@@ -9,10 +9,11 @@ from sglang.test.test_utils import (
     CustomTestCase,
     popen_launch_server,
     terminate_and_kill_process_tree,
+    unified_radix_tree_server_env,
 )
 
-register_cuda_ci(est_time=250, stage="base-b", runner_config="2-gpu-large")
-register_amd_ci(est_time=400, suite="stage-b-test-2-gpu-large-amd")
+register_cuda_ci(est_time=500, stage="base-b", runner_config="2-gpu-large")
+register_amd_ci(est_time=800, suite="stage-b-test-2-gpu-large-amd")
 
 FULL_MODEL = "Qwen/Qwen3-32B"
 
@@ -20,6 +21,7 @@ FULL_MODEL = "Qwen/Qwen3-32B"
 class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
     """Full attention."""
 
+    tree_core_backend = "python"
     kl_threshold = 0.0025
 
     @classmethod
@@ -38,7 +40,7 @@ class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
                 "--page-size",
                 "64",
             ],
-            env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env=unified_radix_tree_server_env(cls.tree_core_backend),
         )
         cls.input_ids = get_input_ids(cls.model, num_samples=18)
 
@@ -47,5 +49,9 @@ class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
         terminate_and_kill_process_tree(cls.process, wait_timeout=60)
 
 
+class TestRustUnifiedFullRadixCache(TestUnifiedFullRadixCache):
+    tree_core_backend = "rust"
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py

@@ -70,9 +70,10 @@ from sglang.test.test_utils import (
     CustomTestCase,
     popen_launch_server,
     terminate_and_kill_process_tree,
+    unified_radix_tree_server_env,
 )
 
-register_cuda_ci(est_time=1150, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=2300, stage="base-b", runner_config="1-gpu-large")
 
 _MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inkling")
 _MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test")
@@ -149,6 +150,8 @@ class TestUnifiedHybridBitExact(CustomTestCase):
     decode-region state reuse in general rather than that regression.
     """
 
+    tree_core_backend = "python"
+
     @classmethod
     def setUpClass(cls):
         cls.model = _MODEL_PATH
@@ -168,7 +171,7 @@ class TestUnifiedHybridBitExact(CustomTestCase):
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=other_args,
-            env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env=unified_radix_tree_server_env(cls.tree_core_backend),
         )
 
     @classmethod
@@ -226,7 +229,7 @@ class TestUnifiedHybridLazyBitExact(TestUnifiedHybridBitExact):
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=other_args,
-            env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env=unified_radix_tree_server_env(cls.tree_core_backend),
         )
 
 
@@ -243,6 +246,8 @@ class TestUnifiedHybridHiCacheBitExact(CustomTestCase):
     cannot produce a non-aligned hit length, which this regression needs.
     """
 
+    tree_core_backend = "python"
+
     @classmethod
     def setUpClass(cls):
         cls.model = _MODEL_PATH
@@ -277,7 +282,7 @@ class TestUnifiedHybridHiCacheBitExact(CustomTestCase):
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=other_args,
-            env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env=unified_radix_tree_server_env(cls.tree_core_backend),
         )
         cls.input_ids = get_input_ids(
             tokenizer_path=cls.model, num_samples=9, trust_remote_code=True
@@ -333,6 +338,8 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
     environment override; a regression there surfaces here as a nonzero KL.
     """
 
+    tree_core_backend = "python"
+
     @classmethod
     def setUpClass(cls):
         cls.model = _MODEL_PATH
@@ -360,10 +367,7 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=other_args,
-            env={
-                **os.environ,
-                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
-            },
+            env=unified_radix_tree_server_env(cls.tree_core_backend),
         )
 
     @classmethod
@@ -396,5 +400,21 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
         self._run(assert_decode_cache_hit)
 
 
+class TestRustUnifiedHybridBitExact(TestUnifiedHybridBitExact):
+    tree_core_backend = "rust"
+
+
+class TestRustUnifiedHybridLazyBitExact(TestUnifiedHybridLazyBitExact):
+    tree_core_backend = "rust"
+
+
+class TestRustUnifiedHybridHiCacheBitExact(TestUnifiedHybridHiCacheBitExact):
+    tree_core_backend = "rust"
+
+
+class TestRustUnifiedHybridMTPBitExact(TestUnifiedHybridMTPBitExact):
+    tree_core_backend = "rust"
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/rust/test_run_mem_cache_rust_tests.py

@@ -0,0 +1,64 @@
+"""Run the standalone mem-cache crate's native Rust unit tests."""
+
+import shutil
+import subprocess
+import unittest
+from pathlib import Path
+
+from sglang.srt.environ import envs
+from sglang.srt.rust_extensions.torch_build import torch_build_configuration
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+BUILD_AND_RUN_TIMEOUT_S = 900
+RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust"
+MEM_CACHE_MANIFEST = RUST_WORKSPACE / "mem-cache" / "Cargo.toml"
+
+register_cpu_ci(est_time=900, suite="base-a-test-cpu")
+
+
+@unittest.skipIf(
+    envs.SGLANG_SKIP_RUST_TESTS.get(),
+    "SGLANG_SKIP_RUST_TESTS is set (no rust/ workspace changes per CI check-changes)",
+)
+class TestMemCacheCargo(CustomTestCase):
+    def test_mem_cache_native_tests(self):
+        self.assertIsNotNone(
+            shutil.which("cargo"),
+            "cargo not found on PATH; install a Rust toolchain "
+            "(scripts/ci/utils/install_rust_protoc.sh)",
+        )
+        self.assertTrue(
+            MEM_CACHE_MANIFEST.is_file(),
+            f"mem-cache manifest not found at {MEM_CACHE_MANIFEST}",
+        )
+        build = torch_build_configuration(
+            compat_header=MEM_CACHE_MANIFEST.parent / "torch_2_13_compat.h",
+            python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
+        )
+        proc = subprocess.run(
+            [
+                "cargo",
+                "test",
+                "--manifest-path",
+                str(MEM_CACHE_MANIFEST),
+                "--locked",
+                "--no-default-features",
+            ],
+            cwd=RUST_WORKSPACE,
+            env=build.environment,
+            capture_output=True,
+            text=True,
+            timeout=BUILD_AND_RUN_TIMEOUT_S,
+        )
+        print(proc.stdout)
+        self.assertEqual(
+            proc.returncode,
+            0,
+            f"mem-cache native tests failed\n"
+            f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/rust/test_run_rust_tests.py

@@ -1,4 +1,4 @@
-"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite."""
+"""Run the repository's native Rust unit tests from the CPU CI suite."""
 
 import shutil
 import subprocess
@@ -11,7 +11,6 @@ from sglang.test.test_utils import CustomTestCase
 
 BUILD_AND_RUN_TIMEOUT_S = 900
 RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust"
-
 register_cpu_ci(est_time=900, suite="base-a-test-cpu")
 
 
@@ -23,6 +22,24 @@ register_cpu_ci(est_time=900, suite="base-a-test-cpu")
     "SGLANG_SKIP_RUST_TESTS is set (no rust/ workspace changes per CI check-changes)",
 )
 class TestCargoWorkspace(CustomTestCase):
+    def _run_cargo(self, args: list[str], *, cwd: Path, env: dict | None = None):
+        proc = subprocess.run(
+            ["cargo", *args],
+            cwd=cwd,
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=BUILD_AND_RUN_TIMEOUT_S,
+        )
+        # Print unconditionally so a green run still shows which tests ran.
+        print(proc.stdout)
+        self.assertEqual(
+            proc.returncode,
+            0,
+            f"`cargo {' '.join(args)}` failed in {cwd}\n"
+            f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
+        )
+
     def test_cargo_test_workspace(self):
         # Not skipUnless: cargo is a hard dependency of the editable install
         # (setuptools-rust builds sglang-grpc), so a missing toolchain is a
@@ -37,21 +54,7 @@ class TestCargoWorkspace(CustomTestCase):
             f"rust workspace manifest not found at {RUST_WORKSPACE}",
         )
 
-        proc = subprocess.run(
-            ["cargo", "test", "--workspace"],
-            cwd=RUST_WORKSPACE,
-            capture_output=True,
-            text=True,
-            timeout=BUILD_AND_RUN_TIMEOUT_S,
-        )
-        # Print unconditionally so a green run still shows which tests ran.
-        print(proc.stdout)
-        self.assertEqual(
-            proc.returncode,
-            0,
-            f"`cargo test --workspace` failed in {RUST_WORKSPACE}\n"
-            f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
-        )
+        self._run_cargo(["test", "--workspace"], cwd=RUST_WORKSPACE)
 
 
 if __name__ == "__main__":

test/registered/rust/test_rust_extension.py

@@ -9,11 +9,12 @@ import time
 import unittest
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from types import ModuleType
+from types import ModuleType, SimpleNamespace
 from unittest import mock
 
 from sglang.srt.rust_extensions import load_rust_extension
 from sglang.srt.rust_extensions import loader as rust_extension
+from sglang.srt.rust_extensions.torch_build import torch_build_configuration
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -84,6 +85,74 @@ crate-type = ["cdylib"]
         fingerprint.assert_not_called()
         cargo_build.assert_not_called()
 
+    def test_bundled_named_variant_never_touches_source_or_cargo(self):
+        bundled = ModuleType("demo._inspection")
+        with (
+            mock.patch.object(
+                rust_extension.importlib, "import_module", return_value=bundled
+            ) as import_module,
+            mock.patch.object(rust_extension, "_discover_crate") as discover,
+            mock.patch.object(rust_extension, "_build_context") as fingerprint,
+            mock.patch.object(rust_extension, "_cargo_build") as cargo_build,
+        ):
+            self.assertIs(
+                load_rust_extension(
+                    "demo._core",
+                    mode="never",
+                    workspace=Path("/workspace/not-present"),
+                    additional_features=("inspection",),
+                    extension_module="demo._inspection",
+                ),
+                bundled,
+            )
+        import_module.assert_called_once_with("demo._inspection")
+        discover.assert_not_called()
+        fingerprint.assert_not_called()
+        cargo_build.assert_not_called()
+
+    def test_auto_ignores_a_stale_bundled_extension_in_a_source_tree(self):
+        with TemporaryDirectory() as directory:
+            root = Path(directory)
+            workspace = self._workspace(root)
+            (workspace / "demo/lib.rs").write_text(
+                "fn source_changed() {}\n", encoding="utf-8"
+            )
+            stale = ModuleType("demo._core")
+            built = ModuleType("demo._core")
+            artifact = root / "libdemo_extension.so"
+            artifact.write_bytes(b"fresh extension")
+            context = rust_extension._BuildContext(
+                "changed-source", "fingerprint", "target"
+            )
+            with (
+                mock.patch.object(
+                    rust_extension, "_import_bundled_extension", return_value=stale
+                ) as bundled_import,
+                mock.patch.object(
+                    rust_extension, "_build_context", return_value=context
+                ),
+                mock.patch.object(
+                    rust_extension, "_source_digest", return_value="changed-source"
+                ),
+                mock.patch.object(
+                    rust_extension, "_cargo_build", return_value=artifact
+                ) as cargo_build,
+                mock.patch.object(
+                    rust_extension, "_load_extension_from_path", return_value=built
+                ),
+            ):
+                self.assertIs(
+                    load_rust_extension(
+                        "demo._core",
+                        mode="auto",
+                        workspace=workspace,
+                        cache_dir=root / "cache",
+                    ),
+                    built,
+                )
+            bundled_import.assert_not_called()
+            cargo_build.assert_called_once()
+
     def test_discovery_reads_crate_manifest_metadata(self):
         with TemporaryDirectory() as directory:
             workspace = self._workspace(Path(directory))
@@ -126,6 +195,18 @@ crate-type = ["cdylib"]
                     changed_flags.target_fingerprint,
                 )
 
+                inspection = rust_extension._build_context(
+                    crate,
+                    features=(*crate.features, "inspection"),
+                    extension_module="demo._inspection",
+                    build_fingerprint={"torch": "2.13"},
+                )
+                self.assertNotEqual(changed_source.fingerprint, inspection.fingerprint)
+                self.assertNotEqual(
+                    changed_source.target_fingerprint,
+                    inspection.target_fingerprint,
+                )
+
     def test_auto_builds_once_then_uses_cache(self):
         with TemporaryDirectory() as directory:
             root = Path(directory)
@@ -155,13 +236,19 @@ crate-type = ["cdylib"]
             ):
                 self.assertIs(
                     rust_extension.load_rust_extension(
-                        "demo._core", workspace=workspace, cache_dir=root / "cache"
+                        "demo._core",
+                        mode="auto",
+                        workspace=workspace,
+                        cache_dir=root / "cache",
                     ),
                     loaded,
                 )
                 self.assertIs(
                     rust_extension.load_rust_extension(
-                        "demo._core", workspace=workspace, cache_dir=root / "cache"
+                        "demo._core",
+                        mode="auto",
+                        workspace=workspace,
+                        cache_dir=root / "cache",
                     ),
                     loaded,
                 )
@@ -270,6 +357,122 @@ crate-type = ["cdylib"]
                 ],
             )
 
+    def test_variant_uses_its_own_module_name_features_and_environment(self):
+        with TemporaryDirectory() as directory:
+            root = Path(directory)
+            workspace = self._workspace(root)
+            artifact = root / "libdemo_extension.so"
+            artifact.write_bytes(b"extension")
+            context = rust_extension._BuildContext("source", "fingerprint", "target")
+            loaded = ModuleType("demo._inspection")
+            environment = {"CUSTOM_BUILD_INPUT": "value"}
+            with (
+                mock.patch.object(
+                    rust_extension, "_import_bundled_extension", return_value=None
+                ) as bundled_import,
+                mock.patch.object(
+                    rust_extension, "_build_context", return_value=context
+                ) as build_context,
+                mock.patch.object(
+                    rust_extension, "_source_digest", return_value="source"
+                ),
+                mock.patch.object(
+                    rust_extension, "_cargo_build", return_value=artifact
+                ) as cargo_build,
+                mock.patch.object(
+                    rust_extension,
+                    "_load_extension_from_path",
+                    return_value=loaded,
+                ) as load_from_path,
+            ):
+                self.assertIs(
+                    load_rust_extension(
+                        "demo._core",
+                        mode="auto",
+                        workspace=workspace,
+                        cache_dir=root / "cache",
+                        additional_features=("inspection",),
+                        extension_module="demo._inspection",
+                        build_environment=environment,
+                        build_fingerprint={"native": "abi"},
+                    ),
+                    loaded,
+                )
+            bundled_import.assert_not_called()
+            self.assertEqual(
+                build_context.call_args.kwargs,
+                {
+                    "features": ("python", "inspection"),
+                    "build_fingerprint": {"native": "abi"},
+                    "extension_module": "demo._inspection",
+                },
+            )
+            self.assertEqual(
+                cargo_build.call_args.kwargs,
+                {
+                    "features": ("python", "inspection"),
+                    "build_environment": environment,
+                },
+            )
+            self.assertEqual(load_from_path.call_args.args[0], "demo._inspection")
+
+    def test_torch_build_configuration_is_versioned_and_relocatable(self):
+        with TemporaryDirectory() as directory:
+            root = Path(directory)
+            torch_root = root / "torch"
+            (torch_root / "lib").mkdir(parents=True)
+            torch_init = torch_root / "__init__.py"
+            torch_init.write_text("", encoding="utf-8")
+            compat_header = root / "compat.h"
+            compat_header.write_text("// compatibility\n", encoding="utf-8")
+            fake_torch = SimpleNamespace(
+                __version__="2.13.0+cu130",
+                __file__=str(torch_init),
+                compiled_with_cxx11_abi=lambda: True,
+                version=SimpleNamespace(cuda="13.0", hip=None),
+            )
+
+            build = torch_build_configuration(
+                compat_header=compat_header,
+                python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
+                torch_module=fake_torch,
+                base_environment={
+                    "PATH": "/usr/bin",
+                    "CXXFLAGS": "-O2",
+                    "RUSTFLAGS": "-Ctarget-cpu=x86-64",
+                },
+            )
+
+            self.assertEqual(build.environment["LIBTORCH_USE_PYTORCH"], "1")
+            self.assertEqual(build.environment["LIBTORCH_BYPASS_VERSION_CHECK"], "1")
+            self.assertIn(str(compat_header), build.environment["CXXFLAGS"])
+            self.assertIn(
+                "$ORIGIN/../../../../torch/lib", build.environment["RUSTFLAGS"]
+            )
+            self.assertIn(str(torch_root / "lib"), build.environment["RUSTFLAGS"])
+            self.assertEqual(build.fingerprint["torch_version"], "2.13.0+cu130")
+            self.assertTrue(build.fingerprint["torch_cxx11_abi"])
+
+            wheel_build = torch_build_configuration(
+                compat_header=compat_header,
+                python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
+                torch_module=fake_torch,
+                base_environment={},
+                include_absolute_rpath=False,
+            )
+            self.assertNotIn(
+                str(torch_root / "lib"), wheel_build.environment["RUSTFLAGS"]
+            )
+            self.assertFalse(wheel_build.fingerprint["include_absolute_rpath"])
+
+            fake_torch.__version__ = "2.14.0"
+            with self.assertRaisesRegex(RuntimeError, "PyTorch 2.11 through 2.13"):
+                torch_build_configuration(
+                    compat_header=compat_header,
+                    python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
+                    torch_module=fake_torch,
+                )
+
     def test_filesystem_lock_serializes_processes(self):
         with TemporaryDirectory() as directory:
             lock_path = Path(directory) / "build.lock"
@@ -337,6 +540,12 @@ crate-type = ["cdylib"]
                 "sglang_mm_core",
                 ("python", "parallel"),
             ),
+            (
+                "sglang.srt.mem_cache.rust_tree_core.mem_cache",
+                "mem_cache",
+                "mem_cache",
+                ("python-extension",),
+            ),
         ):
             crate = rust_extension._discover_crate(
                 rust_extension._RUST_WORKSPACE, python_module

test/registered/unit/disaggregation/test_decode_hicache_tree_core.py

@@ -0,0 +1,107 @@
+"""Unit tests for decode HiCache TreeCore interactions."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import torch
+
+from sglang.srt.disaggregation.decode_hicache_mixin import (
+    DecodeHiCachePreallocMixin,
+    DecodePrefixMatch,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestDecodeHiCacheTreeCore(CustomTestCase):
+    def test_storage_probe_and_prefetch_use_node_handles(self):
+        ongoing_prefetch = {}
+
+        def register_prefetch(req_id, *_args, **_kwargs):
+            ongoing_prefetch[req_id] = object()
+
+        tree_cache = SimpleNamespace(
+            hicache_storage_pass_prefix_keys=True,
+            ongoing_prefetch=ongoing_prefetch,
+            is_backuped=Mock(return_value=True),
+            is_root=Mock(return_value=False),
+            get_last_hash_value=Mock(return_value="h2"),
+            get_prefix_hash_values=Mock(return_value=["h0", "h1"]),
+            query_storage_hit_length=Mock(return_value=2),
+            prefetch_from_storage=Mock(side_effect=register_prefetch),
+        )
+        harness = SimpleNamespace(
+            scheduler=SimpleNamespace(enable_decode_hicache=True),
+            tree_cache=tree_cache,
+        )
+        req = SimpleNamespace(
+            rid="req-0",
+            origin_input_ids=[0, 1, 2, 3, 4, 5, 6, 7],
+            extra_key="model",
+            cache_salt=None,
+        )
+        result = SimpleNamespace(
+            device_indices=torch.tensor([10, 11]),
+            host_hit_length=2,
+            last_device_node=11,
+            last_host_node=22,
+        )
+
+        prefix_match = DecodeHiCachePreallocMixin._build_decode_prefix_match(
+            harness, req, result
+        )
+
+        self.assertEqual(prefix_match.l3_storage_hit_length, 2)
+        tree_cache.query_storage_hit_length.assert_called_once_with(
+            22, [4, 5, 6, 7], "h2", ["h0", "h1"]
+        )
+
+        DecodeHiCachePreallocMixin._start_hicache_prefetch(harness, req, prefix_match)
+
+        self.assertTrue(prefix_match.prefetch_registered)
+        tree_cache.prefetch_from_storage.assert_called_once_with(
+            "req-0",
+            22,
+            [4, 5],
+            "h2",
+            ["h0", "h1"],
+            extra_key="model",
+            cache_salt=None,
+        )
+
+    def test_stale_prefetch_anchor_degrades_to_l2(self):
+        tree_cache = SimpleNamespace(
+            hicache_storage_pass_prefix_keys=True,
+            ongoing_prefetch={},
+            get_last_hash_value=Mock(side_effect=KeyError(22)),
+            get_prefix_hash_values=Mock(),
+            prefetch_from_storage=Mock(),
+        )
+        harness = SimpleNamespace(tree_cache=tree_cache)
+        req = SimpleNamespace(
+            rid="req-0",
+            origin_input_ids=[0, 1, 2, 3, 4, 5],
+            extra_key=None,
+            cache_salt=None,
+        )
+        prefix_match = DecodePrefixMatch(
+            prefix_indices=torch.tensor([10, 11]),
+            l2_host_hit_length=2,
+            l3_storage_hit_length=2,
+            last_device_node=11,
+            last_host_node=22,
+        )
+
+        DecodeHiCachePreallocMixin._start_hicache_prefetch(harness, req, prefix_match)
+
+        self.assertEqual(prefix_match.l3_storage_hit_length, 0)
+        self.assertFalse(prefix_match.prefetch_registered)
+        tree_cache.get_prefix_hash_values.assert_not_called()
+        tree_cache.prefetch_from_storage.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/managers/test_schedule_policy_dfs_weight.py

@@ -0,0 +1,55 @@
+"""Unit tests for DFS-weight schedule-policy delegation."""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.managers.schedule_policy import SchedulePolicy
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestSchedulePolicyDfsWeight(CustomTestCase):
+    def test_orders_requests_by_subtree_weight(self):
+        class Node:
+            def __init__(self):
+                self.children = {}
+
+        root = Node()
+        branch_a = Node()
+        branch_b = Node()
+        leaf_a1 = Node()
+        leaf_a2 = Node()
+        root.children = {"a": branch_a, "b": branch_b}
+        branch_a.children = {"a1": leaf_a1, "a2": leaf_a2}
+
+        class TreeCache:
+            dfs_weight_order = BasePrefixCache.dfs_weight_order
+
+            def __init__(self):
+                self.root_node = root
+
+            @staticmethod
+            def resolve_node_handle(node):
+                return node
+
+        waiting_queue = [
+            SimpleNamespace(last_node=branch_b, name="b"),
+            SimpleNamespace(last_node=leaf_a2, name="a2"),
+            SimpleNamespace(last_node=leaf_a1, name="a1-first"),
+            SimpleNamespace(last_node=leaf_a1, name="a1-second"),
+            SimpleNamespace(last_node=branch_a, name="a-parent"),
+        ]
+
+        SchedulePolicy._sort_by_dfs_weight(waiting_queue, TreeCache())
+
+        self.assertEqual(
+            [req.name for req in waiting_queue],
+            ["a1-first", "a1-second", "a2", "a-parent", "b"],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/rust_unified_tree_core_inspector.py

@@ -0,0 +1,228 @@
+"""Test-only inspection adapter for the Rust Unified TreeCore."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from unified_tree_core_inspection_interface import (
+    UnifiedTreeCoreInspectionInterface,
+)
+
+from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams, MatchResult
+from sglang.srt.mem_cache.rust_tree_core.adapter import (
+    RustUnifiedTreeCore,
+    _fill_evict_result,
+    _match_result_from_binding,
+    _radix_key_buffer,
+)
+from sglang.srt.mem_cache.rust_tree_core.extension import load_tree_core_extension
+from sglang.srt.mem_cache.unified_cache.components import ComponentType, EvictLayer
+from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
+    BaseEvictionResult,
+    NodeId,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(
+    est_time=0, suite="base-a-test-cpu", disabled="Rust TreeCore test inspector"
+)
+
+_inspection_bindings = load_tree_core_extension(inspection=True)
+
+
+class RustUnifiedTreeCoreInspector(
+    RustUnifiedTreeCore, UnifiedTreeCoreInspectionInterface
+):
+    """Rust TreeCore variant used by the shared backend-conformance tests.
+
+    The production adapter deliberately implements only
+    ``UnifiedTreeCoreInterface``. These forwarding methods keep white-box state
+    controls in test code while the binding returns snapshots rather than Rust
+    iterators across the Python boundary.
+    """
+
+    _bindings = _inspection_bindings
+
+    def contains_node(self, node_id: NodeId) -> bool:
+        return self._binding.inspect_contains_node(node_id)
+
+    def get_parent_node_id(self, node_id: NodeId) -> Optional[NodeId]:
+        return self._binding.inspect_get_parent_node_id(node_id)
+
+    def get_child_node_ids(self, node_id: NodeId) -> list[NodeId]:
+        return self._binding.inspect_get_child_node_ids(node_id)
+
+    def get_node_key_length(self, node_id: NodeId) -> int:
+        return self._binding.inspect_get_node_key_length(node_id)
+
+    def get_node_token_ids(self, node_id: NodeId) -> list[int]:
+        return self._binding.inspect_get_node_token_ids(node_id)
+
+    def is_node_key_bigram(self, node_id: NodeId) -> bool:
+        return self._binding.inspect_is_node_key_bigram(node_id)
+
+    def get_component_host_value(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> Optional[torch.Tensor]:
+        return self._binding.inspect_get_component_host_value(
+            node_id, int(component_type)
+        )
+
+    def get_component_device_lock_ref(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> int:
+        return self._binding.inspect_get_component_device_lock_ref(
+            node_id, int(component_type)
+        )
+
+    def get_node_hit_count(self, node_id: NodeId) -> int:
+        return self._binding.inspect_get_node_hit_count(node_id)
+
+    def get_write_through_pending_id(self, node_id: NodeId) -> Optional[int]:
+        return self._binding.inspect_get_write_through_pending_id(node_id)
+
+    def is_node_in_device_lru(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> bool:
+        return self._binding.inspect_is_node_in_device_lru(node_id, int(component_type))
+
+    def is_node_in_host_lru(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> bool:
+        return self._binding.inspect_is_node_in_host_lru(node_id, int(component_type))
+
+    def get_component_device_lru_node_ids(
+        self, component_type: ComponentType
+    ) -> list[NodeId]:
+        return self._binding.inspect_get_component_device_lru_node_ids(
+            int(component_type)
+        )
+
+    def is_device_evictable_leaf(self, node_id: NodeId) -> bool:
+        return self._binding.inspect_is_device_evictable_leaf(node_id)
+
+    def is_host_evictable_leaf(self, node_id: NodeId) -> bool:
+        return self._binding.inspect_is_host_evictable_leaf(node_id)
+
+    def is_device_leaf(self, node_id: NodeId) -> bool:
+        return self._binding.inspect_is_device_leaf(node_id)
+
+    def get_all_node_ids(self) -> list[NodeId]:
+        return self._binding.inspect_get_all_node_ids()
+
+    def component_protected_size(self, component_type: ComponentType) -> int:
+        return self._binding.inspect_component_protected_size(int(component_type))
+
+    def set_node_hash_values(
+        self, node_id: NodeId, hash_values: Optional[list[str]]
+    ) -> None:
+        self._binding.inspect_set_node_hash_values(node_id, hash_values)
+
+    def set_component_device_value_raw(
+        self,
+        node_id: NodeId,
+        component_type: ComponentType,
+        value: Optional[torch.Tensor],
+    ) -> None:
+        self._binding.inspect_set_component_device_value_raw(
+            node_id, int(component_type), value
+        )
+
+    def set_component_host_value_raw(
+        self,
+        node_id: NodeId,
+        component_type: ComponentType,
+        value: Optional[torch.Tensor],
+    ) -> None:
+        self._binding.inspect_set_component_host_value_raw(
+            node_id, int(component_type), value
+        )
+
+    def set_component_device_lock_ref(
+        self, node_id: NodeId, component_type: ComponentType, lock_ref: int
+    ) -> None:
+        assert lock_ref >= 0
+        self._binding.inspect_set_component_device_lock_ref(
+            node_id, int(component_type), lock_ref
+        )
+
+    def remove_node_from_device_lru(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> None:
+        self._binding.inspect_remove_node_from_device_lru(node_id, int(component_type))
+
+    def insert_node_into_host_lru(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> None:
+        self._binding.inspect_insert_node_into_host_lru(node_id, int(component_type))
+
+    def set_component_evictable_size(
+        self, component_type: ComponentType, value: int
+    ) -> None:
+        assert value >= 0
+        self._binding.inspect_set_component_evictable_size(int(component_type), value)
+
+    def set_component_protected_size(
+        self, component_type: ComponentType, value: int
+    ) -> None:
+        assert value >= 0
+        self._binding.inspect_set_component_protected_size(int(component_type), value)
+
+    def update_duplicate_tracking(self, node_id: NodeId) -> None:
+        self._binding.inspect_update_duplicate_tracking(node_id)
+
+    def advance_insert_walk_once(self) -> None:
+        self._binding.inspect_advance_insert_walk_once()
+
+    def evict_component(
+        self,
+        node_id: NodeId,
+        component_type: ComponentType,
+        target: EvictLayer,
+    ) -> BaseEvictionResult:
+        binding_result = self._binding.inspect_evict_component(
+            node_id, int(component_type), int(target)
+        )
+        return _fill_evict_result(binding_result, BaseEvictionResult())
+
+    def validate_cascade_evict(
+        self,
+        node_id: NodeId,
+        component_type: ComponentType,
+        target: EvictLayer,
+    ) -> None:
+        self._binding.inspect_validate_cascade_evict(
+            node_id, int(component_type), int(target)
+        )
+
+    def cleanup_tombstone_ancestors(self, node_id: NodeId) -> BaseEvictionResult:
+        binding_result = self._binding.inspect_cleanup_tombstone_ancestors(node_id)
+        return _fill_evict_result(binding_result, BaseEvictionResult())
+
+    def finalize_component_match_result(
+        self,
+        component_type: ComponentType,
+        result: MatchResult,
+        params: MatchPrefixParams,
+        value_chunks: list[torch.Tensor],
+        best_value_len: int,
+    ) -> MatchResult:
+        binding_result = self._binding.inspect_finalize_component_match_result(
+            int(component_type),
+            result,
+            _radix_key_buffer(params.key),
+            params.key.extra_key,
+            params.key.cache_salt,
+            value_chunks,
+            best_value_len,
+        )
+        return _match_result_from_binding(binding_result)._replace(
+            cache_protected_len=result.cache_protected_len,
+            cache_actions=result.cache_actions,
+        )
+
+    def build_backup_node_ids(
+        self, node_id: NodeId, write_back: bool = False
+    ) -> list[NodeId]:
+        return self._binding.inspect_build_backup_node_ids(node_id, write_back)

test/registered/unit/mem_cache/test_rust_tree_core.py

@@ -0,0 +1,149 @@
+"""Smoke tests for the in-tree Rust TreeCore backend (``rust``).
+
+Requires a Rust toolchain: the extension builds with cargo on first use.
+"""
+
+import shutil
+from array import array
+
+import pytest
+import torch
+from unified_tree_core_inspection_interface import UnifiedTreeCoreInspectionInterface
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=90, suite="base-a-test-cpu")
+
+if shutil.which("cargo") is None:
+    pytest.skip("the rust backend builds with cargo", allow_module_level=True)
+
+from sglang.srt.mem_cache.base_prefix_cache import InsertParams, MatchPrefixParams
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.mem_cache.radix_cache import RadixKey
+from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
+from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core
+
+
+def _tree_core():
+    return create_tree_core(
+        "rust",
+        CacheInitParams(
+            disable=False,
+            req_to_token_pool=None,
+            token_to_kv_pool_allocator=None,
+            page_size=1,
+            tree_components=(ComponentType.FULL,),
+        ),
+        components={},
+    )
+
+
+def _key(token_ids, extra_key=None):
+    return RadixKey(array("q", token_ids), extra_key=extra_key)
+
+
+def _pump_insert(core, params):
+    step = core.begin_insert(params)
+    while step.result is None:
+        step = core.resume_insert()
+    core.end_insert()
+    return step.result
+
+
+def test_registry_resolves_the_rust_backend_lazily():
+    core = _tree_core()
+    assert type(core).__name__ == "RustUnifiedTreeCore"
+    assert not isinstance(core, UnifiedTreeCoreInspectionInterface)
+    assert not any(name.startswith("inspect_") for name in dir(core._binding))
+
+
+def test_insert_then_match_round_trips():
+    core = _tree_core()
+    _pump_insert(
+        core,
+        InsertParams(
+            key=_key([1, 2, 3]), value=torch.tensor([10, 11, 12], dtype=torch.int64)
+        ),
+    )
+    matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3])))
+    assert matched.device_indices.tolist() == [10, 11, 12]
+
+
+def test_lock_moves_tokens_between_evictable_and_protected():
+    core = _tree_core()
+    _pump_insert(
+        core,
+        InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)),
+    )
+    matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2])))
+    core.inc_lock_ref(matched.best_match_node)
+    assert core.protected_size() == 2
+    assert core.evictable_size() == 0
+    core.dec_lock_ref(matched.best_match_node)
+    assert core.evictable_size() == 2
+
+
+def test_namespaces_isolate_the_same_tokens():
+    core = _tree_core()
+    _pump_insert(
+        core,
+        InsertParams(
+            key=_key([1, 2], extra_key="chat"),
+            value=torch.tensor([20, 21], dtype=torch.int64),
+        ),
+    )
+    salted = core.match_prefix(MatchPrefixParams(key=_key([1, 2], extra_key="chat")))
+    assert salted.device_indices.tolist() == [20, 21]
+    unsalted = core.match_prefix(MatchPrefixParams(key=_key([1, 2])))
+    assert unsalted.device_indices.numel() == 0
+
+
+def test_backfill_hashes_existing_nodes_in_parent_order():
+    expected = _tree_core()
+    expected.enable_storage = True
+    _pump_insert(
+        expected,
+        InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)),
+    )
+    _pump_insert(
+        expected,
+        InsertParams(
+            key=_key([1, 2, 3, 4]),
+            value=torch.tensor([10, 11, 12, 13], dtype=torch.int64),
+        ),
+    )
+
+    late = _tree_core()
+    _pump_insert(
+        late,
+        InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)),
+    )
+    _pump_insert(
+        late,
+        InsertParams(
+            key=_key([1, 2, 3, 4]),
+            value=torch.tensor([10, 11, 12, 13], dtype=torch.int64),
+        ),
+    )
+
+    parent = late.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node
+    child = late.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4]))).best_match_node
+    expected_parent = expected.match_prefix(
+        MatchPrefixParams(key=_key([1, 2]))
+    ).best_match_node
+    expected_child = expected.match_prefix(
+        MatchPrefixParams(key=_key([1, 2, 3, 4]))
+    ).best_match_node
+
+    assert late.get_hash_values(parent) == []
+    assert late.get_hash_values(child) == []
+    assert late.backfill_missing_hash_values() == 2
+    assert late.get_hash_values(parent) == expected.get_hash_values(expected_parent)
+    assert late.get_hash_values(child) == expected.get_hash_values(expected_child)
+    assert late.backfill_missing_hash_values() == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/mem_cache/test_rust_unified_radix_cache_bench.py

@@ -0,0 +1,29 @@
+"""Run the UnifiedRadixCache benchmark/fuzz suite with the Rust TreeCore."""
+
+import unittest
+
+import test_unified_radix_cache_bench as shared_suite
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-small")
+
+
+class RustBackendSuite(unittest.TestSuite):
+    """Scope the backend override to this suite and restore it afterward."""
+
+    def run(self, result, debug=False):
+        previous = shared_suite._TREE_CORE_TEST_BACKEND
+        shared_suite._TREE_CORE_TEST_BACKEND = "rust"
+        try:
+            return super().run(result, debug)
+        finally:
+            shared_suite._TREE_CORE_TEST_BACKEND = previous
+
+
+def load_tests(loader, standard_tests, pattern):
+    return RustBackendSuite(loader.loadTestsFromModule(shared_suite))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_rust_unified_radix_cache_unittest.py

@@ -0,0 +1,30 @@
+"""Run the shared UnifiedRadixCache unit suite with the Rust TreeCore."""
+
+import unittest
+
+import test_unified_radix_cache_unittest as shared_suite
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=180, stage="base-b", runner_config="1-gpu-small")
+
+
+class RustBackendSuite(unittest.TestSuite):
+    """Scope the test backend to this suite without polluting discovery."""
+
+    def run(self, result, debug=False):
+        previous = shared_suite._TREE_CORE_TEST_BACKEND
+        shared_suite._TREE_CORE_TEST_BACKEND = "rust"
+        try:
+            return super().run(result, debug)
+        finally:
+            shared_suite._TREE_CORE_TEST_BACKEND = previous
+
+
+def load_tests(loader, standard_tests, pattern):
+    """Reuse the exact cache-level suite while swapping only its test factory."""
+    return RustBackendSuite(loader.loadTestsFromModule(shared_suite))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_unified_radix_cache_bench.py

@@ -14,7 +14,7 @@ import sys
 import time
 import unittest
 from array import array
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Callable
 
@@ -40,6 +40,7 @@ from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=25, suite="stage-b-test-1-gpu-small-amd")
@@ -59,6 +60,7 @@ _BENCH_KV_SIZE = 500_000
 _BENCH_CHUNK_LEN = 256
 
 _DEFAULT_COMPONENTS = (ComponentType.FULL, ComponentType.MAMBA)
+_TREE_CORE_TEST_BACKEND: str | None = None
 
 
 @contextmanager
@@ -226,16 +228,22 @@ def create_bench_cache(
     # --- tree ---
     if tree_cls is None:
         tree_cls = UnifiedRadixCache
-    tree = tree_cls(
-        params=CacheInitParams(
-            req_to_token_pool=req_to_token_pool,
-            token_to_kv_pool_allocator=allocator,
-            page_size=page_size,
-            disable=False,
-            tree_components=components if tree_cls is UnifiedRadixCache else None,
-            sliding_window_size=sliding_window_size if has_swa else None,
-        )
+    backend_override = (
+        envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override(_TREE_CORE_TEST_BACKEND)
+        if _TREE_CORE_TEST_BACKEND is not None and tree_cls is UnifiedRadixCache
+        else nullcontext()
     )
+    with backend_override:
+        tree = tree_cls(
+            params=CacheInitParams(
+                req_to_token_pool=req_to_token_pool,
+                token_to_kv_pool_allocator=allocator,
+                page_size=page_size,
+                disable=False,
+                tree_components=components if tree_cls is UnifiedRadixCache else None,
+                sliding_window_size=sliding_window_size if has_swa else None,
+            )
+        )
 
     _rid = [0]
 
@@ -780,6 +788,10 @@ class _BenchSuite:
             verify=True,
             page_size=cfg["page_size"],
         )
+        backend = (
+            _TREE_CORE_TEST_BACKEND or envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get()
+        )
+        print(f"[{backend}] {r.report()}")
         self.assertGreater(r.num_ops, 0)
         self.assertGreater(r.ops_per_sec, 0)
 
@@ -803,7 +815,7 @@ for _cfg in _CI_BENCH_CONFIGS:
     _name = f"TestBench_{_cfg['label']}"
     globals()[_name] = type(
         _name,
-        (_BenchSuite, unittest.TestCase),
+        (_BenchSuite, CustomTestCase),
         {"bench_cfg": _cfg},
     )
     globals()[_name].__module__ = __name__

test/registered/unit/mem_cache/unified_tree_core_inspection_interface.py

@@ -205,6 +205,11 @@ class UnifiedTreeCoreInspectionInterface(UnifiedTreeCoreInterface):
 
     # ==== Targeted white-box operations ====
 
+    @abstractmethod
+    def advance_insert_walk_once(self) -> None:
+        """Advance one suspended insert walk step without flushing its actions."""
+        ...
+
     @abstractmethod
     def evict_component(
         self,

test/registered/unit/mem_cache/unified_tree_core_inspector.py

@@ -14,6 +14,7 @@ from sglang.srt.mem_cache.unified_cache.components import ComponentType, EvictLa
 from sglang.srt.mem_cache.unified_cache.unified_tree_core import (
     UnifiedLRUList,
     UnifiedTreeCore,
+    _InsertPhase,
 )
 from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
     BaseEvictionResult,
@@ -195,6 +196,15 @@ class UnifiedTreeCoreInspector(UnifiedTreeCore, UnifiedTreeCoreInspectionInterfa
         """Refresh duplicate-host tracking for the node."""
         self._update_duplicate_tracking(self.node_by_id(node_id))
 
+    def advance_insert_walk_once(self) -> None:
+        """Advance one suspended insert walk step without flushing its actions."""
+        state = self._ongoing_insert_walk_state
+        if state is None:
+            raise RuntimeError("no in-flight insert")
+        if state.phase is not _InsertPhase.WALK:
+            raise RuntimeError("in-flight insert is not in walk phase")
+        self._insert_walk_step(state)
+
     def evict_component(
         self,
         node_id: NodeId,

test/registered/unit/server_args/test_server_args.py

@@ -61,7 +61,9 @@ from sglang.srt.arg_groups.serving_hook import (
     ssl_verify_of,
 )
 from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
-from sglang.srt.arg_groups.validation_hook import check_two_batch_overlap
+from sglang.srt.arg_groups.validation_hook import (
+    check_two_batch_overlap,
+)
 from sglang.srt.entrypoints.sidecar import (
     SGLANG_GRPC_ENDPOINT_ENV,
     Sidecar,
@@ -789,6 +791,19 @@ class TestLoadBalanceMethod(unittest.TestCase):
             "mooncake",
         )
 
+    def test_pd_decode_hicache_allows_rust_tree_core(self):
+        server_args = ServerArgs(
+            model_path="dummy",
+            disaggregation_mode="decode",
+            disaggregation_decode_enable_radix_cache=True,
+            disaggregation_transfer_backend="nixl",
+            enable_hierarchical_cache=True,
+        )
+        with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("rust"):
+            handle_pd_disaggregation(server_args)
+
+        self.assertFalse(resolution_result(server_args, "disable_radix_cache"))
+
 
 class TestSkipTokenizerInit(unittest.TestCase):
     def test_skip_tokenizer_worker_counts(self):
@@ -1475,6 +1490,16 @@ class TestHiCacheArgs(unittest.TestCase):
                 expected_decode_backend,
             )
 
+    def test_buffer_only_accepts_both_tree_cores(self):
+        for backend in ("python", "rust"):
+            args = self._make_args(
+                enable_hierarchical_cache=True,
+                hicache_host_memory_mode="buffer_only",
+                hicache_storage_backend="file",
+            )
+            with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override(backend):
+                handle_hicache(args)
+
     def test_hicache_io_backend_and_mem_layout_compatibility(self):
         cases = [
             {