From 067353f67b56ea6143d9cd77b87b32edde74b6ac Mon Sep 17 00:00:00 2001 From: roikoren755 <26850796+roikoren755@users.noreply.github.com> Date: Thu, 12 Mar 2026 10:11:59 +0200 Subject: [PATCH] [Test] Refactor KL divergence and prefix cache branching to kits (#19715) --- python/sglang/test/kits/kl_divergence_kit.py | 42 ++ .../test/kits/prefix_cache_branching_kit.py | 50 +++ .../test/server_fixtures/default_fixture.py | 6 + .../4-gpu-models/test_qwen3_next_models.py | 138 +------ .../test_qwen3_next_models_mtp.py | 360 +++++------------- .../models/test_qwen3_next_models_fp4.py | 73 +--- .../models/test_qwen3_next_models_pcg.py | 57 +-- .../radix_cache/test_swa_radix_cache_kl.py | 70 +--- 8 files changed, 244 insertions(+), 552 deletions(-) create mode 100644 python/sglang/test/kits/kl_divergence_kit.py create mode 100644 python/sglang/test/kits/prefix_cache_branching_kit.py diff --git a/python/sglang/test/kits/kl_divergence_kit.py b/python/sglang/test/kits/kl_divergence_kit.py new file mode 100644 index 000000000..cffa84bcd --- /dev/null +++ b/python/sglang/test/kits/kl_divergence_kit.py @@ -0,0 +1,42 @@ +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_decode_cache_hit_helper, + test_input_output_logprobs_match_prefill_cache_hit_helper, +) + + +class KLDivergenceMixin: + kl_div_thres: float + kl_div_thres_decode: float | None = None + kl_div_thres_prefill: float | None = None + kl_div_max_samples: int = 32 + kl_div_prefill_max_new_tokens: int = 512 + kl_div_decode_max_new_tokens: int = 512 + + @classmethod + def _build_acc_thresholds(cls, threshold): + """Build an ACC_THRESHOLDS dict compatible with kl_test_utils.""" + return {cls.model: {"kl_div": threshold}} + + @classmethod + def test_input_output_logprobs_match_prefill_cache_hit(cls): + test_input_output_logprobs_match_prefill_cache_hit_helper( + base_url=cls.base_url, + ACC_THRESHOLDS=cls._build_acc_thresholds( + cls.kl_div_thres_prefill or cls.kl_div_thres + ), + model_name=cls.model, + max_samples=cls.kl_div_max_samples, + max_new_tokens=cls.kl_div_prefill_max_new_tokens, + ) + + @classmethod + def test_input_output_logprobs_match_decode_cache_hit(cls): + test_input_output_logprobs_match_decode_cache_hit_helper( + base_url=cls.base_url, + ACC_THRESHOLDS=cls._build_acc_thresholds( + cls.kl_div_thres_decode or cls.kl_div_thres + ), + model_name=cls.model, + max_samples=cls.kl_div_max_samples, + max_new_tokens=cls.kl_div_decode_max_new_tokens, + ) diff --git a/python/sglang/test/kits/prefix_cache_branching_kit.py b/python/sglang/test/kits/prefix_cache_branching_kit.py new file mode 100644 index 000000000..780e423a3 --- /dev/null +++ b/python/sglang/test/kits/prefix_cache_branching_kit.py @@ -0,0 +1,50 @@ +import requests + + +class PrefixCacheBranchingMixin: + cache_chunk_size: int + + @classmethod + def send_request_helper(cls, text: str): + response = requests.post( + cls.base_url + "/generate", + json={ + "text": text, + "sampling_params": { + "max_new_tokens": 1, + }, + }, + ) + return response.json() + + @classmethod + def test_prefix_cache_branching(cls): + cls.flush_cache() + branching_pos = 257 + text_prefix = "hi" * branching_pos + suffix_list = [ + "this" * cls.cache_chunk_size * 4, + "here" * cls.cache_chunk_size * 4, + "that" * cls.cache_chunk_size * 4, + ] + cache_hit_list = [False, False, True] + + # First request only prefill the entire sequence + # Second request won't have cache hit, but will cache the branching point + # Third request will have cache hit on the branching point + for i, (suffix, cache_hit) in enumerate( + zip(suffix_list, cache_hit_list, strict=True) + ): + result = cls.send_request_helper(text_prefix + suffix) + cached_tokens = result["meta_info"]["cached_tokens"] + if cache_hit: + expected_cached_tokens = ( + branching_pos // cls.cache_chunk_size * cls.cache_chunk_size + ) + assert ( + cached_tokens == expected_cached_tokens + ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}" + else: + assert ( + cached_tokens == 0 + ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0" diff --git a/python/sglang/test/server_fixtures/default_fixture.py b/python/sglang/test/server_fixtures/default_fixture.py index f5f355046..89cf1c8fc 100644 --- a/python/sglang/test/server_fixtures/default_fixture.py +++ b/python/sglang/test/server_fixtures/default_fixture.py @@ -2,6 +2,8 @@ import logging import time from contextlib import contextmanager +import requests + from sglang.srt.utils import kill_process_tree from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -65,3 +67,7 @@ class DefaultServerBase(CustomTestCase): def tearDownClass(cls): kill_process_tree(cls.process.pid) time.sleep(2) + + @classmethod + def flush_cache(cls): + requests.post(cls.base_url + "/flush_cache") diff --git a/test/registered/4-gpu-models/test_qwen3_next_models.py b/test/registered/4-gpu-models/test_qwen3_next_models.py index 982cb6dce..db3140cf5 100644 --- a/test/registered/4-gpu-models/test_qwen3_next_models.py +++ b/test/registered/4-gpu-models/test_qwen3_next_models.py @@ -1,129 +1,33 @@ import unittest -from types import SimpleNamespace -import requests - -from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.few_shot_gsm8k import run_eval -from sglang.test.kl_test_utils import ( - test_input_output_logprobs_match_decode_cache_hit_helper, - test_input_output_logprobs_match_prefill_cache_hit_helper, -) -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) +from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin +from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin +from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase register_cuda_ci(est_time=350, suite="stage-c-test-4-gpu-h100") QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" -ACC_THRESHOLDS = { - QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93}, -} - -def send_request_helper(base_url: str, text: str): - response = requests.post( - base_url + "/generate", - json={ - "text": text, - "sampling_params": { - "max_new_tokens": 1, - }, - }, - ) - return response.json() - - -class TestQwen3Next(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = QWEN3_NEXT_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", - "--mamba-scheduler-strategy", - "extra_buffer", - "--mamba-track-interval", - "128", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=200, - max_new_tokens=512, - parallel=128, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual( - metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] - ) - - def test_input_output_logprobs_match_prefill_cache_hit(self): - test_input_output_logprobs_match_prefill_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS, - self.model, - max_samples=32, - max_new_tokens=512, - ) - - def test_input_output_logprobs_match_decode_cache_hit(self): - test_input_output_logprobs_match_decode_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS, - self.model, - max_samples=32, - max_new_tokens=512, - ) - - def test_prefix_cache_branching(self): - print("running test_prefix_cache_branching") - requests.get(self.base_url + "/flush_cache") - branching_pos = 257 - text_prefix = "hi" * branching_pos - suffix_list = ["this" * 256, "here" * 256, "that" * 256] - cache_hit_list = [False, False, True] - - # First request only prefill the entire sequence - # Second request won't have cache hit, but will cache the branching point - # Third request will have cache hit on the branching point - for i, (suffix, cache_hit) in enumerate( - zip(suffix_list, cache_hit_list, strict=True) - ): - result = send_request_helper(self.base_url, text_prefix + suffix) - cached_tokens = result["meta_info"]["cached_tokens"] - if cache_hit: - expected_cached_tokens = branching_pos // 64 * 64 - assert ( - cached_tokens == expected_cached_tokens - ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}" - else: - assert ( - cached_tokens == 0 - ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0" - print("test_prefix_cache_branching passed") +class TestQwen3Next( + GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase +): + model = QWEN3_NEXT_MODEL + cache_chunk_size = 64 + gsm8k_accuracy_thres = 0.93 + kl_div_thres = 0.0025 + other_args = [ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + ] if __name__ == "__main__": diff --git a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py index 6f0784a83..215e96125 100644 --- a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py +++ b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py @@ -1,299 +1,109 @@ import unittest -from types import SimpleNamespace - -import requests from sglang.srt.environ import envs -from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.few_shot_gsm8k import run_eval -from sglang.test.kl_test_utils import ( - test_input_output_logprobs_match_decode_cache_hit_helper, - test_input_output_logprobs_match_prefill_cache_hit_helper, -) -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) +from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin +from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin +from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase register_cuda_ci(est_time=500, suite="stage-c-test-4-gpu-h100") QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" -ACC_THRESHOLDS = { - QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93}, -} -# MTP has higher KL divergence threshold -ACC_THRESHOLDS_MTP = { - QWEN3_NEXT_MODEL: {"kl_div": 0.008, "gsm8k": 0.93}, -} +class TestQwen3NextMTP(GSM8KMixin, KLDivergenceMixin, DefaultServerBase): + model = QWEN3_NEXT_MODEL + gsm8k_accuracy_thres = 0.93 + kl_div_thres = 0.0025 + other_args = [ + "--trust-remote-code", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--mem-fraction-static", + "0.8", + "--tp", + "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "no_buffer", + "--disable-radix-cache", + ] -def send_request_helper(base_url: str, text: str): - response = requests.post( - base_url + "/generate", - json={ - "text": text, - "sampling_params": { - "max_new_tokens": 1, - }, - }, - ) - return response.json() +class TestQwen3NextMTPTopk( + GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase +): + model = QWEN3_NEXT_MODEL + cache_chunk_size = 64 + gsm8k_accuracy_thres = 0.93 + kl_div_thres = 0.008 + other_args = [ + "--trust-remote-code", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "5", + "--speculative-eagle-topk", + "4", + "--speculative-num-draft-tokens", + "8", + "--mem-fraction-static", + "0.8", + "--tp", + "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + ] -class TestQwen3NextMTP(CustomTestCase): +# TODO(hzh): After merging the PR that fixes specv2 to correctly return log probs, +# add KLDivergenceMixin back. https://github.com/sgl-project/sglang/pull/18645 +class TestQwen3NextMTPV2(GSM8KMixin, DefaultServerBase): + model = QWEN3_NEXT_MODEL + gsm8k_accuracy_thres = 0.93 + other_args = [ + "--trust-remote-code", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--mem-fraction-static", + "0.8", + "--tp", + "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + ] + @classmethod def setUpClass(cls): - cls.model = QWEN3_NEXT_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=[ - "--trust-remote-code", - "--speculative-algorithm", - "NEXTN", - "--speculative-num-steps", - "3", - "--speculative-eagle-topk", - "1", - "--speculative-num-draft-tokens", - "4", - "--mem-fraction-static", - "0.8", - "--tp", - "4", - "--chunked-prefill-size", - "2048", - "--mamba-scheduler-strategy", - "no_buffer", - "--disable-radix-cache", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=200, - max_new_tokens=512, - parallel=128, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual( - metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] - ) - - def test_input_output_logprobs_match_prefill_cache_hit(self): - test_input_output_logprobs_match_prefill_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS, - self.model, - max_samples=32, - max_new_tokens=512, - ) - - def test_input_output_logprobs_match_decode_cache_hit(self): - test_input_output_logprobs_match_decode_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS, - self.model, - max_samples=32, - max_new_tokens=512, - ) - - -class TestQwen3NextMTPTopk(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = QWEN3_NEXT_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=[ - "--trust-remote-code", - "--speculative-algorithm", - "NEXTN", - "--speculative-num-steps", - "5", - "--speculative-eagle-topk", - "4", - "--speculative-num-draft-tokens", - "8", - "--mem-fraction-static", - "0.8", - "--tp", - "4", - "--chunked-prefill-size", - "2048", - "--mamba-scheduler-strategy", - "extra_buffer", - "--mamba-track-interval", - "128", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=200, - max_new_tokens=512, - parallel=128, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual( - metrics["accuracy"], ACC_THRESHOLDS_MTP[self.model]["gsm8k"] - ) - - def test_input_output_logprobs_match_prefill_cache_hit(self): - test_input_output_logprobs_match_prefill_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS_MTP, - self.model, - max_samples=32, - max_new_tokens=512, - ) - - def test_input_output_logprobs_match_decode_cache_hit(self): - test_input_output_logprobs_match_decode_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS_MTP, - self.model, - max_samples=32, - max_new_tokens=512, - ) - - def test_prefix_cache_branching(self): - print("running test_prefix_cache_branching") - requests.get(self.base_url + "/flush_cache") - branching_pos = 257 - text_prefix = "hi" * branching_pos - suffix_list = ["this" * 256, "here" * 256, "that" * 256] - cache_hit_list = [False, False, True] - - # First request only prefill the entire sequence - # Second request won't have cache hit, but will cache the branching point - # Third request will have cache hit on the branching point - for i, (suffix, cache_hit) in enumerate( - zip(suffix_list, cache_hit_list, strict=True) - ): - result = send_request_helper(self.base_url, text_prefix + suffix) - cached_tokens = result["meta_info"]["cached_tokens"] - if cache_hit: - expected_cached_tokens = branching_pos // 64 * 64 - assert ( - cached_tokens == expected_cached_tokens - ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}" - else: - assert ( - cached_tokens == 0 - ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0" - print("test_prefix_cache_branching passed") - - -class TestQwen3NextMTPV2(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = QWEN3_NEXT_MODEL envs.SGLANG_ENABLE_SPEC_V2.set(True) - 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=[ - "--trust-remote-code", - "--speculative-algorithm", - "NEXTN", - "--speculative-num-steps", - "3", - "--speculative-eagle-topk", - "1", - "--speculative-num-draft-tokens", - "4", - "--mem-fraction-static", - "0.8", - "--tp", - "4", - "--chunked-prefill-size", - "2048", - "--mamba-scheduler-strategy", - "extra_buffer", - "--mamba-track-interval", - "128", - ], - ) + super().setUpClass() @classmethod def tearDownClass(cls): envs.SGLANG_ENABLE_SPEC_V2.set(False) - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=200, - max_new_tokens=512, - parallel=128, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual( - metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] - ) - - # TODO(hzh): After merging the PR that fixes specv2 to correctly return log probs, re-open the tests below. https://github.com/sgl-project/sglang/pull/18645 - # def test_input_output_logprobs_match(self): - # test_input_output_logprobs_match_helper( - # self.base_url, - # ACC_THRESHOLDS, - # self.model, - # max_samples=32, - # max_new_tokens=512, - # ) - - # def test_input_output_logprobs_match_prefill_cache_hit(self): - # test_input_output_logprobs_match_prefill_cache_hit_helper( - # self.base_url, - # ACC_THRESHOLDS, - # self.model, - # max_samples=32, - # max_new_tokens=512, - # ) - - # def test_input_output_logprobs_match_decode_cache_hit(self): - # test_input_output_logprobs_match_decode_cache_hit_helper( - # self.base_url, - # ACC_THRESHOLDS, - # self.model, - # max_samples=32, - # max_new_tokens=512, - # ) + super().tearDownClass() if __name__ == "__main__": diff --git a/test/registered/models/test_qwen3_next_models_fp4.py b/test/registered/models/test_qwen3_next_models_fp4.py index 14be5ffa0..f08d99be1 100644 --- a/test/registered/models/test_qwen3_next_models_fp4.py +++ b/test/registered/models/test_qwen3_next_models_fp4.py @@ -1,70 +1,33 @@ import unittest -from types import SimpleNamespace -from sglang.srt.utils import get_device_sm, kill_process_tree +from sglang.srt.utils import get_device_sm from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.few_shot_gsm8k import run_eval -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) +from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase register_cuda_ci(est_time=500, suite="nightly-4-gpu-b200", nightly=True) QWEN3_NEXT_MODEL_FP4 = "nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4" -ACC_THRESHOLDS = { - QWEN3_NEXT_MODEL_FP4: {"kl_div": 0.0025, "gsm8k": 0.93}, -} - @unittest.skipIf( get_device_sm() < 100, "Test requires CUDA SM 100 or higher (Blackwell)" ) -class TestQwen3NextFp4(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = QWEN3_NEXT_MODEL_FP4 - 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", - "--quantization", - "modelopt_fp4", - "--mamba-scheduler-strategy", - "extra_buffer", - "--mamba-track-interval", - "128", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=200, - max_new_tokens=512, - parallel=128, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual( - metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] - ) +class TestQwen3NextFp4(GSM8KMixin, DefaultServerBase): + model = QWEN3_NEXT_MODEL_FP4 + gsm8k_accuracy_thres = 0.93 + other_args = [ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--quantization", + "modelopt_fp4", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + ] if __name__ == "__main__": diff --git a/test/registered/models/test_qwen3_next_models_pcg.py b/test/registered/models/test_qwen3_next_models_pcg.py index 7156ad042..1bb5f003b 100644 --- a/test/registered/models/test_qwen3_next_models_pcg.py +++ b/test/registered/models/test_qwen3_next_models_pcg.py @@ -3,17 +3,10 @@ Qwen3 Next piecewise CUDA graph tests. """ import unittest -from types import SimpleNamespace -from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.few_shot_gsm8k import run_eval -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) +from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase register_cuda_ci( est_time=400, @@ -22,46 +15,14 @@ register_cuda_ci( QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" -ACC_THRESHOLDS = { - QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93}, -} - -class TestQwen3NextPiecewiseCudaGraph(CustomTestCase): - - @classmethod - def setUpClass(cls): - cls.model = QWEN3_NEXT_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", - "4", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=200, - max_new_tokens=512, - parallel=128, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual( - metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] - ) +class TestQwen3NextPiecewiseCudaGraph(GSM8KMixin, DefaultServerBase): + model = QWEN3_NEXT_MODEL + gsm8k_accuracy_thres = 0.93 + other_args = [ + "--tp", + "4", + ] if __name__ == "__main__": diff --git a/test/registered/radix_cache/test_swa_radix_cache_kl.py b/test/registered/radix_cache/test_swa_radix_cache_kl.py index f51d863d4..617219ceb 100644 --- a/test/registered/radix_cache/test_swa_radix_cache_kl.py +++ b/test/registered/radix_cache/test_swa_radix_cache_kl.py @@ -1,69 +1,25 @@ import unittest -from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.kl_test_utils import ( - test_input_output_logprobs_match_decode_cache_hit_helper, - test_input_output_logprobs_match_prefill_cache_hit_helper, -) -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) +from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase MODEL = "openai/gpt-oss-20b" -ACC_THRESHOLDS = { - MODEL: {"kl_div": 0.002}, -} - register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu") -class TestSWARadixCacheKL(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = MODEL - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - # Use a lower mem-fraction-static to avoid OOM during input logprobs - # gathering. With PCG enabled, more memory is reserved for CUDA graph - # captures, so the static fraction should be lower. - other_args=[ - "--tp-size", - "1", - "--mem-fraction-static", - "0.70", - "--disable-piecewise-cuda-graph", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_input_output_logprobs_match_prefill_cache_hit(self): - test_input_output_logprobs_match_prefill_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS, - self.model, - max_samples=32, - max_new_tokens=512, - ) - - def test_input_output_logprobs_match_decode_cache_hit(self): - test_input_output_logprobs_match_decode_cache_hit_helper( - self.base_url, - ACC_THRESHOLDS, - self.model, - max_samples=32, - max_new_tokens=2048, - ) +class TestSWARadixCacheKL(KLDivergenceMixin, DefaultServerBase): + model = MODEL + kl_div_thres = 0.002 + kl_div_decode_max_new_tokens = 2048 + other_args = [ + "--tp-size", + "1", + "--mem-fraction-static", + "0.70", + "--disable-piecewise-cuda-graph", + ] if __name__ == "__main__":