minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Radix Cache] Add Rust TreeCore backend with shared parity tests (#32710)

Browse Source 
Co-authored-by: alphabetc1 <2508695655@qq.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Jialin Ouyang
2026-09-01 00:26:20 +08:00
committed by GitHub
co-authored by alphabetc1 ispobock
parent 52e1c24744
commit 9cf157c252
72 changed files with 39973 additions and 396 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
+49 -19
View File
@@ -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 -3
View File
@@ -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
+28 -8
View File
@@ -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()