[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
+123 -36
View File
@@ -3,32 +3,80 @@ name: PR Test (MLX)
on:
pull_request:
branches: [ main ]
types: [opened, synchronize, reopened, labeled]
workflow_dispatch:
inputs:
# target_stage is the job id (e.g. stage-b-e2e-test-mlx), distinct from
# the run_suite.py suite name it dispatches (e.g. stage-b-e2e-mlx).
target_stage:
description: "Select a stage to run from dropdown (leave empty for auto-detect)"
required: false
type: choice
default: ''
options:
- ''
- stage-a-unit-test-mlx
- stage-b-e2e-test-mlx
concurrency:
group: pr-test-mlx-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
mlx-unit-test:
# Label-gated on 'apple-silicon', reusing the repo's run-ci label-gate mechanism
# (pr-test-rust.yml). The gate keys on the PR's current label set, so any
# event on a labeled PR re-runs the tests and a later unrelated label can't
# mask a prior result. The non-pull_request clause covers push / dispatch.
# ==================== Check Changes ==================== #
check-changes:
runs-on: ubuntu-latest
outputs:
changes_exist: ${{ steps.filter.outputs.main_package == 'true'
|| steps.filter.outputs.sgl_kernel == 'true'
|| inputs.target_stage != '' }}
main_package: ${{ steps.filter.outputs.main_package }}
sgl_kernel: ${{ steps.filter.outputs.sgl_kernel }}
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Detect file changes
id: filter
uses: dorny/paths-filter@v3
with:
# Ignored on pull_request; workflow_dispatch has no PR base without it.
base: main
filters: |
main_package:
- "python/sglang/!(multimodal_gen)/**/!(*.md)"
- "python/pyproject_other.toml"
- "test/**/!(*.md)"
- ".github/workflows/pr-test-mlx.yml"
sgl_kernel:
- "sgl-kernel/**/!(*.md|THIRDPARTYNOTICES.txt|LICENSE)"
- ".github/workflows/pr-test-mlx.yml"
# ==================== PR Gate ==================== #
pr-gate:
needs: check-changes
if: needs.check-changes.outputs.changes_exist == 'true'
uses: ./.github/workflows/pr-gate.yml
secrets: inherit
stage-a-unit-test-mlx:
needs: [check-changes, pr-gate]
# !cancelled() lets dispatch past a skipped pr-gate; the pull_request
# path still requires pr-gate success, keeping the run-ci gate intact.
if: |
github.event_name != 'pull_request' ||
contains(github.event.pull_request.labels.*.name, 'apple-silicon')
!cancelled() && (
(needs.pr-gate.result == 'success' &&
(needs.check-changes.outputs.main_package == 'true' ||
needs.check-changes.outputs.sgl_kernel == 'true')) ||
(github.event_name == 'workflow_dispatch' &&
inputs.target_stage == 'stage-a-unit-test-mlx')
)
runs-on: macos-26
timeout-minutes: 60
env:
SGLANG_IS_IN_CI: true
# use_mlx() needs this var (and mlx importable); without it the profiler
# takes the MPS branch and the MLX capture path goes untested.
# use_mlx() needs this; without it the profiler takes the MPS branch.
SGLANG_USE_MLX: 1
# Forbid HF downloads so the model-free guarantee is enforced, not assumed.
# Forbid HF downloads to enforce the model-free guarantee.
HF_HUB_OFFLINE: 1
steps:
- name: Checkout code
@@ -54,14 +102,11 @@ jobs:
timeout-minutes: 30
run: |
test -f python/pyproject_other.toml || { echo "alt pyproject_other.toml missing"; exit 1; }
# Swap in the Apple Silicon project metadata, then install the srt_mps
# extra (mlx, mlx-lm, runtime_common, torch) into an isolated uv venv.
# The all_mps diffusion chain is not needed here.
# Swap in the Apple Silicon pyproject; srt_mps skips the all_mps diffusion chain.
rm -f python/pyproject.toml
mv python/pyproject_other.toml python/pyproject.toml
uv venv
uv pip install -e "python[srt_mps]"
uv pip install pytest
uv pip install -e "python[srt_mps,test]"
- name: Report MLX / torch versions
run: |
@@ -71,27 +116,69 @@ jobs:
- name: Run model-free MLX unit tests
timeout-minutes: 15
run: |
# Model-free MLX unit tests via pytest, like MUSA's unit-test jobs
# (its model and server suites use run_suite.py). None load a model:
# signature contracts, mocked Metal capture, dummy ServerArgs patching,
# and quant-config dicts.
uv run python -m pytest -v \
test/registered/unit/hardware_backend/mlx/test_runner_init_contract.py \
test/registered/unit/hardware_backend/mlx/test_metal_profiler.py \
test/registered/unit/hardware_backend/mlx/test_attention_patching.py \
"test/registered/unit/hardware_backend/mlx/test_quantization.py::TestMlxQuantizationOverride"
uv run python test/run_suite.py --hw mlx --suite stage-a-unit-test-mlx
stage-b-e2e-test-mlx:
needs: [check-changes, pr-gate]
# Manual-only: no self hosted Apple Silicon runner is generally available; hosted runners are too small for e2e.
if: |
!cancelled() && inputs.target_stage == 'stage-b-e2e-test-mlx'
runs-on: [self-hosted, macOS, ARM64]
timeout-minutes: 120
env:
SGLANG_IS_IN_CI: true
SGLANG_USE_MLX: 1
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
cache-dependency-glob: 'python/pyproject_other.toml'
- name: Verify Apple Silicon runner
run: |
echo "uname -m: $(uname -m)"
python3 -c "import platform; assert platform.machine()=='arm64', platform.machine(); print('machine:', platform.machine(), 'system:', platform.system())"
- name: Install dependencies (MLX / srt_mps extra)
timeout-minutes: 30
run: |
test -f python/pyproject_other.toml || { echo "alt pyproject_other.toml missing"; exit 1; }
rm -f python/pyproject.toml
mv python/pyproject_other.toml python/pyproject.toml
uv venv --python 3.11
uv pip install -e "python[srt_mps,test]"
# Needs a warm HF cache or network and roughly 24 GB of unified memory.
- name: Run MLX e2e smoke tests
timeout-minutes: 90
run: |
uv run python test/run_suite.py --hw mlx --suite stage-b-e2e-mlx
pr-test-mlx-finish:
needs: [mlx-unit-test]
needs:
[
pr-gate,
check-changes,
stage-a-unit-test-mlx,
stage-b-e2e-test-mlx,
]
if: always()
runs-on: ubuntu-latest
steps:
- name: Check dependent job status
- name: Check all dependent job statuses
run: |
result="${{ needs.mlx-unit-test.result }}"
echo "mlx-unit-test: $result"
if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then
echo "The MLX unit-test job failed."
exit 1
fi
json_needs='${{ toJson(needs) }}'
job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]')
for job in $job_names; do
result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result')
echo "$job: $result"
if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then
echo "The above jobs failed."
exit 1
fi
done
echo "All jobs completed successfully"
+16
View File
@@ -16,6 +16,7 @@ __all__ = [
"register_npu_ci",
"register_xpu_ci",
"register_musa_ci",
"register_mlx_ci",
"ut_parse_one_file",
]
@@ -35,6 +36,7 @@ class HWBackend(Enum):
NPU = auto()
XPU = auto()
MUSA = auto()
MLX = auto()
@dataclass
@@ -148,6 +150,19 @@ def register_musa_ci(
return None
def register_mlx_ci(
est_time: float,
suite: Optional[str] = None,
nightly: bool = False,
disabled: Optional[str] = None,
*,
stage: Optional[str] = None,
runner_config: Optional[str] = None,
):
"""Marker for MLX CI registration (parsed via AST; runtime no-op)."""
return None
REGISTER_MAPPING = {
"register_cpu_ci": HWBackend.CPU,
"register_cuda_ci": HWBackend.CUDA,
@@ -156,6 +171,7 @@ REGISTER_MAPPING = {
"register_npu_ci": HWBackend.NPU,
"register_xpu_ci": HWBackend.XPU,
"register_musa_ci": HWBackend.MUSA,
"register_mlx_ci": HWBackend.MLX,
}
@@ -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,
}