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

test/registered/mlx/models_e2e/test_qwen2_moe_mlx_correctness.py

@@ -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()

test/registered/mlx/models_e2e/test_qwen3_moe_mlx_correctness.py

@@ -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()

test/registered/unit/hardware_backend/mlx/test_mlx_reference_correctness.py

@@ -0,0 +1,255 @@
+"""Reference-equivalence correctness for the MLX MoE execution path.
+
+The existing ``test/registered/models/test_qwen{2,3}_moe_mlx_correctness.py`` are
+black-box smoke tests: they only check that a served model emits *plausible* text
+("Paris" appears, "4" appears). A subtly broken KV cache, slot allocator, or
+prefill-chunking change can still pass those.
+
+This is the strong guard. The SGLang MLX backend executes models by wrapping
+``mlx_lm`` (``mlx_lm.load`` + the model's own forward, with attention patched for
+SGLang's cache) and decodes greedily (``mx.argmax``). So the ground truth for "did
+SGLang corrupt the output?" is raw, *unpatched* ``mlx_lm`` greedy generation on the
+same prompt. We record those reference tokens, then drive ``MlxModelRunner``
+(prefill + decode_batch) and assert the tokens match exactly, up to and including
+EOS. Empirically the agreement is exact (not approximate), so any divergence here
+is a real regression.
+
+A second test pins batching isolation: a prompt decoded inside a multi-request
+``decode_batch`` must yield the same tokens as when decoded alone -- guarding the
+slot/cache bookkeeping that single-request black-box tests never touch.
+
+Memory safety: the runner patches the model in-place, so the reference needs its
+own *separate* ``mlx_lm.load``. Loading both copies at once doubles resident weight
+memory (~2x model) and can trigger an unrecoverable Metal command-buffer OOM that
+hard-reboots a memory-constrained Mac. So we load the reference, record tokens,
+fully release it (``del`` + ``gc`` + ``mx.clear_cache``), and only then build the
+runner -- keeping peak at a single model copy. A pre-flight free-memory check skips
+(never crashes) when there isn't headroom for even one copy.
+
+MLX-gated like its siblings: registered on the CPU suite but skipped wherever the
+``mlx`` package is absent (all current CI runners), so it runs for real only on
+Apple Silicon. Override the model with ``SGLANG_MLX_TEST_MODEL`` (e.g. a local path
+or ``mlx-community/Qwen3-30B-A3B-4bit`` for qwen3_moe); bump
+``SGLANG_MLX_TEST_MIN_FREE_GB`` accordingly for larger models.
+"""
+
+from __future__ import annotations
+
+import gc
+import importlib.util
+import os
+import unittest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+_HAS_MLX = (
+    importlib.util.find_spec("mlx") is not None
+    and importlib.util.find_spec("mlx_lm") is not None
+)
+_SKIP_REASON = "requires mlx + mlx_lm (Apple Silicon only)"
+
+# Default to the portable mlx-community 4-bit qwen2_moe repo; override to a local
+# dir or another MoE arch (e.g. Qwen3-30B-A3B-4bit) via the env var.
+MODEL_PATH = os.environ.get(
+    "SGLANG_MLX_TEST_MODEL", "mlx-community/Qwen1.5-MoE-A2.7B-Chat-4bit"
+)
+MEM_FRACTION_STATIC = float(os.environ.get("SGLANG_MLX_TEST_MEM_FRACTION", "0.7"))
+# Skip (do NOT crash) unless this much system memory is free: enough for ONE
+# resident model copy plus activations. An MLX Metal OOM is uncatchable and can
+# reboot the machine, so this guard must fail safe. ~12 GB suits the default
+# 14.3B-param 4-bit qwen2_moe; raise it for qwen3_moe (~18+).
+MIN_FREE_GB = float(os.environ.get("SGLANG_MLX_TEST_MIN_FREE_GB", "12"))
+
+# Short prompts with deterministic, quickly-terminating greedy answers.
+PROMPTS = [
+    "List the first 10 prime numbers, comma separated.",
+    "What is the capital of France? Answer in one word.",
+    "Write one short sentence about the ocean.",
+]
+MAX_NEW_TOKENS = 64  # hard cap; most prompts hit EOS well before this
+BATCH_HORIZON = 24  # fixed step count for the batching-isolation test
+
+
+def _available_gb():
+    try:
+        import psutil
+
+        return psutil.virtual_memory().available / 1024**3
+    except Exception:
+        return None  # psutil absent -> skip the pre-flight check
+
+
+@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
+class TestMlxReferenceCorrectness(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        import mlx.core as mx
+        from mlx_lm import load
+
+        avail = _available_gb()
+        if avail is not None and avail < MIN_FREE_GB:
+            raise unittest.SkipTest(
+                f"insufficient free memory: {avail:.1f} GB < {MIN_FREE_GB} GB needed "
+                f"to safely load {MODEL_PATH} (override SGLANG_MLX_TEST_MIN_FREE_GB)"
+            )
+
+        # --- Phase 1: reference tokens from UNPATCHED mlx_lm (one copy resident) ---
+        try:
+            ref_model, cls.tokenizer = load(
+                MODEL_PATH, tokenizer_config={"trust_remote_code": True}
+            )
+        except Exception as exc:  # not cached / offline / bad path
+            raise unittest.SkipTest(f"could not load {MODEL_PATH}: {exc}")
+
+        eos = getattr(cls.tokenizer, "eos_token_ids", None) or {
+            cls.tokenizer.eos_token_id
+        }
+        cls.eos_ids = set(eos)
+
+        cls.cases = []  # (prompt, prompt_ids, reference_token_ids)
+        for prompt in PROMPTS:
+            prompt_ids = list(
+                cls.tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}], add_generation_prompt=True
+                )
+            )
+            ref_ids = cls._reference_greedy(
+                ref_model, cls.tokenizer, prompt_ids, MAX_NEW_TOKENS
+            )
+            cls.cases.append((prompt, prompt_ids, ref_ids))
+
+        # --- Release the reference BEFORE building the runner (cap peak at 1x) ---
+        del ref_model
+        gc.collect()
+        mx.clear_cache()
+        active_gb = mx.get_active_memory() / 1024**3
+        if active_gb > 2.0:
+            # Weights were not released; loading the runner now would double resident
+            # memory and risk an OOM. Bail safely instead.
+            raise unittest.SkipTest(
+                f"reference model not released (active={active_gb:.1f} GB); "
+                "skipping to avoid a double-resident OOM"
+            )
+
+        # --- Phase 2: SGLang runner (one copy resident) ---
+        from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
+
+        cls.runner = MlxModelRunner(
+            model_path=MODEL_PATH,
+            trust_remote_code=True,
+            disable_radix_cache=True,  # per-request contiguous caches; no big pool
+            mem_fraction_static=MEM_FRACTION_STATIC,
+        )
+        cls.runner.init_cache_pools(req_to_token_pool=None)
+
+    @classmethod
+    def tearDownClass(cls):
+        runner = getattr(cls, "runner", None)
+        if runner is not None:
+            runner.clear()
+        cls.runner = None
+        gc.collect()
+        try:
+            import mlx.core as mx
+
+            mx.clear_cache()
+        except Exception:
+            pass
+
+    # --- helpers ----------------------------------------------------------
+
+    @staticmethod
+    def _reference_greedy(model, tokenizer, prompt_ids, max_new):
+        """Ground-truth token ids from raw, unpatched mlx_lm greedy generation."""
+        import mlx.core as mx
+        from mlx_lm import stream_generate
+        from mlx_lm.sample_utils import make_sampler
+
+        sampler = make_sampler(temp=0.0)  # greedy / argmax
+        out = []
+        for resp in stream_generate(
+            model, tokenizer, mx.array(prompt_ids), max_tokens=max_new, sampler=sampler
+        ):
+            out.append(int(resp.token))
+        return out
+
+    def _prefill(self, rid, prompt_ids):
+        return int(
+            self.runner.prefill(
+                req_id=rid,
+                new_token_ids=list(prompt_ids),
+                full_token_ids=list(prompt_ids),
+                prefix_slot_ids=[],
+                new_slot_ids=[],
+                req_pool_idx=0,
+            )
+        )
+
+    def _decode(self, rids):
+        return [int(t) for t in self.runner.decode_batch(rids)]
+
+    def _sglang_greedy(self, rid, prompt_ids, max_new):
+        """SGLang MLX greedy generation, stopping at EOS like the reference."""
+        tok = self._prefill(rid, prompt_ids)
+        out = [tok]
+        while len(out) < max_new and tok not in self.eos_ids:
+            tok = self._decode([rid])[0]
+            out.append(tok)
+        self.runner.remove_request(rid)
+        return out
+
+    def _diff_msg(self, prompt, ref, sgl):
+        horizon = min(len(ref), len(sgl))
+        first = next((j for j in range(horizon) if ref[j] != sgl[j]), horizon)
+        return (
+            f"\nprompt: {prompt!r}"
+            f"\n  first divergence @ index {first} (len ref={len(ref)} sgl={len(sgl)})"
+            f"\n  ref text: {self.tokenizer.decode(ref)!r}"
+            f"\n  sgl text: {self.tokenizer.decode(sgl)!r}"
+        )
+
+    # --- tests ------------------------------------------------------------
+
+    def test_greedy_matches_reference_exact(self):
+        """SGLang MLX greedy output == unpatched mlx_lm greedy output, token-for-token."""
+        for i, (prompt, prompt_ids, ref) in enumerate(self.cases):
+            sgl = self._sglang_greedy(f"ref-{i}", prompt_ids, MAX_NEW_TOKENS)
+            self.assertEqual(sgl, ref, self._diff_msg(prompt, ref, sgl))
+
+    def test_batched_decode_matches_solo(self):
+        """A request's tokens are identical whether decoded alone or in a batch.
+
+        Pins slot/cache isolation: ``decode_batch`` over several concurrent
+        requests must not let one request's state bleed into another's.
+        """
+        ids_list = [prompt_ids for (_, prompt_ids, _) in self.cases]
+
+        # Solo: prefill, decode a fixed horizon, remove -- one request at a time.
+        solo = []
+        for i, ids in enumerate(ids_list):
+            seq = [self._prefill(f"solo-{i}", ids)]
+            for _ in range(BATCH_HORIZON - 1):
+                seq.append(self._decode([f"solo-{i}"])[0])
+            self.runner.remove_request(f"solo-{i}")
+            solo.append(seq)
+
+        # Batched: prefill all, then advance them together in one decode_batch.
+        rids = [f"batch-{i}" for i in range(len(ids_list))]
+        batched = [[self._prefill(rid, ids)] for rid, ids in zip(rids, ids_list)]
+        for _ in range(BATCH_HORIZON - 1):
+            for j, t in enumerate(self._decode(rids)):
+                batched[j].append(t)
+        for rid in rids:
+            self.runner.remove_request(rid)
+
+        for i, (prompt, _, _) in enumerate(self.cases):
+            self.assertEqual(
+                batched[i], solo[i], self._diff_msg(prompt, solo[i], batched[i])
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)