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

[UnifiedRadixTree]: Add KL accuracy CI for UnifiedTree with HiCache (#24346)

Browse Source
This commit is contained in:
Zhangheng
2026-05-04 20:18:10 +08:00
committed by GitHub
parent 84f3b44916
commit 05aed5e1d5
3 changed files with 63 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/radix_cache/test_unified_radix_cache_kl.py
+55 -52
View File
@@ -34,7 +34,7 @@ MAMBA_TRACK_INTERVAL = 128
SWA_MODEL = "openai/gpt-oss-20b"
FULL_MODEL = "Qwen/Qwen3-32B"
register_cuda_ci(est_time=632, suite="stage-c-test-4-gpu-h100")
register_cuda_ci(est_time=760, suite="stage-c-test-4-gpu-h100")
class UnifiedRadixTreeTestMixin:
@@ -235,8 +235,8 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""SWA hybrid + UnifiedRadixCache."""
kl_threshold = 0.03
gsm8k_threshold = 0.75
mmlu_threshold = 0.75
gsm8k_threshold = 0.7
mmlu_threshold = 0.7
@classmethod
def setUpClass(cls):
@@ -253,7 +253,7 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"0.7",
"--disable-piecewise-cuda-graph",
],
env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "0"},
env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18)
@@ -262,54 +262,57 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
kill_process_tree(cls.process.pid)
# TODO(hzh): Currently, HiCache is not fully compatible with the CI CUDA13 environment; we need to wait for the fix before re-enabling the tests below.
# class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
# """Mamba hybrid + UnifiedRadixCache."""
#
# kl_threshold = 0.003
# prefill_cache_assert = staticmethod(
# make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE)
# )
# decode_cache_assert = staticmethod(
# make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL)
# )
#
# @classmethod
# def setUpClass(cls):
# cls.model = MAMBA_MODEL
# cls.base_url = DEFAULT_URL_FOR_TEST
# cls.process = popen_launch_server(
# cls.model,
# cls.base_url,
# timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
# other_args=[
# "--tp-size",
# "4",
# "--chunked-prefill-size",
# "2048",
# "--mem-fraction-static",
# "0.85",
# "--mamba-scheduler-strategy",
# "extra_buffer",
# "--mamba-track-interval",
# str(MAMBA_TRACK_INTERVAL),
# "--enable-hierarchical-cache",
# "--hicache-ratio",
# "1.5",
# "--hicache-write-policy",
# "write_through",
# "--hicache-io-backend",
# "direct",
# "--hicache-mem-layout",
# "page_first_direct",
# ],
# env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
# )
# cls.input_ids = get_input_ids(cls.model, num_samples=18)
#
# @classmethod
# def tearDownClass(cls):
# kill_process_tree(cls.process.pid)
class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""Mamba hybrid + HiCache + UnifiedRadixCache."""
kl_threshold = 0.003
prefill_cache_assert = staticmethod(
make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE)
)
decode_cache_assert = staticmethod(
make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL)
)
@classmethod
def setUpClass(cls):
cls.model = MAMBA_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"--chunked-prefill-size",
"2048",
"--mem-fraction-static",
"0.85",
"--mamba-scheduler-strategy",
"extra_buffer",
"--mamba-track-interval",
str(MAMBA_TRACK_INTERVAL),
"--enable-hierarchical-cache",
"--hicache-ratio",
"4",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"direct",
"--hicache-mem-layout",
"page_first_direct",
"--max-total-tokens",
"12000",
"--max-running-requests",
"4",
],
env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":