From a77283fb02843669e1f534a24e59baa39e38de22 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Tue, 1 Sep 2026 03:03:44 -0700 Subject: [PATCH] [Rust] Rename mem-cache to sglang-radix-tree (#37290) --- .github/workflows/_pr-test-rust-ext-build.yml | 10 ++++----- .gitignore | 2 +- .pre-commit-config.yaml | 4 ++-- python/setup.py | 4 ++-- .../srt/mem_cache/rust_tree_core/extension.py | 2 +- rust/Cargo.toml | 6 ++--- .../.gitignore | 0 .../Cargo.lock | 22 +++++++++---------- .../Cargo.toml | 2 +- .../README.md | 4 ++-- .../src/components/full.rs | 0 .../src/components/mamba.rs | 0 .../src/components/mod.rs | 0 .../src/components/swa.rs | 0 .../src/lib.rs | 0 .../src/node.rs | 0 .../src/python_bindings.rs | 0 .../src/tests/components/base.rs | 0 .../src/tests/components/full.rs | 0 .../src/tests/components/mamba.rs | 0 .../src/tests/components/swa.rs | 0 .../src/tests/node.rs | 0 .../src/tests/test_utils.rs | 0 .../src/tests/unified_lru_list.rs | 0 .../src/tests/unified_tree_core.rs | 0 .../src/unified_lru_list.rs | 0 .../src/unified_tree_core.rs | 0 .../torch_2_13_compat.h | 0 ... test_run_sglang_radix_tree_rust_tests.py} | 18 +++++++-------- test/registered/rust/test_rust_extension.py | 2 +- 30 files changed, 38 insertions(+), 38 deletions(-) rename rust/{mem-cache => sglang-radix-tree}/.gitignore (100%) rename rust/{mem-cache => sglang-radix-tree}/Cargo.lock (99%) rename rust/{mem-cache => sglang-radix-tree}/Cargo.toml (97%) rename rust/{mem-cache => sglang-radix-tree}/README.md (97%) rename rust/{mem-cache => sglang-radix-tree}/src/components/full.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/components/mamba.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/components/mod.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/components/swa.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/lib.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/node.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/python_bindings.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/components/base.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/components/full.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/components/mamba.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/components/swa.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/node.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/test_utils.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/unified_lru_list.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/tests/unified_tree_core.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/unified_lru_list.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/src/unified_tree_core.rs (100%) rename rust/{mem-cache => sglang-radix-tree}/torch_2_13_compat.h (100%) rename test/registered/rust/{test_run_mem_cache_rust_tests.py => test_run_sglang_radix_tree_rust_tests.py} (73%) diff --git a/.github/workflows/_pr-test-rust-ext-build.yml b/.github/workflows/_pr-test-rust-ext-build.yml index 7631186e2..19d88f992 100644 --- a/.github/workflows/_pr-test-rust-ext-build.yml +++ b/.github/workflows/_pr-test-rust-ext-build.yml @@ -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", ], diff --git a/.gitignore b/.gitignore index d0dcf786d..7c068cb49 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index eba95baa0..cc6d6c1ed 100644 --- a/.pre-commit-config.yaml +++ b/.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 && 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 diff --git a/python/setup.py b/python/setup.py index 96b489d95..d014ac2c8 100644 --- a/python/setup.py +++ b/python/setup.py @@ -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) diff --git a/python/sglang/srt/mem_cache/rust_tree_core/extension.py b/python/sglang/srt/mem_cache/rust_tree_core/extension.py index 955fe7906..d2b05dd46 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/extension.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/extension.py @@ -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" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 553ae362a..765d64dc7 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -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" diff --git a/rust/mem-cache/.gitignore b/rust/sglang-radix-tree/.gitignore similarity index 100% rename from rust/mem-cache/.gitignore rename to rust/sglang-radix-tree/.gitignore diff --git a/rust/mem-cache/Cargo.lock b/rust/sglang-radix-tree/Cargo.lock similarity index 99% rename from rust/mem-cache/Cargo.lock rename to rust/sglang-radix-tree/Cargo.lock index d07a7c7e1..9028c9fc9 100644 --- a/rust/mem-cache/Cargo.lock +++ b/rust/sglang-radix-tree/Cargo.lock @@ -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" diff --git a/rust/mem-cache/Cargo.toml b/rust/sglang-radix-tree/Cargo.toml similarity index 97% rename from rust/mem-cache/Cargo.toml rename to rust/sglang-radix-tree/Cargo.toml index 69057cb0e..426fe97d4 100644 --- a/rust/mem-cache/Cargo.toml +++ b/rust/sglang-radix-tree/Cargo.toml @@ -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" diff --git a/rust/mem-cache/README.md b/rust/sglang-radix-tree/README.md similarity index 97% rename from rust/mem-cache/README.md rename to rust/sglang-radix-tree/README.md index 6479e493a..2a5ec00be 100644 --- a/rust/mem-cache/README.md +++ b/rust/sglang-radix-tree/README.md @@ -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" \ diff --git a/rust/mem-cache/src/components/full.rs b/rust/sglang-radix-tree/src/components/full.rs similarity index 100% rename from rust/mem-cache/src/components/full.rs rename to rust/sglang-radix-tree/src/components/full.rs diff --git a/rust/mem-cache/src/components/mamba.rs b/rust/sglang-radix-tree/src/components/mamba.rs similarity index 100% rename from rust/mem-cache/src/components/mamba.rs rename to rust/sglang-radix-tree/src/components/mamba.rs diff --git a/rust/mem-cache/src/components/mod.rs b/rust/sglang-radix-tree/src/components/mod.rs similarity index 100% rename from rust/mem-cache/src/components/mod.rs rename to rust/sglang-radix-tree/src/components/mod.rs diff --git a/rust/mem-cache/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs similarity index 100% rename from rust/mem-cache/src/components/swa.rs rename to rust/sglang-radix-tree/src/components/swa.rs diff --git a/rust/mem-cache/src/lib.rs b/rust/sglang-radix-tree/src/lib.rs similarity index 100% rename from rust/mem-cache/src/lib.rs rename to rust/sglang-radix-tree/src/lib.rs diff --git a/rust/mem-cache/src/node.rs b/rust/sglang-radix-tree/src/node.rs similarity index 100% rename from rust/mem-cache/src/node.rs rename to rust/sglang-radix-tree/src/node.rs diff --git a/rust/mem-cache/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs similarity index 100% rename from rust/mem-cache/src/python_bindings.rs rename to rust/sglang-radix-tree/src/python_bindings.rs diff --git a/rust/mem-cache/src/tests/components/base.rs b/rust/sglang-radix-tree/src/tests/components/base.rs similarity index 100% rename from rust/mem-cache/src/tests/components/base.rs rename to rust/sglang-radix-tree/src/tests/components/base.rs diff --git a/rust/mem-cache/src/tests/components/full.rs b/rust/sglang-radix-tree/src/tests/components/full.rs similarity index 100% rename from rust/mem-cache/src/tests/components/full.rs rename to rust/sglang-radix-tree/src/tests/components/full.rs diff --git a/rust/mem-cache/src/tests/components/mamba.rs b/rust/sglang-radix-tree/src/tests/components/mamba.rs similarity index 100% rename from rust/mem-cache/src/tests/components/mamba.rs rename to rust/sglang-radix-tree/src/tests/components/mamba.rs diff --git a/rust/mem-cache/src/tests/components/swa.rs b/rust/sglang-radix-tree/src/tests/components/swa.rs similarity index 100% rename from rust/mem-cache/src/tests/components/swa.rs rename to rust/sglang-radix-tree/src/tests/components/swa.rs diff --git a/rust/mem-cache/src/tests/node.rs b/rust/sglang-radix-tree/src/tests/node.rs similarity index 100% rename from rust/mem-cache/src/tests/node.rs rename to rust/sglang-radix-tree/src/tests/node.rs diff --git a/rust/mem-cache/src/tests/test_utils.rs b/rust/sglang-radix-tree/src/tests/test_utils.rs similarity index 100% rename from rust/mem-cache/src/tests/test_utils.rs rename to rust/sglang-radix-tree/src/tests/test_utils.rs diff --git a/rust/mem-cache/src/tests/unified_lru_list.rs b/rust/sglang-radix-tree/src/tests/unified_lru_list.rs similarity index 100% rename from rust/mem-cache/src/tests/unified_lru_list.rs rename to rust/sglang-radix-tree/src/tests/unified_lru_list.rs diff --git a/rust/mem-cache/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs similarity index 100% rename from rust/mem-cache/src/tests/unified_tree_core.rs rename to rust/sglang-radix-tree/src/tests/unified_tree_core.rs diff --git a/rust/mem-cache/src/unified_lru_list.rs b/rust/sglang-radix-tree/src/unified_lru_list.rs similarity index 100% rename from rust/mem-cache/src/unified_lru_list.rs rename to rust/sglang-radix-tree/src/unified_lru_list.rs diff --git a/rust/mem-cache/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs similarity index 100% rename from rust/mem-cache/src/unified_tree_core.rs rename to rust/sglang-radix-tree/src/unified_tree_core.rs diff --git a/rust/mem-cache/torch_2_13_compat.h b/rust/sglang-radix-tree/torch_2_13_compat.h similarity index 100% rename from rust/mem-cache/torch_2_13_compat.h rename to rust/sglang-radix-tree/torch_2_13_compat.h diff --git a/test/registered/rust/test_run_mem_cache_rust_tests.py b/test/registered/rust/test_run_sglang_radix_tree_rust_tests.py similarity index 73% rename from test/registered/rust/test_run_mem_cache_rust_tests.py rename to test/registered/rust/test_run_sglang_radix_tree_rust_tests.py index 55758516c..4efbec94b 100644 --- a/test/registered/rust/test_run_mem_cache_rust_tests.py +++ b/test/registered/rust/test_run_sglang_radix_tree_rust_tests.py @@ -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}", ) diff --git a/test/registered/rust/test_rust_extension.py b/test/registered/rust/test_rust_extension.py index 85ea9c2d1..2b93cf714 100644 --- a/test/registered/rust/test_rust_extension.py +++ b/test/registered/rust/test_rust_extension.py @@ -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",), ),