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
python/sglang/srt/mem_cache/unified_radix_cache.py
+1 -12
View File
@@ -1244,12 +1244,6 @@ class UnifiedRadixCache(BasePrefixCache):
self.dec_lock_ref(ancestor_node)
return None
logger.info(
"load_back: kv_tokens=%d, node_id=%d",
kv_tokens,
last_hit_node.id,
)
# Load H→D
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
aux_xfers.extend(anchor_kv_shared_indices_xfers)
@@ -1367,14 +1361,9 @@ class UnifiedRadixCache(BasePrefixCache):
req = params.req
if last_node.evicted or params.host_hit_length > 0:
logger.info(
"init_load_back triggered: node_id=%d, host_hit_length=%d",
last_node.id,
params.host_hit_length,
)
loading_values = self.load_back(last_node, mem_quota, req=req)
if loading_values is not None:
logger.info(
logger.debug(
"init_load_back success: loaded %d tokens for node %d",
len(loading_values),
last_node.id,
python/sglang/test/kl_test_utils.py
+7 -2
View File
@@ -118,8 +118,13 @@ def compare_kl_divergence(
# Common request helpers
def _flush_cache(base_url):
requests.post(base_url + "/flush_cache")
def _flush_cache(base_url, timeout_s=30):
response = requests.post(
base_url + "/flush_cache",
params={"timeout": timeout_s},
timeout=timeout_s + 10,
)
response.raise_for_status()
def _generate(
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__":