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

[UnifiedTree]: Split unified tree kl ci into multiple files to reduce GPU usage. (#26666)

Browse Source
This commit is contained in:
Zhangheng
2026-05-29 02:46:34 -07:00
committed by GitHub
parent 4585f8eb95
commit eb5d4827e8
7 changed files with 192 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py
-77
View File
@@ -1,77 +0,0 @@
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/radix_cache/test_unified_radix_cache_kl_mamba.py
-66
View File
@@ -1,66 +0,0 @@
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
from sglang.test.kl_multiturn_utils import (
get_input_ids,
make_mamba_decode_assert,
make_mamba_prefill_assert,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=250, stage="base-c", runner_config="4-gpu-h100")
MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
MAMBA_CHUNK_SIZE = 64
MAMBA_TRACK_INTERVAL = 128
class TestUnifiedMambaRadixCache(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),
],
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__":
unittest.main()
test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py → test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
+2 -67
View File
@@ -3,83 +3,18 @@ import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
from sglang.test.kl_multiturn_utils import (
get_input_ids,
make_mamba_decode_assert,
make_mamba_prefill_assert,
)
from sglang.test.kl_multiturn_utils import get_input_ids
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
)
MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
MAMBA_CHUNK_SIZE = 64
MAMBA_TRACK_INTERVAL = 128
DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_LAUNCH_TIMEOUT = 3600
register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200")
class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""Mamba hybrid + HiCache + UnifiedRadixCache."""
kl_threshold = 0.005
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-mamba-cache-size",
"500",
"--max-running-requests",
"4",
"--weight-loader-prefetch-checkpoints",
],
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)
register_cuda_ci(est_time=768, stage="base-c", runner_config="4-gpu-h100")
def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label):
test/registered/radix_cache/test_unified_radix_cache_kl_full.py → test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
+2 -2
View File
@@ -11,7 +11,7 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=250, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=250, stage="base-b", runner_config="2-gpu-large")
FULL_MODEL = "Qwen/Qwen3-32B"
@@ -31,7 +31,7 @@ class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"2",
"--mem-fraction-static",
"0.80",
"--page-size",
test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py → test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hicache_nightly.py
View File
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py
+186
View File
@@ -0,0 +1,186 @@
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.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
from sglang.test.kl_multiturn_utils import (
get_input_ids,
make_mamba_decode_assert,
make_mamba_prefill_assert,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=768, stage="base-c", runner_config="4-gpu-h100")
MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
MAMBA_CHUNK_SIZE = 64
MAMBA_TRACK_INTERVAL = 128
class TestUnifiedMambaRadixCache(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),
],
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)
# ─── Mamba + HiCache L2 ──────────────────────────────────────────────────────
class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""Mamba hybrid + HiCache L2 + UnifiedRadixCache."""
kl_threshold = 0.005
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-mamba-cache-size",
"500",
"--max-running-requests",
"4",
"--weight-loader-prefetch-checkpoints",
],
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)
# ─── Mamba + HiCache L3 (file backend) ───────────────────────────────────────
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/radix_cache/test_unified_radix_cache_kl_swa.py → test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py
+2 -2
View File
@@ -12,7 +12,7 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=250, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=250, stage="base-b", runner_config="2-gpu-large")
SWA_MODEL = "openai/gpt-oss-20b"
@@ -38,7 +38,7 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"2",
"--mem-fraction-static",
"0.7",
"--disable-piecewise-cuda-graph",