[Apple Silicon] [CI] Move the MLX lane to the check-changes + pr-gate composite (#30121)

This commit is contained in:
Jae B.
2026-07-10 22:00:15 -07:00
committed by GitHub
parent e1d51be91f
commit a91c2e6596
15 changed files with 288 additions and 145 deletions
@@ -5,7 +5,7 @@ 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.ci.ci_register import register_cpu_ci, register_mlx_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
@@ -15,9 +15,10 @@ from sglang.test.test_utils import (
)
# 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.
# only on Apple Silicon. Also registered under stage-b-e2e-mlx, which the
# macOS CI lane (pr-test-mlx.yml) only dispatches via a gated workflow_dispatch.
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
register_mlx_ci(est_time=1, suite="stage-b-e2e-mlx")
_HAS_MLX = importlib.util.find_spec("mlx") is not None
@@ -5,7 +5,7 @@ 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.ci.ci_register import register_cpu_ci, register_mlx_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
@@ -15,9 +15,10 @@ from sglang.test.test_utils import (
)
# 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.
# only on Apple Silicon. Also registered under stage-b-e2e-mlx, which the
# macOS CI lane (pr-test-mlx.yml) only dispatches via a gated workflow_dispatch.
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
register_mlx_ci(est_time=1, suite="stage-b-e2e-mlx")
_HAS_MLX = importlib.util.find_spec("mlx") is not None
@@ -7,9 +7,10 @@ import unittest
from collections import deque
from types import SimpleNamespace
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
_HAS_MLX = importlib.util.find_spec("mlx") is not None
_SKIP_REASON = "requires mlx"
@@ -20,9 +20,10 @@ import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
register_mlx_ci(est_time=5, suite="stage-a-unit-test-mlx")
_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
_HAS_MLX = importlib.util.find_spec("mlx") is not None
@@ -19,10 +19,11 @@ from __future__ import annotations
import importlib.util
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
register_mlx_ci(est_time=4, suite="stage-a-unit-test-mlx")
_HAS_MLX = (
importlib.util.find_spec("mlx") is not None
@@ -0,0 +1,106 @@
"""Unit tests for ``MlxQuantizationConfig.override_quantization_method``.
The override is a classmethod over a dict; no mlx / Apple Silicon dependency.
Runs on every CI platform and guards #25119 from regression.
"""
from __future__ import annotations
import unittest
from sglang.srt.layers.quantization.mlx import MlxQuantizationConfig
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
class TestMlxQuantizationOverride(unittest.TestCase):
"""Pure-logic tests for ``MlxQuantizationConfig.override_quantization_method``.
The override is a classmethod over a dict; no mlx / Apple Silicon
dependency. Runs on every CI platform and guards #25119 from regression.
"""
def test_mlx_q4_dict_config_autodetect(self):
"""Bare {group_size, bits=4} dict maps to mlx_q4."""
result = MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, None
)
self.assertEqual(result, "mlx_q4")
def test_mlx_q8_dict_config_autodetect(self):
"""Bare {group_size, bits=8} dict maps to mlx_q8."""
result = MlxQuantizationConfig.override_quantization_method(
{"group_size": 32, "bits": 8}, None
)
self.assertEqual(result, "mlx_q8")
def test_non_mlx_dict_not_matched(self):
"""Dicts with an explicit quant_method belong to that method, not ours."""
# modelopt-style: explicit quant_method takes priority.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"quant_method": "modelopt", "bits": 4, "group_size": 64}, None
)
)
# gptq-style: same.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"quant_method": "gptq", "bits": 4, "group_size": 128}, None
)
)
def test_non_dict_not_matched(self):
"""Non-dict inputs and malformed dicts return None."""
# None / string inputs.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(None, None)
)
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method("mlx_q4", None)
)
# Missing keys.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method({"bits": 4}, None)
)
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method({"group_size": 64}, None)
)
# Non-integer values.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"bits": "4", "group_size": 64}, None
)
)
# Unsupported bit-width.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"bits": 2, "group_size": 64}, None
)
)
def test_user_quant_explicit_defers_to_user(self):
"""When the user passes --quantization explicitly, defer to that choice."""
# User chose mlx_q8 explicitly, even though config dict shape suggests q4
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, "mlx_q8"
)
)
# User chose mlx_q4 explicitly with matching config
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, "mlx_q4"
)
)
# User chose something completely different
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, "fp8"
)
)
if __name__ == "__main__":
unittest.main()
@@ -40,10 +40,11 @@ import importlib.util
import os
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
register_mlx_ci(est_time=1, suite="stage-b-e2e-mlx")
_HAS_MLX = (
importlib.util.find_spec("mlx") is not None
@@ -16,9 +16,10 @@ import importlib.util
import inspect
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
_HAS_MLX = importlib.util.find_spec("mlx") is not None
_SKIP_REASON = "requires mlx"
@@ -17,14 +17,18 @@ import importlib.util
import platform
import unittest
from sglang.srt.layers.quantization.mlx import MlxQuantizationConfig
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
# Registered with the CPU suite (runtime no-op marker, parsed via AST).
# On non-Apple-Silicon CI runners the entire TestCase class skips via the
# @skipUnless guard below, so this registration is the harmless "yes this
# test exists" signal the registry requires.
# Registered on the CPU suite but skipped wherever mlx is absent; runs for real
# only on Apple Silicon. Also registered under stage-b-e2e-mlx, not stage-a:
# this class loads real HF models (Qwen/Qwen3-0.6B, mlx-community/Qwen3-0.6B-4bit)
# via MlxModelRunner, and stage-a's job env sets HF_HUB_OFFLINE=1 to enforce a
# model-free guarantee (.github/workflows/pr-test-mlx.yml) -- confirmed this
# fails with LocalEntryNotFoundError on a runner with no pre-warmed cache. The
# macOS CI lane (pr-test-mlx.yml) only dispatches stage-b-e2e-mlx via a gated
# workflow_dispatch, matching the models_e2e correctness tests' convention.
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
register_mlx_ci(est_time=10, suite="stage-b-e2e-mlx")
_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
_HAS_MLX = (
@@ -187,92 +191,5 @@ class TestMlxQuantization(unittest.TestCase):
self._reset_mlx_memory()
class TestMlxQuantizationOverride(unittest.TestCase):
"""Pure-logic tests for ``MlxQuantizationConfig.override_quantization_method``.
The override is a classmethod over a dict; no mlx / Apple Silicon
dependency. Runs on every CI platform and guards #25119 from regression.
"""
def test_mlx_q4_dict_config_autodetect(self):
"""Bare {group_size, bits=4} dict maps to mlx_q4."""
result = MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, None
)
self.assertEqual(result, "mlx_q4")
def test_mlx_q8_dict_config_autodetect(self):
"""Bare {group_size, bits=8} dict maps to mlx_q8."""
result = MlxQuantizationConfig.override_quantization_method(
{"group_size": 32, "bits": 8}, None
)
self.assertEqual(result, "mlx_q8")
def test_non_mlx_dict_not_matched(self):
"""Dicts with an explicit quant_method belong to that method, not ours."""
# modelopt-style: explicit quant_method takes priority.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"quant_method": "modelopt", "bits": 4, "group_size": 64}, None
)
)
# gptq-style: same.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"quant_method": "gptq", "bits": 4, "group_size": 128}, None
)
)
def test_non_dict_not_matched(self):
"""Non-dict inputs and malformed dicts return None."""
# None / string inputs.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(None, None)
)
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method("mlx_q4", None)
)
# Missing keys.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method({"bits": 4}, None)
)
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method({"group_size": 64}, None)
)
# Non-integer values.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"bits": "4", "group_size": 64}, None
)
)
# Unsupported bit-width.
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"bits": 2, "group_size": 64}, None
)
)
def test_user_quant_explicit_defers_to_user(self):
"""When the user passes --quantization explicitly, defer to that choice."""
# User chose mlx_q8 explicitly, even though config dict shape suggests q4
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, "mlx_q8"
)
)
# User chose mlx_q4 explicitly with matching config
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, "mlx_q4"
)
)
# User chose something completely different
self.assertIsNone(
MlxQuantizationConfig.override_quantization_method(
{"group_size": 64, "bits": 4}, "fp8"
)
)
if __name__ == "__main__":
unittest.main()
@@ -11,9 +11,10 @@ import importlib.util
import inspect
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
_HAS_MLX = importlib.util.find_spec("mlx") is not None
_SKIP_REASON = "requires mlx"
@@ -18,9 +18,10 @@ import platform
import unittest
from unittest.mock import MagicMock
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
register_mlx_ci(est_time=5, suite="stage-a-unit-test-mlx")
_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
_HAS_MLX = importlib.util.find_spec("mlx") is not None
@@ -35,11 +35,13 @@ from types import SimpleNamespace
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
# AST-parsed "this test exists" marker; actual execution is gated by the
# @skipUnless guard below (mirrors test_quantization.py in this directory).
# CPU marker is AST-parsed "this test exists"; actual CPU-side execution is
# gated by the @skipUnless guard below. MLX marker runs for real on the MLX
# lane's stage-a (model-free: mocks the runner, loads no model).
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
register_mlx_ci(est_time=10, suite="stage-a-unit-test-mlx")
_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
_HAS_MLX = importlib.util.find_spec("mlx") is not None
+6
View File
@@ -22,6 +22,7 @@ HW_MAPPING = {
"musa": HWBackend.MUSA,
"npu": HWBackend.NPU,
"xpu": HWBackend.XPU,
"mlx": HWBackend.MLX,
}
# Per-commit test suites (run on every PR).
@@ -106,6 +107,10 @@ PER_COMMIT_SUITES = {
"stage-a-test-1-gpu-xpu",
"stage-b-test-1-gpu-xpu",
],
HWBackend.MLX: [
"stage-a-unit-test-mlx",
"stage-b-e2e-mlx",
],
}
# Nightly test suites (run nightly, organized by GPU configuration)
@@ -194,6 +199,7 @@ _SUITE_CHECKED_BACKENDS = {
HWBackend.CPU,
HWBackend.MUSA,
HWBackend.XPU,
HWBackend.MLX,
}