From eb5d4827e83dd5b4b215e42cebfb7a829a6d651b Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Fri, 29 May 2026 17:46:34 +0800 Subject: [PATCH] [UnifiedTree]: Split unified tree kl ci into multiple files to reduce GPU usage. (#26666) --- ...st_unified_radix_cache_kl_hicache_part2.py | 77 -------- .../test_unified_radix_cache_kl_mamba.py | 66 ------- .../test_unified_radix_cache_kl_dsv4.py} | 69 +------ .../test_unified_radix_cache_kl_full.py | 4 +- ..._unified_radix_cache_kl_hicache_nightly.py | 0 .../test_unified_radix_cache_kl_mamba.py | 186 ++++++++++++++++++ .../test_unified_radix_cache_kl_swa.py | 4 +- 7 files changed, 192 insertions(+), 214 deletions(-) delete mode 100644 test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py delete mode 100644 test/registered/radix_cache/test_unified_radix_cache_kl_mamba.py rename test/registered/radix_cache/{test_unified_radix_cache_kl_hicache.py => unified_radix_tree/test_unified_radix_cache_kl_dsv4.py} (59%) rename test/registered/radix_cache/{ => unified_radix_tree}/test_unified_radix_cache_kl_full.py (93%) rename test/registered/radix_cache/{ => unified_radix_tree}/test_unified_radix_cache_kl_hicache_nightly.py (100%) create mode 100644 test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py rename test/registered/radix_cache/{ => unified_radix_tree}/test_unified_radix_cache_kl_swa.py (93%) diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py deleted file mode 100644 index 99331421b..000000000 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_part2.py +++ /dev/null @@ -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() diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_mamba.py b/test/registered/radix_cache/test_unified_radix_cache_kl_mamba.py deleted file mode 100644 index 2e540881f..000000000 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_mamba.py +++ /dev/null @@ -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() diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py similarity index 59% rename from test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py rename to test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index 9a108c9c7..d73c68f02 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py @@ -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): diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_full.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py similarity index 93% rename from test/registered/radix_cache/test_unified_radix_cache_kl_full.py rename to test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py index cb8e6e48d..1f66f6d9c 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_full.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py @@ -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", diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hicache_nightly.py similarity index 100% rename from test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py rename to test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hicache_nightly.py diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py new file mode 100644 index 000000000..4512a1aa3 --- /dev/null +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py @@ -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() diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_swa.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py similarity index 93% rename from test/registered/radix_cache/test_unified_radix_cache_kl_swa.py rename to test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py index 1f8b95239..1846a9c05 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_swa.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py @@ -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",