[Rust] Rename mem-cache to sglang-radix-tree (#37290)

This commit is contained in:
Jialin Ouyang
2026-09-01 03:03:44 -07:00
committed by GitHub
parent c16a8fc899
commit a77283fb02
30 changed files with 38 additions and 38 deletions
@@ -96,7 +96,7 @@ jobs:
run: bash scripts/ci/utils/ensure_zstd.sh
# The setup hook and torch helper select and configure what gets built;
# pyproject.toml pins the libtorch ABI used by mem-cache.
# pyproject.toml pins the libtorch ABI used by sglang-radix-tree.
- name: Restore built modules
id: cache
uses: actions/cache/restore@v4
@@ -218,8 +218,8 @@ jobs:
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
cargo clean --release --manifest-path rust/sglang-radix-tree/Cargo.toml \
-p torch-sys -p sglang-radix-tree 2>/dev/null || true
(cd python && SGLANG_BUILD_RUST_EXTS=all python setup.py build_rust --inplace)
python - <<'PY'
import importlib.util
@@ -245,7 +245,7 @@ jobs:
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",
compat_header=root / "rust/sglang-radix-tree/torch_2_13_compat.h",
python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
torch_module=torch,
include_absolute_rpath=False,
@@ -257,7 +257,7 @@ jobs:
"--release",
"--locked",
"--manifest-path",
"rust/mem-cache/Cargo.toml",
"rust/sglang-radix-tree/Cargo.toml",
"--features",
"python-extension,inspection",
],
+1 -1
View File
@@ -231,7 +231,7 @@ work_dirs/
# Rust lib
Cargo.lock
!rust/Cargo.lock
!rust/mem-cache/Cargo.lock
!rust/sglang-radix-tree/Cargo.lock
# Generated vision test fixtures (regenerate with: python scripts/generate_vision_golden.py)
sgl-model-gateway/tests/fixtures/golden/
+2 -2
View File
@@ -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 && cargo clippy --manifest-path mem-cache/Cargo.toml --all-targets --no-default-features --features tch/doc-only -- -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 sglang-radix-tree/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 && cargo fmt --manifest-path mem-cache/Cargo.toml'
entry: bash -c 'cd rust && cargo fmt && cargo fmt --manifest-path sglang-radix-tree/Cargo.toml'
language: system
files: ^rust/.*\.rs$
pass_filenames: false
+2 -2
View File
@@ -148,8 +148,8 @@ def _discovered_rust_extensions():
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.
# built. Alternate platform pyprojects filter the Rust TreeCore extension
# 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)
@@ -10,7 +10,7 @@ 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"
_CRATE_DIR = Path(__file__).resolve().parents[5] / "rust" / "sglang-radix-tree"
_TORCH_COMPAT_HEADER = _CRATE_DIR / "torch_2_13_compat.h"
+3 -3
View File
@@ -5,12 +5,12 @@ members = [
"sglang-mm",
"sglang-server"
]
exclude = ["mem-cache"]
exclude = ["sglang-radix-tree"]
[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"]
# graph (sglang-radix-tree currently uses the tch-compatible PyO3 0.22 API).
extension-manifests = ["sglang-radix-tree/Cargo.toml"]
[workspace.package]
version = "0.1.0"
@@ -322,17 +322,6 @@ dependencies = [
"rawpointer",
]
[[package]]
name = "mem_cache"
version = "0.1.0"
dependencies = [
"hashbrown",
"pyo3",
"sha2",
"tch",
"thiserror",
]
[[package]]
name = "memchr"
version = "2.8.3"
@@ -653,6 +642,17 @@ dependencies = [
"zmij",
]
[[package]]
name = "sglang-radix-tree"
version = "0.1.0"
dependencies = [
"hashbrown",
"pyo3",
"sha2",
"tch",
"thiserror",
]
[[package]]
name = "sha1"
version = "0.10.7"
@@ -1,5 +1,5 @@
[package]
name = "mem_cache"
name = "sglang-radix-tree"
version = "0.1.0"
edition = "2024"
description = "Rust tree core for the Unified Radix Cache"
@@ -1,4 +1,4 @@
# mem-cache
# sglang-radix-tree
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.
@@ -20,7 +20,7 @@ and `torch_2_13_compat.h` covers two alignment APIs removed in PyTorch 2.13.
```bash
# Build (libtorch from the installed torch package):
cd rust/mem-cache
cd rust/sglang-radix-tree
LIBTORCH_USE_PYTORCH=1 \
LIBTORCH_BYPASS_VERSION_CHECK=1 \
CXXFLAGS="-include $PWD/torch_2_13_compat.h" \
@@ -1,4 +1,4 @@
"""Run the standalone mem-cache crate's native Rust unit tests."""
"""Run the standalone sglang-radix-tree crate's native Rust unit tests."""
import shutil
import subprocess
@@ -12,7 +12,7 @@ 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"
SGLANG_RADIX_TREE_MANIFEST = RUST_WORKSPACE / "sglang-radix-tree" / "Cargo.toml"
register_cpu_ci(est_time=900, suite="base-a-test-cpu")
@@ -21,19 +21,19 @@ register_cpu_ci(est_time=900, suite="base-a-test-cpu")
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):
class TestSGLangRadixTreeCargo(CustomTestCase):
def test_sglang_radix_tree_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}",
SGLANG_RADIX_TREE_MANIFEST.is_file(),
f"sglang-radix-tree manifest not found at {SGLANG_RADIX_TREE_MANIFEST}",
)
build = torch_build_configuration(
compat_header=MEM_CACHE_MANIFEST.parent / "torch_2_13_compat.h",
compat_header=SGLANG_RADIX_TREE_MANIFEST.parent / "torch_2_13_compat.h",
python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
)
proc = subprocess.run(
@@ -41,7 +41,7 @@ class TestMemCacheCargo(CustomTestCase):
"cargo",
"test",
"--manifest-path",
str(MEM_CACHE_MANIFEST),
str(SGLANG_RADIX_TREE_MANIFEST),
"--locked",
"--no-default-features",
],
@@ -55,7 +55,7 @@ class TestMemCacheCargo(CustomTestCase):
self.assertEqual(
proc.returncode,
0,
f"mem-cache native tests failed\n"
f"sglang-radix-tree native tests failed\n"
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
)
+1 -1
View File
@@ -542,7 +542,7 @@ crate-type = ["cdylib"]
),
(
"sglang.srt.mem_cache.rust_tree_core.mem_cache",
"mem_cache",
"sglang-radix-tree",
"mem_cache",
("python-extension",),
),