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