From 125e17efd547650c6c2e16e7171fd476a830c38c Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 7 Dec 2025 23:55:00 +0800 Subject: [PATCH] Add small model test for spec v2 + dp + trtllm_mla (#14576) --- .../srt/model_executor/forward_batch_info.py | 5 +- .../test_eagle_infer_beta_dp_attention.py | 133 +++++++++++++----- test/srt/ep/test_deepep_large.py | 4 + 3 files changed, 101 insertions(+), 41 deletions(-) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 3437f2ef8..3a85e6a7e 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -52,6 +52,7 @@ from sglang.srt.layers.dp_attention import ( set_is_extend_in_batch, ) from sglang.srt.utils import get_compiler_backend, is_npu, support_triton +from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.layers.attention.base_attn_backend import AttentionBackend @@ -731,9 +732,7 @@ class ForwardBatch: for i in range(sync_group_size): # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim. # there is no reduce-scatter in LM logprob, so we do not need to adjust the padded length for logprob - global_num_tokens[i] = ( - (global_num_tokens[i] - 1) // attn_tp_size + 1 - ) * attn_tp_size + global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_tp_size) dp_padding_mode = DpPaddingMode.get_dp_padding_mode( self.is_extend_in_batch, global_num_tokens diff --git a/test/manual/test_eagle_infer_beta_dp_attention.py b/test/manual/test_eagle_infer_beta_dp_attention.py index 382196a18..199726be1 100644 --- a/test/manual/test_eagle_infer_beta_dp_attention.py +++ b/test/manual/test_eagle_infer_beta_dp_attention.py @@ -1,12 +1,15 @@ -import os 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.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k from sglang.test.test_utils import ( + DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST, + DEFAULT_MODEL_NAME_FOR_TEST_MLA, + DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, @@ -15,14 +18,75 @@ from sglang.test.test_utils import ( write_github_step_summary, ) -FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4" + +def test_gsm8k(base_url: str): + requests.get(base_url + "/flush_cache") + + 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(base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + server_info = requests.get(base_url + "/get_server_info") + avg_spec_accept_length = server_info.json()["internal_states"][0][ + "avg_spec_accept_length" + ] + + print(f"{metrics=}") + print(f"{avg_spec_accept_length=}") + return metrics, avg_spec_accept_length -class TestEagleDPAttnServerBase(CustomTestCase): +class TestEagleDPAttnServerSmall(CustomTestCase): @classmethod def setUpClass(cls): - os.environ["SGLANG_ENABLE_SPEC_V2"] = "1" - cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--tp-size", + "2", + "--dp-size", + "2", + "--enable-dp-attention", + "--speculative-draft-model-path", + DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN, + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + ] + with envs.SGLANG_ENABLE_SPEC_V2.override(True): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_a_gsm8k(self): + metrics, avg_spec_accept_length = test_gsm8k(self.base_url) + self.assertGreater(metrics["accuracy"], 0.64) + self.assertGreater(avg_spec_accept_length, 1.4) + + +class TestEagleDPAttnServerLarge(CustomTestCase): + # FIXME: move this large mode test into nightly tests + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST cls.base_url = DEFAULT_URL_FOR_TEST other_args = [ "--tp-size", @@ -46,52 +110,45 @@ class TestEagleDPAttnServerBase(CustomTestCase): "4", "--kv-cache-dtype", "fp8_e4m3", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', ] - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=other_args, - ) + with envs.SGLANG_ENABLE_SPEC_V2.override(True): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) @classmethod def tearDownClass(cls): kill_process_tree(cls.process.pid) - if "SGLANG_ENABLE_SPEC_V2" in os.environ: - del os.environ["SGLANG_ENABLE_SPEC_V2"] - def test_a_gsm8k( - self, - ): # Append an "a" to make this test run first (alphabetically) to warm up the server - requests.get(self.base_url + "/flush_cache") - - 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_few_shot_gsm8k(args) - print(f"{metrics=}") - - server_info = requests.get(self.base_url + "/get_server_info") - avg_spec_accept_length = server_info.json()["internal_states"][0][ - "avg_spec_accept_length" - ] - print(f"{avg_spec_accept_length=}") + def test_a_gsm8k(self): + metrics, avg_spec_accept_length = test_gsm8k(self.base_url) + self.assertGreater(metrics["accuracy"], 0.94) + self.assertGreater(avg_spec_accept_length, 2.04) if is_in_ci(): write_github_step_summary( f"### test_gsm8k (deepseek-v3-fp4 mtp)\n" f'{metrics["accuracy"]=:.3f}\n' f"{avg_spec_accept_length=:.2f}\n" ) - self.assertGreater(metrics["accuracy"], 0.94) - self.assertGreater(avg_spec_accept_length, 2.04) if __name__ == "__main__": - unittest.main() + # Force the unittest to run the small test first + s = unittest.TestSuite() + small_test = unittest.defaultTestLoader.loadTestsFromTestCase( + TestEagleDPAttnServerSmall + ) + large_test = unittest.defaultTestLoader.loadTestsFromTestCase( + TestEagleDPAttnServerLarge + ) + s.addTest(small_test) + s.addTest(large_test) + + runner = unittest.TextTestRunner() + runner.run(s) diff --git a/test/srt/ep/test_deepep_large.py b/test/srt/ep/test_deepep_large.py index 8a5234605..a5f3c2f95 100644 --- a/test/srt/ep/test_deepep_large.py +++ b/test/srt/ep/test_deepep_large.py @@ -47,6 +47,8 @@ class TestDeepseek(CustomTestCase): "--max-running-requests", "2048", "--disable-radix-cache", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', ], ) @@ -111,6 +113,8 @@ class TestDeepseekMTP(CustomTestCase): "--speculative-num-draft-tokens", "2", "--disable-radix-cache", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', ], )