From 45ef8344128da4cb2f88b00afadf8bdc1f51092e Mon Sep 17 00:00:00 2001 From: Jonah Bernard <96398205+Jonahcb@users.noreply.github.com> Date: Mon, 5 Jan 2026 11:26:11 -0800 Subject: [PATCH] Add MoE Integration Tests For CUTLASS Coverage (#16280) --- ...oe_runners.py => test_moe_runners_1gpu.py} | 26 +++- .../layers/moe/test_moe_runners_4gpu.py | 116 ++++++++++++++++++ 2 files changed, 139 insertions(+), 3 deletions(-) rename test/manual/layers/moe/{test_moe_runners.py => test_moe_runners_1gpu.py} (88%) create mode 100644 test/manual/layers/moe/test_moe_runners_4gpu.py diff --git a/test/manual/layers/moe/test_moe_runners.py b/test/manual/layers/moe/test_moe_runners_1gpu.py similarity index 88% rename from test/manual/layers/moe/test_moe_runners.py rename to test/manual/layers/moe/test_moe_runners_1gpu.py index 4b4047d4b..158e48aa6 100644 --- a/test/manual/layers/moe/test_moe_runners.py +++ b/test/manual/layers/moe/test_moe_runners_1gpu.py @@ -1,3 +1,4 @@ +import os import unittest from types import SimpleNamespace @@ -8,7 +9,6 @@ from sglang.test.test_utils import ( DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4, DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE, DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT, - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, popen_launch_server, @@ -17,7 +17,7 @@ from sglang.test.test_utils import ( class TestMoERunner(CustomTestCase): BASE_URL = DEFAULT_URL_FOR_TEST - TIMEOUT = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + TIMEOUT = 6000 DEFAULT_EVAL_KWARGS = { "eval_name": "mmlu", "num_examples": 5, @@ -131,6 +131,20 @@ class TestMoERunner(CustomTestCase): "pytorch", ], }, + "moe_runner_cutlass_fp8": { + "model": DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE, + "timeout": 3600, + "other_args": [ + "--trust-remote-code", + "--moe-runner-backend", + "cutlass", + "--attention-backend", + "triton", + "--sampling-backend", + "pytorch", + "--disable-cuda-graph", + ], + }, "moe_runner_speculative": { "model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT, "other_args": [ @@ -159,12 +173,18 @@ class TestMoERunner(CustomTestCase): model = config["model"] other_args = config.get("other_args", []) eval_kwargs = self.DEFAULT_EVAL_KWARGS + env = dict(os.environ) + env["SGLANG_ENABLE_JIT_DEEPGEMM"] = "1" + env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0" + env.update(config.get("env_overrides", {})) + timeout = config.get("timeout", self.TIMEOUT) process = popen_launch_server( model, self.BASE_URL, - timeout=self.TIMEOUT, + timeout=timeout, other_args=other_args, + env=env, ) try: args = SimpleNamespace( diff --git a/test/manual/layers/moe/test_moe_runners_4gpu.py b/test/manual/layers/moe/test_moe_runners_4gpu.py new file mode 100644 index 000000000..58a218874 --- /dev/null +++ b/test/manual/layers/moe/test_moe_runners_4gpu.py @@ -0,0 +1,116 @@ +import os +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestMoERunner4GPU(CustomTestCase): + BASE_URL = DEFAULT_URL_FOR_TEST + TIMEOUT = 6000 + DEFAULT_EVAL_KWARGS = { + "eval_name": "mmlu", + "num_examples": 5, + "num_threads": 1, + } + + CONFIGS = { + "moe_runner_cutlass_w4a8": { + "model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8", # FP8 W8A8 MoE model + "other_args": [ + "--trust-remote-code", + "--moe-runner-backend", + "cutlass", + "--attention-backend", + "triton", + "--sampling-backend", + "pytorch", + "--tp-size", + "4", + ], + }, + "moe_runner_cutlass_w4a8_deepep_normal": { + "model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8", # FP8 W8A8 MoE model + "other_args": [ + "--trust-remote-code", + "--moe-runner-backend", + "cutlass", + "--moe-a2a-backend", + "deepep", + "--deepep-mode", + "normal", + "--attention-backend", + "triton", + "--sampling-backend", + "pytorch", + "--tp-size", + "4", + ], + }, + "moe_runner_cutlass_w4a8_deepep_ll": { + "model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8", # FP8 W8A8 MoE model + "env_overrides": {"SGLANG_DEEPEP_BF16_DISPATCH": "1"}, + "other_args": [ + "--trust-remote-code", + "--moe-runner-backend", + "cutlass", + "--moe-a2a-backend", + "deepep", + "--deepep-mode", + "low_latency", + "--attention-backend", + "triton", + "--sampling-backend", + "pytorch", + "--tp-size", + "4", + ], + }, + } + + def _run_config(self, config: dict) -> None: + model = config["model"] + other_args = config.get("other_args", []) + eval_kwargs = self.DEFAULT_EVAL_KWARGS + env = dict(os.environ) + env["SGLANG_ENABLE_JIT_DEEPGEMM"] = "1" + env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0" + env.update(config.get("env_overrides", {})) + timeout = config.get("timeout", self.TIMEOUT) + + process = popen_launch_server( + model, + self.BASE_URL, + timeout=timeout, + other_args=other_args, + env=env, + ) + try: + args = SimpleNamespace( + base_url=self.BASE_URL, + model=model, + **eval_kwargs, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], 0.48) + finally: + kill_process_tree(process.pid) + + +for _name, _cfg in TestMoERunner4GPU.CONFIGS.items(): + setattr( + TestMoERunner4GPU, + f"test_{_name}", + (lambda self, cfg=_cfg: self._run_config(cfg)), + ) + + +if __name__ == "__main__": + unittest.main()