diff --git a/python/sglang/test/kits/pd_parity_kit.py b/python/sglang/test/kits/pd_parity_kit.py new file mode 100644 index 000000000..19b5ecb20 --- /dev/null +++ b/python/sglang/test/kits/pd_parity_kit.py @@ -0,0 +1,70 @@ +import time + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.server_fixtures.disaggregation_fixture import assert_process_healthy +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + popen_launch_server, +) + + +class PDLogprobParityMixin: + # Mix in before the PD server fixture, which owns the P/D launches. + reference_parallel_args = [] + baseline_args = [] + + @staticmethod + def generate(base_url): + response = requests.post( + base_url + "/generate", + json={ + "input_ids": [1] + [100 + i % 1000 for i in range(256)], + "sampling_params": { + "temperature": 0, + "max_new_tokens": 4, + "ignore_eos": True, + }, + "return_logprob": True, + "top_logprobs_num": 5, + }, + timeout=120, + ) + response.raise_for_status() + return response.json()["meta_info"] + + def test_logprob_parity(self): + baseline = popen_launch_server( + self.model, + self.lb_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=self.reference_parallel_args + + ["--trust-remote-code"] + + self.baseline_args, + env=self.extra_prefill_env, + ) + try: + reference = self.generate(self.lb_url) + finally: + kill_process_tree(baseline.pid, wait_timeout=60) + time.sleep(5) + + self.launch_all() + disaggregated = self.generate(self.lb_url) + + reference_logprobs = reference["output_token_logprobs"] + disaggregated_logprobs = disaggregated["output_token_logprobs"] + self.assertEqual( + [item[1] for item in reference_logprobs], + [item[1] for item in disaggregated_logprobs], + ) + self.assertEqual(len(reference_logprobs), 4) + for reference_item, disaggregated_item in zip( + reference_logprobs, disaggregated_logprobs + ): + self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05) + + assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url) + assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url) + assert_process_healthy(self, "decode", self.process_decode, self.decode_url) diff --git a/test/registered/disaggregation/test_disaggregation_kimi_linear.py b/test/registered/disaggregation/test_disaggregation_kimi_linear.py index 49c93e6ec..ebe3295e1 100644 --- a/test/registered/disaggregation/test_disaggregation_kimi_linear.py +++ b/test/registered/disaggregation/test_disaggregation_kimi_linear.py @@ -1,24 +1,17 @@ -import time import unittest -import requests - -from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin from sglang.test.server_fixtures.disaggregation_fixture import ( PDDisaggregationServerBase, - assert_process_healthy, -) -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - popen_launch_server, ) register_cuda_ci(est_time=480, stage="base-c", runner_config="4-gpu-h100") KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random" SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"} -SERVER_ARGS = [ + +DETERMINISTIC_ARGS = [ "--skip-tokenizer-init", "--random-seed", "1", @@ -34,74 +27,19 @@ SERVER_ARGS = [ ] -class TestKimiLinearHeterogeneousTPDisaggregation(PDDisaggregationServerBase): +class TestKimiLinearHeterogeneousTPDisaggregation( + PDLogprobParityMixin, PDDisaggregationServerBase +): + model = KIMI_LINEAR_MODEL + extra_prefill_env = SERVER_ENV + extra_decode_env = SERVER_ENV prefill_tp_size = 2 decode_tp_size = 1 decode_base_gpu_id = 2 reference_parallel_args = ["--tp-size", "2"] - extra_prefill_args = SERVER_ARGS - extra_decode_args = SERVER_ARGS - extra_prefill_env = SERVER_ENV - extra_decode_env = SERVER_ENV - - @classmethod - def setUpClass(cls): - super().setUpClass() - cls.model = KIMI_LINEAR_MODEL - - @staticmethod - def generate(base_url): - response = requests.post( - base_url + "/generate", - json={ - "input_ids": [1] + [100 + i % 1000 for i in range(256)], - "sampling_params": { - "temperature": 0, - "max_new_tokens": 4, - "ignore_eos": True, - }, - "return_logprob": True, - "top_logprobs_num": 5, - }, - timeout=120, - ) - response.raise_for_status() - return response.json()["meta_info"] - - def test_logprob_parity(self): - baseline = popen_launch_server( - self.model, - self.lb_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=self.reference_parallel_args - + ["--trust-remote-code"] - + SERVER_ARGS, - env=SERVER_ENV, - ) - try: - reference = self.generate(self.lb_url) - finally: - kill_process_tree(baseline.pid, wait_timeout=60) - time.sleep(5) - - self.launch_all() - disaggregated = self.generate(self.lb_url) - - reference_logprobs = reference["output_token_logprobs"] - disaggregated_logprobs = disaggregated["output_token_logprobs"] - self.assertEqual( - [item[1] for item in reference_logprobs], - [item[1] for item in disaggregated_logprobs], - ) - self.assertEqual(len(reference_logprobs), 4) - for reference_item, disaggregated_item in zip( - reference_logprobs, disaggregated_logprobs - ): - self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05) - - assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url) - assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url) - assert_process_healthy(self, "decode", self.process_decode, self.decode_url) + baseline_args = DETERMINISTIC_ARGS + extra_prefill_args = DETERMINISTIC_ARGS + extra_decode_args = DETERMINISTIC_ARGS class TestKimiLinearPipelineDisaggregation(TestKimiLinearHeterogeneousTPDisaggregation): @@ -109,7 +47,7 @@ class TestKimiLinearPipelineDisaggregation(TestKimiLinearHeterogeneousTPDisaggre decode_tp_size = 1 decode_base_gpu_id = 2 reference_parallel_args = ["--tp-size", "1", "--pp-size", "2"] - extra_prefill_args = SERVER_ARGS + ["--pp-size", "2"] + extra_prefill_args = DETERMINISTIC_ARGS + ["--pp-size", "2"] if __name__ == "__main__": diff --git a/test/registered/disaggregation/test_disaggregation_unified_memory.py b/test/registered/disaggregation/test_disaggregation_unified_memory.py index 02e2cac4c..3f63c1a87 100644 --- a/test/registered/disaggregation/test_disaggregation_unified_memory.py +++ b/test/registered/disaggregation/test_disaggregation_unified_memory.py @@ -1,53 +1,24 @@ -"""PD disaggregation with --enable-unified-memory (MLA hybrid-Mamba). - -Guards the unified-memory PD transfer scheme end to end: whole page-envelope -KV registration (`UnifiedMLATokenToKVPool.get_contiguous_buf_infos`), whole -slot-envelope KDA/mamba state transfer, virtual->physical index translation at -the prefill send / decode prealloc sites, and the compaction move gate. A -regression in any of them shifts the decode-side KV/state bytes and breaks -logprob parity with the non-PD unified-memory reference. - -`--attention-backend` is deliberately NOT pinned, matching -`models_e2e/test_kimi_linear_unified_memory.py`, which documents that pinning -hides defects reachable only under the resolved default. The transferred bytes -are backend-independent, so the default (fa3 on this suite's H100 runner) covers -this file's subject either way. The linear-attn/Mamba backends stay pinned to -triton -- the page-major layout requires them. - -`--enable-deterministic-inference` is deliberately NOT set. It would only guard -against batch-shape-dependent kernel variation, and the reference and P+D paths -run the same shapes: measured, two fresh servers on separate GPUs produce -bit-identical logits without it. Setting it would narrow the test to the -batch-invariant op set and a non-default sampling backend -- a less -representative config -- and couple a PD-transfer test to the deterministic code -path, so a defect there would fail this file for an unrelated reason. -""" - -import time import unittest -import requests - -from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin from sglang.test.server_fixtures.disaggregation_fixture import ( PDDisaggregationServerBase, - assert_process_healthy, -) -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - popen_launch_server, ) -register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-h100") +register_cuda_ci(est_time=900, stage="base-b", runner_config="2-gpu-large") KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random" SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"} -SERVER_ARGS = [ + +# --attention-backend and --enable-deterministic-inference are deliberately +# absent: bytes are backend-independent and both paths run identical shapes. +UNIFIED_MEMORY_ARGS = [ "--skip-tokenizer-init", "--random-seed", "1", "--enable-unified-memory", + # Page-major layout requires triton for both. "--linear-attn-backend", "triton", "--mamba-backend", @@ -63,75 +34,19 @@ SERVER_ARGS = [ ] -class TestUnifiedMemoryDisaggregation(PDDisaggregationServerBase): +class TestUnifiedMemoryDisaggregation(PDLogprobParityMixin, PDDisaggregationServerBase): """1 prefill + 1 decode, both with --enable-unified-memory, vs a non-PD unified-memory reference server.""" + model = KIMI_LINEAR_MODEL + extra_prefill_env = SERVER_ENV + extra_decode_env = SERVER_ENV prefill_tp_size = 1 decode_tp_size = 1 decode_base_gpu_id = 1 - extra_prefill_args = SERVER_ARGS - extra_decode_args = SERVER_ARGS - extra_prefill_env = SERVER_ENV - extra_decode_env = SERVER_ENV - baseline_args = SERVER_ARGS - - @classmethod - def setUpClass(cls): - super().setUpClass() - cls.model = KIMI_LINEAR_MODEL - - @staticmethod - def generate(base_url): - response = requests.post( - base_url + "/generate", - json={ - "input_ids": [1] + [100 + i % 1000 for i in range(256)], - "sampling_params": { - "temperature": 0, - "max_new_tokens": 4, - "ignore_eos": True, - }, - "return_logprob": True, - "top_logprobs_num": 5, - }, - timeout=120, - ) - response.raise_for_status() - return response.json()["meta_info"] - - def test_logprob_parity(self): - baseline = popen_launch_server( - self.model, - self.lb_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--trust-remote-code"] + self.baseline_args, - env=SERVER_ENV, - ) - try: - reference = self.generate(self.lb_url) - finally: - kill_process_tree(baseline.pid, wait_timeout=60) - time.sleep(5) - - self.launch_all() - disaggregated = self.generate(self.lb_url) - - reference_logprobs = reference["output_token_logprobs"] - disaggregated_logprobs = disaggregated["output_token_logprobs"] - self.assertEqual( - [item[1] for item in reference_logprobs], - [item[1] for item in disaggregated_logprobs], - ) - self.assertEqual(len(reference_logprobs), 4) - for reference_item, disaggregated_item in zip( - reference_logprobs, disaggregated_logprobs - ): - self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05) - - assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url) - assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url) - assert_process_healthy(self, "decode", self.process_decode, self.decode_url) + baseline_args = UNIFIED_MEMORY_ARGS + extra_prefill_args = UNIFIED_MEMORY_ARGS + extra_decode_args = UNIFIED_MEMORY_ARGS class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregation): @@ -142,10 +57,10 @@ class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregat chunk size so any parity break isolates to the PD transfer. """ - _chunked_args = SERVER_ARGS + ["--chunked-prefill-size", "64"] + _chunked_args = UNIFIED_MEMORY_ARGS + ["--chunked-prefill-size", "64"] + baseline_args = _chunked_args extra_prefill_args = _chunked_args extra_decode_args = _chunked_args - baseline_args = _chunked_args if __name__ == "__main__": diff --git a/test/registered/ep/test_deepep_small.py b/test/registered/ep/test_deepep_small.py index 56df75dc9..79dd30b78 100644 --- a/test/registered/ep/test_deepep_small.py +++ b/test/registered/ep/test_deepep_small.py @@ -16,7 +16,7 @@ from sglang.test.test_utils import ( popen_launch_server, ) -register_cuda_ci(est_time=478, stage="base-c", runner_config="4-gpu-h100") +register_cuda_ci(est_time=270, stage="base-c", runner_config="4-gpu-h100") class TestPureDP(CustomTestCase): @@ -66,51 +66,6 @@ class TestPureDP(CustomTestCase): self.assertGreater(metrics["score"], 0.60) -class TestHybridDPTP(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA - 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", - "--tp", - "4", - "--enable-dp-attention", - "--dp", - "2", - "--moe-a2a-backend", - "deepep", - "--cuda-graph-max-bs-decode", - "128", - "--max-running-requests", - "256", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=200, - num_threads=128, - ) - metrics = run_eval(args) - print(metrics) - - self.assertGreater(metrics["score"], 0.60) - - class TestTP(CustomTestCase): @classmethod def setUpClass(cls): @@ -153,55 +108,6 @@ class TestTP(CustomTestCase): self.assertGreater(metrics["score"], 0.60) -@unittest.skip("covered in test_deepep_large.py") -class TestNoGatherdBuffer(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA - 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", - "--tp", - "4", - "--enable-dp-attention", - "--dp", - "4", - "--moe-dense-tp-size", - "1", - "--enable-dp-lm-head", - "--moe-a2a-backend", - "deepep", - "--cuda-graph-max-bs-decode", - "32", - "--max-running-requests", - "512", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=200, - num_threads=128, - ) - metrics = run_eval(args) - print(metrics) - - self.assertGreater(metrics["score"], 0.60) - - class TestTBO(CustomTestCase): @classmethod def setUpClass(cls): @@ -254,176 +160,6 @@ class TestTBO(CustomTestCase): self.assertGreater(metrics["score"], 0.60) -class TestTBOWithTPAttn(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA - 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", - "--tp", - "4", - "--moe-a2a-backend", - "deepep", - "--enable-two-batch-overlap", - "--cuda-graph-max-bs-decode", - "128", - "--max-running-requests", - "512", - "--mem-fraction-static", # temp fix as DeepEP buffer is too large. - "0.7", - ], - env={ - **os.environ, - "SGLANG_TBO_DEBUG": "1", - }, - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=200, - num_threads=128, - ) - metrics = run_eval(args) - print(metrics) - - self.assertGreater(metrics["score"], 0.60) - - -# There exists bug when using MTP + TBO + attn_tp_size > 1, currently skip that case. -# @unittest.skip("covered in TestMTPWithTPAttnAndTBO") -class TestTBOWithTPAttnAndDenseDP(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA - 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", - "--tp", - "4", - "--moe-dense-tp-size", - "1", - "--moe-a2a-backend", - "deepep", - "--enable-two-batch-overlap", - "--cuda-graph-max-bs-decode", - "128", - "--max-running-requests", - "512", - "--mem-fraction-static", # temp fix as DeepEP buffer is too large. - "0.7", - ], - env={ - **os.environ, - "SGLANG_TBO_DEBUG": "1", - }, - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=200, - num_threads=128, - ) - metrics = run_eval(args) - print(metrics) - - self.assertGreater(metrics["score"], 0.60) - - -@unittest.skip("covered in TestMTPWithTBO") -class TestMTP(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA - 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", - "--tp", - "4", - "--enable-dp-attention", - "--dp", - "2", - "--enable-dp-lm-head", - "--moe-a2a-backend", - "deepep", - "--speculative-algo", - "EAGLE", - "--speculative-draft-model-path", - DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN, - "--speculative-num-steps", - "2", - "--speculative-eagle-topk", - "3", - "--speculative-num-draft-tokens", - "3", - "--cuda-graph-max-bs-decode", - "32", - "--max-running-requests", - "64", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=200, - num_threads=128, - ) - metrics = run_eval(args) - print(metrics) - - self.assertGreater(metrics["score"], 0.60) - - server_info = requests.get(self.base_url + "/server_info") - avg_spec_accept_length = server_info.json()["internal_states"][0][ - "avg_spec_accept_length" - ] - print( - f"###test_gsm8k (deepseek-v3 mtp + dp + tbo):\n" - f"accuracy={metrics['score']=:.3f}\n" - f"{avg_spec_accept_length=:.3f}\n" - ) - self.assertGreater(avg_spec_accept_length, 2.1) - - class TestMTPWithTBO(CustomTestCase): @classmethod def setUpClass(cls): diff --git a/test/registered/ep/test_deepep_small_extra.py b/test/registered/ep/test_deepep_small_extra.py new file mode 100644 index 000000000..763f620eb --- /dev/null +++ b/test/registered/ep/test_deepep_small_extra.py @@ -0,0 +1,167 @@ +import os +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.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST_MLA, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=210, stage="extra-b", runner_config="4-gpu-h100") + + +class TestHybridDPTP(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + 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", + "--tp", + "4", + "--enable-dp-attention", + "--dp", + "2", + "--moe-a2a-backend", + "deepep", + "--cuda-graph-max-bs-decode", + "128", + "--max-running-requests", + "256", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(metrics) + + self.assertGreater(metrics["score"], 0.60) + + +class TestTBOWithTPAttn(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + 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", + "--tp", + "4", + "--moe-a2a-backend", + "deepep", + "--enable-two-batch-overlap", + "--cuda-graph-max-bs-decode", + "128", + "--max-running-requests", + "512", + "--mem-fraction-static", # temp fix as DeepEP buffer is too large. + "0.7", + ], + env={ + **os.environ, + "SGLANG_TBO_DEBUG": "1", + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(metrics) + + self.assertGreater(metrics["score"], 0.60) + + +# There exists bug when using MTP + TBO + attn_tp_size > 1, currently skip that case. +# @unittest.skip("covered in TestMTPWithTPAttnAndTBO") +class TestTBOWithTPAttnAndDenseDP(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + 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", + "--tp", + "4", + "--moe-dense-tp-size", + "1", + "--moe-a2a-backend", + "deepep", + "--enable-two-batch-overlap", + "--cuda-graph-max-bs-decode", + "128", + "--max-running-requests", + "512", + "--mem-fraction-static", # temp fix as DeepEP buffer is too large. + "0.7", + ], + env={ + **os.environ, + "SGLANG_TBO_DEBUG": "1", + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(metrics) + + self.assertGreater(metrics["score"], 0.60) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/hicache/test_hicache_storage_3fs_backend.py b/test/registered/hicache/test_hicache_storage_3fs_backend.py index d064a0d2b..8f1c49384 100644 --- a/test/registered/hicache/test_hicache_storage_3fs_backend.py +++ b/test/registered/hicache/test_hicache_storage_3fs_backend.py @@ -13,7 +13,7 @@ from test_hicache_storage_file_backend import HiCacheStorageBaseMixin from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=300, stage="base-c", runner_config="4-gpu-h100") +register_cuda_ci(est_time=300, stage="base-b", runner_config="2-gpu-large") register_amd_ci(est_time=300, suite="base-b-test-2-gpu-large") diff --git a/test/registered/models_e2e/test_qwen3_next_models.py b/test/registered/models_e2e/test_qwen3_next_models.py index 80953409e..b67736130 100644 --- a/test/registered/models_e2e/test_qwen3_next_models.py +++ b/test/registered/models_e2e/test_qwen3_next_models.py @@ -1,4 +1,3 @@ -import os import unittest from sglang.test.ci.ci_register import register_cuda_ci @@ -7,7 +6,7 @@ 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, stage="base-c", runner_config="4-gpu-h100") +register_cuda_ci(est_time=260, stage="base-c", runner_config="4-gpu-h100") QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" @@ -54,45 +53,5 @@ class TestQwen3NextLazyExtraBufferLargePage( other_args = _make_args(page_size=2, track_interval=2) -class TestQwen3NextLazyExtraBufferAllocFail(KLDivergenceMixin, DefaultServerBase): - model = QWEN3_NEXT_MODEL - cache_chunk_size = 64 - kl_div_thres = 0.002 - other_args = _make_args(page_size=1, track_interval=2) - - @classmethod - def setUpClass(cls): - os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1" - os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1" - super().setUpClass() - - @classmethod - def tearDownClass(cls): - super().tearDownClass() - os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None) - os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None) - - -class TestQwen3NextLazyExtraBufferLargePageAllocFail( - KLDivergenceMixin, DefaultServerBase -): - model = QWEN3_NEXT_MODEL - cache_chunk_size = 64 - kl_div_thres = 0.002 - other_args = _make_args(page_size=2, track_interval=2) - - @classmethod - def setUpClass(cls): - os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1" - os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1" - super().setUpClass() - - @classmethod - def tearDownClass(cls): - super().tearDownClass() - os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None) - os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None) - - if __name__ == "__main__": unittest.main() diff --git a/test/registered/models_e2e/test_qwen3_next_models_extra.py b/test/registered/models_e2e/test_qwen3_next_models_extra.py new file mode 100644 index 000000000..57d997599 --- /dev/null +++ b/test/registered/models_e2e/test_qwen3_next_models_extra.py @@ -0,0 +1,76 @@ +import os +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase + +register_cuda_ci(est_time=250, stage="extra-b", runner_config="4-gpu-h100") + +QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" + +_COMMON_ARGS = [ + "--trust-remote-code", + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "extra_buffer_lazy", + "--attention-backend", + "triton", +] + + +def _make_args(*, page_size=1, track_interval=2): + return [ + *_COMMON_ARGS, + "--mamba-track-interval", + str(track_interval), + "--page-size", + str(page_size), + ] + + +class TestQwen3NextLazyExtraBufferAllocFail(KLDivergenceMixin, DefaultServerBase): + model = QWEN3_NEXT_MODEL + cache_chunk_size = 64 + kl_div_thres = 0.002 + other_args = _make_args(page_size=1, track_interval=2) + + @classmethod + def setUpClass(cls): + os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1" + os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1" + super().setUpClass() + + @classmethod + def tearDownClass(cls): + super().tearDownClass() + os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None) + os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None) + + +class TestQwen3NextLazyExtraBufferLargePageAllocFail( + KLDivergenceMixin, DefaultServerBase +): + model = QWEN3_NEXT_MODEL + cache_chunk_size = 64 + kl_div_thres = 0.002 + other_args = _make_args(page_size=2, track_interval=2) + + @classmethod + def setUpClass(cls): + os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1" + os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1" + super().setUpClass() + + @classmethod + def tearDownClass(cls): + super().tearDownClass() + os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None) + os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/pp/test_pp_gemma4.py b/test/registered/pp/test_pp_gemma4.py new file mode 100644 index 000000000..7a352b1d2 --- /dev/null +++ b/test/registered/pp/test_pp_gemma4.py @@ -0,0 +1,153 @@ +import time +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.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP, + DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + is_in_amd_ci, + is_in_ci, + popen_launch_server, +) + +# tp=1 pp=2 -- two GPUs. +register_cuda_ci(est_time=220, stage="base-b", runner_config="2-gpu-large") + + +@unittest.skipIf( + is_in_amd_ci(), + "Gemma4 PP not yet validated on AMD", +) +class TestGemma4PPAccuracy(unittest.TestCase): + """End-to-end PP=2 accuracy gate for Gemma4 multimodal. + + Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so + sglang auto-selects the triton attention backend; no manual flag needed. + The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an + H100's 80 GB. + """ + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP + cls.base_url = "http://127.0.0.1:23333" + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + 1, + "--pp-size", + 2, + "--trust-remote-code", + "--enable-multimodal", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + # Gemma4 is instruction-tuned and doesn't follow few-shot completion + # prompts well — use the chat API (default in run_eval), which scores + # ~0.98 on this model vs ~0.44 with api="completion". + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + num_examples=200, + num_threads=32, + ) + metrics = run_eval(args) + print(f"{metrics=}") + + # Chat-API baseline ~0.98; gate well below to absorb sample-noise + # without missing a real PP-routing regression (pre-PP-fix the model + # produced garbage outputs scoring ≈ 0). + self.assertGreaterEqual(metrics["score"], 0.90) + # Wait a little bit so that the memory check happens. + time.sleep(4) + + @unittest.skipIf(is_in_ci(), "To reduce the CI execution time.") + def test_mmmu(self): + # Multimodal accuracy gate covering the vision_tower → embed_vision + # (first rank) → PP-proxy handoff → LM tail (last rank) chain. + # Measured 0.71 on 200 examples; full eval (~900 questions) takes + # ~5-7 min on H100 so this is manual-only. + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mmmu", + num_examples=None, + num_threads=32, + ) + metrics = run_eval(args) + print(f"{metrics=}") + # Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in + # 0.69-0.73. Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary + # samples ≈ 0.015) while still catching mid-grade vision/PP + # regressions, not just complete breakage. + self.assertGreater(metrics["score"], 0.65) + + +@unittest.skipIf( + is_in_amd_ci(), + "Gemma4 PP not yet validated on AMD", +) +class TestGemma4PLEPPAccuracy(unittest.TestCase): + """PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path). + + 26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP + test never crosses the PLE branch. Cuda graph + PLE corrupts outputs + (the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema + drops ``per_layer_inputs``), so this test pins the eager configuration. + """ + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP + cls.base_url = "http://127.0.0.1:23339" + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + 1, + "--pp-size", + 2, + "--trust-remote-code", + "--enable-multimodal", + # Required for PLE under PP — see Gemma4TextModel guard. + "--disable-cuda-graph", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + # Eager-path baseline ~0.92; gate 0.80 catches PLE breakage + # (corruption collapses score to ~0). + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + num_examples=100, + num_threads=32, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], 0.80) + time.sleep(4) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/pp/test_pp_single_node.py b/test/registered/pp/test_pp_single_node.py index a807d33d0..8ca9c9cb2 100644 --- a/test/registered/pp/test_pp_single_node.py +++ b/test/registered/pp/test_pp_single_node.py @@ -23,18 +23,15 @@ from sglang.test.run_eval import run_eval from sglang.test.test_utils import ( DEFAULT_MLA_MODEL_NAME_FOR_TEST, DEFAULT_MODEL_NAME_FOR_TEST, - DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP, - DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, is_in_amd_ci, - is_in_ci, popen_launch_server, run_bench_one_batch_server, ) -register_cuda_ci(est_time=500, stage="base-c", runner_config="4-gpu-h100") +register_cuda_ci(est_time=280, stage="base-c", runner_config="4-gpu-h100") register_amd_ci(est_time=500, suite="stage-c-test-4-gpu-amd") @@ -145,137 +142,6 @@ class TestDPAttentionDP2PP2(CustomTestCase): self.assertGreater(metrics["score"], 0.8) -@unittest.skipIf( - is_in_amd_ci(), - "Gemma4 PP not yet validated on AMD", -) -class TestGemma4PPAccuracy(unittest.TestCase): - """End-to-end PP=2 accuracy gate for Gemma4 multimodal. - - Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so - sglang auto-selects the triton attention backend; no manual flag needed. - The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an - H100's 80 GB. - """ - - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP - cls.base_url = "http://127.0.0.1:23333" - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--tp-size", - 1, - "--pp-size", - 2, - "--trust-remote-code", - "--enable-multimodal", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - # Gemma4 is instruction-tuned and doesn't follow few-shot completion - # prompts well — use the chat API (default in run_eval), which scores - # ~0.98 on this model vs ~0.44 with api="completion". - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - num_examples=200, - num_threads=32, - ) - metrics = run_eval(args) - print(f"{metrics=}") - - # Chat-API baseline ~0.98; gate well below to absorb sample-noise - # without missing a real PP-routing regression (pre-PP-fix the model - # produced garbage outputs scoring ≈ 0). - self.assertGreaterEqual(metrics["score"], 0.90) - # Wait a little bit so that the memory check happens. - time.sleep(4) - - @unittest.skipIf(is_in_ci(), "To reduce the CI execution time.") - def test_mmmu(self): - # Multimodal accuracy gate covering the vision_tower → embed_vision - # (first rank) → PP-proxy handoff → LM tail (last rank) chain. - # Measured 0.71 on 200 examples; full eval (~900 questions) takes - # ~5-7 min on H100 so this is manual-only. - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="mmmu", - num_examples=None, - num_threads=32, - ) - metrics = run_eval(args) - print(f"{metrics=}") - # Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in - # 0.69-0.73. Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary - # samples ≈ 0.015) while still catching mid-grade vision/PP - # regressions, not just complete breakage. - self.assertGreater(metrics["score"], 0.65) - - -@unittest.skipIf( - is_in_amd_ci(), - "Gemma4 PP not yet validated on AMD", -) -class TestGemma4PLEPPAccuracy(unittest.TestCase): - """PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path). - - 26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP - test never crosses the PLE branch. Cuda graph + PLE corrupts outputs - (the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema - drops ``per_layer_inputs``), so this test pins the eager configuration. - """ - - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP - cls.base_url = "http://127.0.0.1:23339" - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--tp-size", - 1, - "--pp-size", - 2, - "--trust-remote-code", - "--enable-multimodal", - # Required for PLE under PP — see Gemma4TextModel guard. - "--disable-cuda-graph", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - # Eager-path baseline ~0.92; gate 0.80 catches PLE breakage - # (corruption collapses score to ~0). - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - num_examples=100, - num_threads=32, - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreaterEqual(metrics["score"], 0.80) - time.sleep(4) - - class TestPPMixedChunk(CustomTestCase): @classmethod def setUpClass(cls): diff --git a/test/registered/rl/test_release_memory_occupation.py b/test/registered/rl/test_release_memory_occupation.py deleted file mode 100644 index 8628114b1..000000000 --- a/test/registered/rl/test_release_memory_occupation.py +++ /dev/null @@ -1,477 +0,0 @@ -"""Test memory release and resume operations for SGLang engine in hybrid RL training. - -This test suite evaluates the SGLang engine's memory management capabilities, focusing -on releasing and resuming memory occupation for KV cache and model weights. It simulates -an RL workflow where the SGLang engine acts as a rollout engine for experience collection. -The process involves initializing the engine, sending a small number of requests to simulate -rollout, releasing memory to mimic offloading during RL training, resuming memory occupation, -updating weights with a trained HuggingFace model, and verifying the updated weights. - -Detailed in our proposal (https://github.com/sgl-project/sglang/pull/7099), two test cases -are included: - -1. Basic Release and Resume: Uses a lower mem_fraction_static (0.6) to control memory allocation -and avoid OOM errors carefully. This test simulates a scenario without multi-stage memory management, -ensuring the engine can release and resume memory occupation while maintaining functionality after -weight updates. - -2. Multi-Stage Release and Resume: Employs a higher mem_fraction_static (0.85) to simulate higher -memory pressure, leveraging multi-stage memory management. It sequentially releases and resumes -KV cache and model weights, verifying memory deallocation and reallocation at each stage, and -ensuring correct weight updates and text generation. - -3. Tensor Parallel Tests: Tests memory release and resume operations with different tensor parallel -configurations (tp=1, tp=2) to ensure proper memory management in distributed settings. For different -data parallel size, we test it in verl. - -NOTE: This test is temporarily disabled. -""" - -import os -import time -import unittest - -from transformers import AutoModelForCausalLM - -import sglang as sgl -from sglang.srt.constants import ( - GPU_MEMORY_TYPE_CUDA_GRAPH, - GPU_MEMORY_TYPE_KV_CACHE, - GPU_MEMORY_TYPE_WEIGHTS, -) -from sglang.srt.utils import get_device -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.test_utils import ( - DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST, - DEFAULT_SMALL_MODEL_NAME_FOR_TEST, - DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE, - DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE, - DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT, - CustomTestCase, - empty_gpu_cache, - get_gpu_count, - get_gpu_memory_gb, -) - -register_cuda_ci( - est_time=200, - stage="base-c", - runner_config="4-gpu-h100", - disabled="Temporarily disabled - needs investigation", -) - -# (temporarily) set to true to observe memory usage in nvidia-smi more clearly -_DEBUG_EXTRA = False - - -class TestReleaseMemoryOccupation(CustomTestCase): - def _setup_engine( - self, - model_name, - mem_fraction_static=0.8, - tp_size=1, - ep_size=1, - enable_weights_cpu_backup=False, - ): - """Common setup for engine and HF model.""" - - os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1" - engine = sgl.Engine( - model_path=model_name, - random_seed=42, - enable_memory_saver=True, - mem_fraction_static=mem_fraction_static, - tp_size=tp_size, - ep_size=ep_size, - enable_weights_cpu_backup=enable_weights_cpu_backup, - # disable_cuda_graph=True, # for debugging only - ) - - return engine - - def _common_test_params(self): - """Common test parameters.""" - return { - "prompt": "Today is a sunny day and I like", - "sampling_params": {"temperature": 0, "max_new_tokens": 8}, - "expect_output_before_update_weights": " to spend it outdoors. I decided to", - "expect_output_after_update_weights": " to go for a walk. I like", - "prompt_moe": "The weather is nice today, and I want to", - "sampling_params_moe": {"temperature": 0, "max_new_tokens": 16}, - "expect_output_before_update_weights_moe": " go to the park. I have a picnic basket, a book, and a", - "expect_output_after_update_weights_moe": " go to the park. I have a lot of things to do, but I", - "prompt_hybrid_mamba": "The weather is nice today, and I want to", - "sampling_params_hybrid_mamba": {"temperature": 0, "max_new_tokens": 16}, - "expect_output_before_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can", - "expect_output_after_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can", - } - - def _test_initial_generation( - self, engine, prompt, sampling_params, expect_output_before_update_weights - ): - """Test initial generation and memory allocation.""" - print("generate (#1)") - outputs = engine.generate(prompt, sampling_params)["text"] - self.assertEqual(outputs, expect_output_before_update_weights) - - if _DEBUG_EXTRA: - time.sleep(3) - - def test_release_and_resume_occupation(self): - # Without multi-stage release and resume, we need to carefully control the memory fraction to avoid OOM - model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST - assert get_gpu_count() >= 2, "Need at least 2 GPUs for tensor parallel tests" - - for tp_size in [1, 2]: - - print(f"Testing tp_size={tp_size} for test_release_and_resume_occupation") - engine = self._setup_engine( - model_name=model_name, mem_fraction_static=0.6, tp_size=tp_size - ) - params = self._common_test_params() - - self._test_initial_generation( - engine, - params["prompt"], - params["sampling_params"], - params["expect_output_before_update_weights"], - ) - - t = time.perf_counter() - gpu_memory_usage_before_release = get_gpu_memory_gb() - engine.release_memory_occupation() - gpu_memory_usage_after_release = get_gpu_memory_gb() - - self.assertLess( - gpu_memory_usage_after_release, - gpu_memory_usage_before_release, - ) - - print( - f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB" - ) - - if _DEBUG_EXTRA: - time.sleep(3) - - t = time.perf_counter() - engine.resume_memory_occupation() - print( - f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB" - ) - - hf_model_new = AutoModelForCausalLM.from_pretrained( - DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE, - torch_dtype="bfloat16", - device_map=get_device(), - ) - engine.update_weights_from_tensor(list(hf_model_new.named_parameters())) - - # destroy the hf model - del hf_model_new - empty_gpu_cache() - - print("generate (#2)") - outputs = engine.generate(params["prompt"], params["sampling_params"])[ - "text" - ] - self.assertEqual(outputs, params["expect_output_after_update_weights"]) - engine.shutdown() - - def test_release_and_resume_occupation_with_weights_cpu_backup(self): - # Test release and resume occupation with weights CPU backup - model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST - - print("Testing test_release_and_resume_occupation_with_weights_cpu_backup") - engine = self._setup_engine( - model_name=model_name, - mem_fraction_static=0.6, - enable_weights_cpu_backup=True, - ) - params = self._common_test_params() - - self._test_initial_generation( - engine, - params["prompt"], - params["sampling_params"], - params["expect_output_before_update_weights"], - ) - - t = time.perf_counter() - gpu_memory_usage_before_release = get_gpu_memory_gb() - engine.release_memory_occupation() - gpu_memory_usage_after_release = get_gpu_memory_gb() - - self.assertLess( - gpu_memory_usage_after_release, - gpu_memory_usage_before_release, - ) - - print( - f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB" - ) - - if _DEBUG_EXTRA: - time.sleep(3) - - t = time.perf_counter() - engine.resume_memory_occupation() - print( - f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB" - ) - - print("generate post resume") - outputs = engine.generate(params["prompt"], params["sampling_params"])["text"] - self.assertEqual(outputs, params["expect_output_before_update_weights"]) - engine.shutdown() - - def test_multi_stage_release_and_resume(self): - # With multi-stage release and resume, we can set the memory fraction to 0.85 without concern of OOM - model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST - - for tp_size in [1, 2]: - if tp_size == 2 and get_gpu_count() < 2: - continue - - print(f"Testing tp_size={tp_size} for test_multi_stage_release_and_resume") - os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1" - engine = sgl.Engine( - model_path=model_name, - random_seed=42, - enable_memory_saver=True, - mem_fraction_static=0.85, # Higher memory pressure - tp_size=tp_size, - ) - params = self._common_test_params() - - self._test_initial_generation( - engine, - params["prompt"], - params["sampling_params"], - params["expect_output_before_update_weights"], - ) - - t = time.perf_counter() - gpu_memory_usage_before_release = get_gpu_memory_gb() - engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE]) - - gpu_memory_usage_after_release_kv_cache = get_gpu_memory_gb() - - self.assertLess( - gpu_memory_usage_after_release_kv_cache, - gpu_memory_usage_before_release, - ) - - engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS]) - gpu_memory_usage_after_release_weights = get_gpu_memory_gb() - - self.assertLess( - gpu_memory_usage_after_release_weights, - gpu_memory_usage_after_release_kv_cache, - ) - - engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]) - gpu_memory_usage_after_release_cuda_graph = get_gpu_memory_gb() - - self.assertLess( - gpu_memory_usage_after_release_cuda_graph, - gpu_memory_usage_after_release_weights, - ) - - print(f"Release took {time.perf_counter() - t:.2f}s") - print( - f"Memory: {gpu_memory_usage_before_release:.1f} → {gpu_memory_usage_after_release_kv_cache:.1f} → {gpu_memory_usage_after_release_weights:.1f} → {gpu_memory_usage_after_release_cuda_graph:.1f} GB" - ) - - if _DEBUG_EXTRA: - time.sleep(3) - - t = time.perf_counter() - gpu_memory_usage_before_resume = get_gpu_memory_gb() - - # gpu_memory_usage_after_release_weights and gpu_memory_usage_before_resume should be close - - self.assertAlmostEqual( - gpu_memory_usage_after_release_weights, - gpu_memory_usage_before_resume, - delta=3.0, - ) - print(f"Resume weights took {time.perf_counter() - t:.2f}s") - - engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]) - gpu_memory_usage_after_resume_cuda_graph = get_gpu_memory_gb() - - self.assertGreater( - gpu_memory_usage_after_resume_cuda_graph, - gpu_memory_usage_before_resume, - ) - - engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS]) - gpu_memory_usage_after_resume_weights = get_gpu_memory_gb() - - self.assertGreater( - gpu_memory_usage_after_resume_weights, - gpu_memory_usage_after_resume_cuda_graph, - ) - - # Update weights from a trained model to serving engine, and then destroy the trained model - hf_model_new = AutoModelForCausalLM.from_pretrained( - DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE, - torch_dtype="bfloat16", - device_map=get_device(), - ) - gpu_memory_usage_after_loaded_hf_model = get_gpu_memory_gb() - engine.update_weights_from_tensor(list(hf_model_new.named_parameters())) - - # destroy the hf model - del hf_model_new - empty_gpu_cache() - engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE]) - - gpu_memory_usage_after_resume_kv_cache = get_gpu_memory_gb() - self.assertGreater( - gpu_memory_usage_after_resume_kv_cache, - gpu_memory_usage_after_resume_weights, - ) - - print(f"Resume + update took {time.perf_counter() - t:.2f}s") - print( - f"Memory: {gpu_memory_usage_before_resume:.1f} → {gpu_memory_usage_after_resume_cuda_graph:.1f} → {gpu_memory_usage_after_resume_weights:.1f} → {gpu_memory_usage_after_loaded_hf_model:.1f} → {gpu_memory_usage_after_resume_kv_cache:.1f} GB" - ) - - print("generate (#2)") - outputs = engine.generate(params["prompt"], params["sampling_params"])[ - "text" - ] - self.assertEqual(outputs, params["expect_output_after_update_weights"]) - engine.shutdown() - - def test_moe_model_release_and_resume(self): - # Test with MoE model - model_name = DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT - - tp_size = ep_size = 2 - - print( - f"Testing tp_size={tp_size} and ep_size={ep_size} for test_moe_model_release_and_resume" - ) - engine = sgl.Engine( - model_path=model_name, - random_seed=42, - enable_memory_saver=True, - mem_fraction_static=0.5, - tp_size=tp_size, - ep_size=ep_size, - ) - params = self._common_test_params() - - self._test_initial_generation( - engine, - params["prompt_moe"], - params["sampling_params_moe"], - params["expect_output_before_update_weights_moe"], - ) - - t = time.perf_counter() - gpu_memory_usage_before_release = get_gpu_memory_gb() - engine.release_memory_occupation() - gpu_memory_usage_after_release = get_gpu_memory_gb() - self.assertLess( - gpu_memory_usage_after_release, - gpu_memory_usage_before_release, - ) - - print( - f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB" - ) - - if _DEBUG_EXTRA: - time.sleep(3) - - t = time.perf_counter() - engine.resume_memory_occupation() - print( - f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB" - ) - - hf_model_new = AutoModelForCausalLM.from_pretrained( - DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE, - torch_dtype="bfloat16", - device_map=get_device(), - ) - engine.update_weights_from_tensor(list(hf_model_new.named_parameters())) - - # destroy the hf model - del hf_model_new - empty_gpu_cache() - - print("generate (#2)") - outputs = engine.generate(params["prompt_moe"], params["sampling_params_moe"])[ - "text" - ] - self.assertEqual(outputs, params["expect_output_after_update_weights_moe"]) - engine.shutdown() - - def test_hybrid_mamba_model_release_and_resume(self): - # Test with Hybrid Mamba model - model_name = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST - - tp_size = 4 - - print( - f"Testing tp_size={tp_size} for test_hybrid_mamba_model_release_and_resume" - ) - engine = sgl.Engine( - model_path=model_name, - random_seed=42, - enable_memory_saver=True, - tp_size=tp_size, - ) - params = self._common_test_params() - - self._test_initial_generation( - engine, - params["prompt_hybrid_mamba"], - params["sampling_params_hybrid_mamba"], - params["expect_output_before_update_weights_hybrid_mamba"], - ) - - t = time.perf_counter() - gpu_memory_usage_before_release = get_gpu_memory_gb() - engine.release_memory_occupation() - gpu_memory_usage_after_release = get_gpu_memory_gb() - self.assertLess( - gpu_memory_usage_after_release, - gpu_memory_usage_before_release, - ) - - print( - f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB" - ) - - if _DEBUG_EXTRA: - time.sleep(3) - - t = time.perf_counter() - engine.resume_memory_occupation() - print( - f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB" - ) - - engine.update_weights_from_disk(model_name) - - # destroy the hf model - empty_gpu_cache() - - print("generate (#2)") - outputs = engine.generate( - params["prompt_hybrid_mamba"], params["sampling_params_hybrid_mamba"] - )["text"] - self.assertEqual( - outputs, params["expect_output_after_update_weights_hybrid_mamba"] - ) - engine.shutdown() - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/layers/test_flashinfer_comm_fusion.py b/test/registered/unit/layers/test_flashinfer_comm_fusion.py index 28ccf1f82..9eaf5c8fc 100644 --- a/test/registered/unit/layers/test_flashinfer_comm_fusion.py +++ b/test/registered/unit/layers/test_flashinfer_comm_fusion.py @@ -10,9 +10,9 @@ from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-h100") -register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-b200") -register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-gb300") +# Collectives are mocked and world_size is a plain int, so the world_size=4 +# cases need one real CUDA device. +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") class _FakeWorkspace: