diff --git a/test/registered/distributed/test_disaggregation_dsv4.py b/test/registered/distributed/test_disaggregation_dsv4.py new file mode 100644 index 000000000..d36844e2d --- /dev/null +++ b/test/registered/distributed/test_disaggregation_dsv4.py @@ -0,0 +1,140 @@ +import unittest +from types import SimpleNamespace + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + popen_launch_pd_server, + try_cached_model, +) + +register_cuda_ci(est_time=250, suite="stage-c-test-dsv4-8-gpu-h200") + +DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8" + +DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}' + +DSV4_FLASH_ENV = { + "SGLANG_DSV4_FP4_EXPERTS": "0", + "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", +} + +_EAGLE_SPEC_ARGS = [ + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "1", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "2", +] + + +class TestDisaggregationDSV4(PDDisaggregationServerBase): + @classmethod + def setUpClass(cls): + super().setUpClass() + + cls.model = try_cached_model(DSV4_FLASH_MODEL) + + cls.start_prefill() + cls.start_decode() + + cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill) + cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode) + + cls.launch_lb() + + @classmethod + def start_prefill(cls): + prefill_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "prefill", + "--disaggregation-bootstrap-port", + cls.bootstrap_port, + "--tp", + 4, + "--dp", + 4, + "--enable-dp-attention", + "--moe-a2a-backend", + "deepep", + "--deepep-config", + DEEPEP_CONFIG, + "--cuda-graph-max-bs", + "128", + "--max-running-requests", + "256", + "--mem-fraction-static", + "0.7", + *_EAGLE_SPEC_ARGS, + ] + prefill_args += cls.transfer_backend + cls.rdma_devices + cls.process_prefill = popen_launch_pd_server( + cls.model, + cls.prefill_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=prefill_args, + env=DSV4_FLASH_ENV, + ) + + @classmethod + def start_decode(cls): + decode_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "decode", + "--disaggregation-bootstrap-port", + cls.bootstrap_port, + "--tp", + 4, + "--dp", + 4, + "--enable-dp-attention", + "--base-gpu-id", + 4, + "--moe-a2a-backend", + "deepep", + "--deepep-config", + DEEPEP_CONFIG, + "--cuda-graph-max-bs", + "128", + "--max-running-requests", + "256", + "--mem-fraction-static", + "0.7", + *_EAGLE_SPEC_ARGS, + ] + decode_args += cls.transfer_backend + cls.rdma_devices + cls.process_decode = popen_launch_pd_server( + cls.model, + cls.decode_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=decode_args, + env=DSV4_FLASH_ENV, + ) + + 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(f"Evaluation metrics: {metrics}") + + self.assertGreater(metrics["score"], 0.95) + + +if __name__ == "__main__": + unittest.main()