[Rust] Rename mem-cache to sglang-radix-tree (#37290)
This commit is contained in:
@@ -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
@@ -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/
|
||||
|
||||
@@ -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
@@ -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
@@ -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"
|
||||
|
||||
+11
-11
@@ -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" \
|
||||
+9
-9
@@ -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}",
|
||||
)
|
||||
|
||||
@@ -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",),
|
||||
),
|
||||
|
||||
Reference in New Issue
Block a user