[Apple Silicon] [CI] Move the MLX lane to the check-changes + pr-gate composite (#30121)
This commit is contained in:
@@ -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,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
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user