[MLX] Add correctness tests for qwen2_moe and qwen3_moe (#29440)

This commit is contained in:
Siming Deng
2026-07-06 20:32:12 -07:00
committed by GitHub
parent df06e03662
commit 4145e595cf
3 changed files with 499 additions and 0 deletions
@@ -0,0 +1,122 @@
import importlib.util
import os
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
try_cached_model,
)
# Registered on the CPU suite but skipped wherever mlx is absent; runs for real
# only on Apple Silicon. The macOS CI lane (pr-test-mlx.yml) is model-free, so
# this serving test is not wired into it and still runs only locally.
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
_HAS_MLX = importlib.util.find_spec("mlx") is not None
# qwen2_moe architecture (Qwen2MoeForCausalLM), served on the MLX backend.
# The model runs through mlx_lm's own qwen2_moe implementation; the SGLang MLX
# backend does not require any srt/models file for it. This test is a black-box
# correctness guard for the served model.
#
# Default is the MLX-community 4-bit repo so the test is portable. Override with
# SGLANG_MLX_TEST_MODEL to point at a local copy, e.g.
# SGLANG_MLX_TEST_MODEL=models/Qwen1.5-MoE-A2.7B-Chat-4bit
MODEL_PATH = os.environ.get(
"SGLANG_MLX_TEST_MODEL", "mlx-community/Qwen1.5-MoE-A2.7B-Chat-4bit"
)
# mem-fraction is tuned conservatively for a 24 GB Apple Silicon machine.
MEM_FRACTION_STATIC = os.environ.get("SGLANG_MLX_TEST_MEM_FRACTION", "0.7")
@unittest.skipUnless(_HAS_MLX, "requires mlx (Apple Silicon only)")
class TestQwen2MoeMlxCorrectness(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL_PATH)
cls.base_url = DEFAULT_URL_FOR_TEST
env = os.environ.copy()
env["SGLANG_USE_MLX"] = "1"
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp-size",
"1",
"--disable-radix-cache",
"--disable-cuda-graph",
"--mem-fraction-static",
MEM_FRACTION_STATIC,
"--max-running-requests",
"1",
"--context-length",
"2048",
],
env=env,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process is not None:
kill_process_tree(cls.process.pid)
def _chat(self, messages, max_tokens=32, temperature=0):
resp = requests.post(
f"{self.base_url}/v1/chat/completions",
json={
"model": MODEL_PATH,
"messages": messages,
"temperature": temperature,
"max_tokens": max_tokens,
},
timeout=120,
)
resp.raise_for_status()
return resp.json()["choices"][0]["message"]["content"].strip()
def test_basic_generation_nonempty(self):
text = self._chat(
[
{"role": "system", "content": "You are a concise assistant."},
{"role": "user", "content": "Say hello briefly."},
],
max_tokens=16,
)
self.assertIsInstance(text, str)
self.assertGreater(len(text), 0)
def test_simple_arithmetic(self):
text = self._chat(
[
{"role": "system", "content": "You are a concise assistant."},
{"role": "user", "content": "What is 2+2? Reply with just the number."},
],
max_tokens=8,
)
self.assertIn("4", text)
def test_simple_fact(self):
text = self._chat(
[
{"role": "system", "content": "You are a concise assistant."},
{"role": "user", "content": "What is the capital of France? One word."},
],
max_tokens=8,
)
self.assertIn("Paris", text)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,122 @@
import importlib.util
import os
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
try_cached_model,
)
# Registered on the CPU suite but skipped wherever mlx is absent; runs for real
# only on Apple Silicon. The macOS CI lane (pr-test-mlx.yml) is model-free, so
# this serving test is not wired into it and still runs only locally.
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
_HAS_MLX = importlib.util.find_spec("mlx") is not None
# qwen3_moe architecture (Qwen3MoeForCausalLM), served on the MLX backend.
# The model runs through mlx_lm's own qwen3_moe implementation; the SGLang MLX
# backend does not require any srt/models file for it. This test is a black-box
# correctness guard for the served model. Qwen3 is a hybrid-thinking model, so
# thinking is disabled to keep outputs short and deterministic.
#
# Default is the MLX-community 4-bit repo so the test is portable. Override with
# SGLANG_MLX_TEST_MODEL to point at a local copy, e.g.
# SGLANG_MLX_TEST_MODEL=models/Qwen3-30B-A3B-4bit
MODEL_PATH = os.environ.get("SGLANG_MLX_TEST_MODEL", "mlx-community/Qwen3-30B-A3B-4bit")
# mem-fraction is tuned conservatively for a 24 GB Apple Silicon machine.
MEM_FRACTION_STATIC = os.environ.get("SGLANG_MLX_TEST_MEM_FRACTION", "0.9")
@unittest.skipUnless(_HAS_MLX, "requires mlx (Apple Silicon only)")
class TestQwen3MoeMlxCorrectness(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL_PATH)
cls.base_url = DEFAULT_URL_FOR_TEST
env = os.environ.copy()
env["SGLANG_USE_MLX"] = "1"
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp-size",
"1",
"--disable-radix-cache",
"--disable-cuda-graph",
"--mem-fraction-static",
MEM_FRACTION_STATIC,
"--max-running-requests",
"1",
"--context-length",
"2048",
],
env=env,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process is not None:
kill_process_tree(cls.process.pid)
def _chat(self, messages, max_tokens=32, temperature=0):
resp = requests.post(
f"{self.base_url}/v1/chat/completions",
json={
"model": MODEL_PATH,
"messages": messages,
"temperature": temperature,
"max_tokens": max_tokens,
"chat_template_kwargs": {"enable_thinking": False},
},
timeout=120,
)
resp.raise_for_status()
return resp.json()["choices"][0]["message"]["content"].strip()
def test_basic_generation_nonempty(self):
text = self._chat(
[
{"role": "system", "content": "You are a concise assistant."},
{"role": "user", "content": "Say hello briefly."},
],
max_tokens=16,
)
self.assertIsInstance(text, str)
self.assertGreater(len(text), 0)
def test_simple_arithmetic(self):
text = self._chat(
[
{"role": "system", "content": "You are a concise assistant."},
{"role": "user", "content": "What is 2+2? Reply with just the number."},
],
max_tokens=8,
)
self.assertIn("4", text)
def test_simple_fact(self):
text = self._chat(
[
{"role": "system", "content": "You are a concise assistant."},
{"role": "user", "content": "What is the capital of France? One word."},
],
max_tokens=8,
)
self.assertIn("Paris", text)
if __name__ == "__main__":
unittest.main()