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

[UnifiedRadixTree]: Support L3 HiStorage framework (#26062)

Browse Source
This commit is contained in:
Zhangheng
2026-05-26 22:38:00 +08:00
committed by GitHub
parent c47f0e7cdd
commit 38f32c38ab
12 changed files with 1148 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py
+1 -1
View File
@@ -24,7 +24,7 @@ MAMBA_TRACK_INTERVAL = 128
DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_LAUNCH_TIMEOUT = 3600
register_cuda_ci(est_time=745, stage="base-c", runner_config="8-gpu-h200")
register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200")
class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py
+42 -17
View File
@@ -28,8 +28,8 @@ GLM5_LAUNCH_TIMEOUT = 3600
register_cuda_ci(est_time=900, suite="nightly-8-gpu-h200", nightly=True)
class GSM8KTwoPassMixin:
"""Mixin: run GSM8K twice with flush in between, verify accuracy diff.
class AccuracyTwoPassMixin:
"""Mixin: run an eval twice with flush in between, verify accuracy diff.
Subclass must provide:
- self.base_url
@@ -38,9 +38,14 @@ class GSM8KTwoPassMixin:
gsm8k_threshold: float = 0.90
num_gsm8k_questions: int = 200
max_accuracy_diff: float = 0.02
gsm8k_parallel: int = 40
mmlu_threshold: float = 0.75
num_mmlu_examples: int = 200
mmlu_num_threads: int = 32
max_accuracy_diff: float = 0.02
def _run_gsm8k(self):
from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
@@ -57,6 +62,19 @@ class GSM8KTwoPassMixin:
metrics = run_few_shot_gsm8k(args)
return metrics["accuracy"]
def _run_mmlu(self):
from sglang.test.run_eval import run_eval as run_simple_eval
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mmlu",
num_examples=self.num_mmlu_examples,
num_threads=self.mmlu_num_threads,
)
metrics = run_simple_eval(args)
return metrics["score"]
def _flush_cache(self):
response = requests.post(
self.base_url + "/flush_cache",
@@ -65,45 +83,52 @@ class GSM8KTwoPassMixin:
)
response.raise_for_status()
def test_gsm8k_two_passes(self):
"""Run GSM8K twice with flush in between, verify accuracy diff <= max_accuracy_diff."""
def _two_pass(self, name: str, run_fn, threshold: float):
# First pass
acc1 = self._run_gsm8k()
print(f"[{self.__class__.__name__}] GSM8K pass 1 accuracy: {acc1:.3f}")
acc1 = run_fn()
print(f"[{self.__class__.__name__}] {name} pass 1 accuracy: {acc1:.3f}")
self.assertGreaterEqual(
acc1,
self.gsm8k_threshold,
f"Pass 1 accuracy {acc1:.3f} < threshold {self.gsm8k_threshold}",
threshold,
f"{name} pass 1 accuracy {acc1:.3f} < threshold {threshold}",
)
# Flush cache
self._flush_cache()
# Second pass
acc2 = self._run_gsm8k()
print(f"[{self.__class__.__name__}] GSM8K pass 2 accuracy: {acc2:.3f}")
acc2 = run_fn()
print(f"[{self.__class__.__name__}] {name} pass 2 accuracy: {acc2:.3f}")
self.assertGreaterEqual(
acc2,
self.gsm8k_threshold,
f"Pass 2 accuracy {acc2:.3f} < threshold {self.gsm8k_threshold}",
threshold,
f"{name} pass 2 accuracy {acc2:.3f} < threshold {threshold}",
)
# Verify diff
# Verify diff (only fail when 2nd pass regressed)
if acc1 > acc2:
diff = abs(acc1 - acc2)
print(
f"[{self.__class__.__name__}] Accuracy diff: {diff:.3f} "
f"[{self.__class__.__name__}] {name} accuracy diff: {diff:.3f} "
f"(max allowed: {self.max_accuracy_diff})"
)
self.assertLessEqual(
diff,
self.max_accuracy_diff,
f"Accuracy diff {diff:.3f} exceeds max {self.max_accuracy_diff} "
f"{name} accuracy diff {diff:.3f} exceeds max {self.max_accuracy_diff} "
f"(pass1={acc1:.3f}, pass2={acc2:.3f})",
)
def test_gsm8k_two_passes(self):
"""Run GSM8K twice with flush in between, verify accuracy diff <= max_accuracy_diff."""
self._two_pass("GSM8K", self._run_gsm8k, self.gsm8k_threshold)
class TestGLM5HiCacheL3GSM8K(GSM8KTwoPassMixin, CustomTestCase):
def test_mmlu_two_passes(self):
"""Run MMLU twice with flush in between, verify accuracy diff <= max_accuracy_diff."""
self._two_pass("MMLU", self._run_mmlu, self.mmlu_threshold)
class TestGLM5HiCacheL3Accuracy(AccuracyTwoPassMixin, CustomTestCase):
"""GLM-5.1-FP8 + HiCache L3 (file backend), with HiRadixTree."""
@classmethod
test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py
+77
View File
@@ -0,0 +1,77 @@
import os
import shutil
import tempfile
import unittest
from test_unified_radix_cache_kl_hicache_nightly import AccuracyTwoPassMixin
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
MAMBA_TRACK_INTERVAL = 128
register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200")
class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
"""Mamba hybrid + HiCache L3 (file backend) + UnifiedRadixCache."""
@classmethod
def setUpClass(cls):
cls.model = MAMBA_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.hicache_dir = tempfile.mkdtemp(prefix="hicache_l3_mamba_")
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",
"2",
"--hicache-write-policy",
"write_through",
"--hicache-storage-prefetch-policy",
"wait_complete",
"--hicache-io-backend",
"direct",
"--hicache-mem-layout",
"page_first_direct",
"--hicache-storage-backend",
"file",
"--max-mamba-cache-size",
"500",
"--weight-loader-prefetch-checkpoints",
],
env={
"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if os.path.isdir(cls.hicache_dir):
shutil.rmtree(cls.hicache_dir, ignore_errors=True)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
-7
View File
@@ -2274,13 +2274,6 @@ class UnifiedRadixCacheSuite:
tokens = self._swa_anchor_chain_tokens(len(chain))
return tree, chain, n, y, x, tokens
def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self):
tree, _, n, y, x, tokens = self._swa_anchor_setup()
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertIs(result.best_match_node, x)
self.assertIs(result.last_device_node, n.parent)
self.assertIs(result.last_host_node, y)
def test_hicache_swa_load_back_anchored_on_best_match_node(self):
tree, _, _, y, x, _ = self._swa_anchor_setup()
ps = self.cfg.page_size