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

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