diff --git a/docs_new/docs/hardware-platforms/apple_metal.mdx b/docs_new/docs/hardware-platforms/apple_metal.mdx index 9e71786d9..9b45fd2db 100644 --- a/docs_new/docs/hardware-platforms/apple_metal.mdx +++ b/docs_new/docs/hardware-platforms/apple_metal.mdx @@ -60,6 +60,30 @@ SGLANG_USE_MLX=1 python -m sglang.launch_server \ 2. `--disable-cuda-graph` - Disables usage of CUDA graph, which is not relevant for Apple Metal. 3. `--disable-overlap-schedule` - Disables overlap scheduling (enabled/not present by default) achieved using MLX's `async_eval()` +## Quantization + +The MLX backend supports two quantization paths on Apple Silicon: + +1. **Pre-quantized HF repos.** Any `mlx-community/-4bit` (or `-8bit`) repo loads directly through `mlx_lm.load(...)` — no extra flag needed. + ```bash + SGLANG_USE_MLX=1 python -m sglang.launch_server \ + --model-path mlx-community/Qwen3-0.6B-4bit \ + --disable-cuda-graph + ``` +2. **On-the-fly quantization.** For any fp16 model, pass `--quantization mlx_q4` or `--quantization mlx_q8` to have sglang quantize the weights at load time via `mlx_lm.utils.quantize_model` (group size 64, the mlx-community default). The quantized weights stay in process memory; the on-disk model is untouched. + ```bash + SGLANG_USE_MLX=1 python -m sglang.launch_server \ + --model-path Qwen/Qwen3-0.6B \ + --quantization mlx_q4 \ + --disable-cuda-graph + ``` + Expected log line: + ``` + Quantizing MLX model on-the-fly: bits=4 group_size=64 (preset=mlx_q4) + Quantization complete in 0.13s — active mem: 1.11 GB -> 0.31 GB (71.9% reduction) + ``` + The MLX backend silently ignores `--quantization mlx_q4` when the model is already quantized in its HF config (path 1), so the same flag is safe to pass either way. + ## Benchmarking with Requests diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index 4c47d67f0..1fb5bb33b 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner.py @@ -20,7 +20,9 @@ from dataclasses import dataclass import mlx.core as mx import psutil +from mlx.utils import tree_flatten from mlx_lm import load as mlx_lm_load +from mlx_lm.utils import quantize_model as mlx_lm_quantize_model from sglang.srt.hardware_backend.mlx.kv_cache import ( BatchedDecodeContext, @@ -90,6 +92,13 @@ class MlxPendingDecode: caches: list # list[list[ContiguousKVCache]] +_MLX_QUANTIZATION_PRESETS: dict[str, tuple[int, int]] = { + # name -> (bits, group_size). group_size=64 matches the mlx-community convention. + "mlx_q4": (4, 64), + "mlx_q8": (8, 64), +} + + class MlxModelRunner: """MLX model runner with radix-cache prefix sharing.""" @@ -100,6 +109,7 @@ class MlxModelRunner: disable_radix_cache: bool = False, pool_size: int | None = None, mem_fraction_static: float = 0.8, + quantization: str | None = None, ): self.model_path = model_path self.trust_remote_code = trust_remote_code @@ -108,6 +118,11 @@ class MlxModelRunner: self._mem_fraction_static = mem_fraction_static # Counter used to trigger periodic mx.clear_cache() calls. self._decode_step_ct: int = 0 + # On-the-fly quantization preset (e.g. "mlx_q4"). None = no on-load quantization. + # Pre-quantized HF repos (e.g. mlx-community/Qwen3-0.6B-4bit) load correctly + # regardless of this setting — mlx_lm.load() detects the config and instantiates + # QuantizedLinear modules directly. + self._quantization: str | None = quantization self._load_model() @@ -180,14 +195,64 @@ class MlxModelRunner: ] def _load_model(self): - """Load model using mlx_lm.""" + """Load model using mlx_lm. If ``self._quantization`` requests a preset + (e.g. ``mlx_q4``), quantize fp16 weights in-place via + :func:`mlx_lm.utils.quantize_model` after load. + """ logger.info(f"Loading MLX model: {self.model_path}") start_time = time.time() - self.model, _ = mlx_lm_load( + # We need the config dict to pass into quantize_model so it knows tied/embedding + # layout. return_config=True is cheap and ignored when no quantization is requested. + loaded = mlx_lm_load( self.model_path, tokenizer_config={"trust_remote_code": self.trust_remote_code}, + return_config=True, ) + self.model, _tokenizer, config = loaded + + if self._quantization in _MLX_QUANTIZATION_PRESETS: + bits, group_size = _MLX_QUANTIZATION_PRESETS[self._quantization] + # Skip if the model was already loaded quantized (pre-quantized HF repo); + # mlx_lm.load detects the config and instantiates QuantizedLinear directly, + # so applying the preset on top would be redundant. + if "quantization" in (config or {}): + logger.info( + "MLX model is already quantized by the HF repo; " + f"ignoring --quantization={self._quantization}" + ) + else: + # Read weight-tensor totals from MLX array metadata (shape + dtype). + # This is zero-cost — neither materializes the lazy fp16 weights nor + # forces them to be peak-resident in memory at once (which on a 64 GB + # Mac running a 32 B model would put us within a few GB of OOM). + bytes_before = sum( + p.size * p.itemsize + for _, p in tree_flatten(self.model.parameters()) + ) + q_start = time.time() + logger.info( + f"Quantizing MLX model on-the-fly: bits={bits} " + f"group_size={group_size} (preset={self._quantization})" + ) + self.model, _new_config = mlx_lm_quantize_model( + self.model, + config or {}, + group_size=group_size, + bits=bits, + ) + bytes_after = sum( + p.size * p.itemsize + for _, p in tree_flatten(self.model.parameters()) + ) + q_time = time.time() - q_start + pct_reduction = (1 - bytes_after / max(bytes_before, 1)) * 100 + logger.info( + f"Quantization complete in {q_time:.2f}s — " + f"weight bytes: {bytes_before / 1024**3:.2f} GB -> " + f"{bytes_after / 1024**3:.2f} GB ({pct_reduction:.1f}% reduction)" + ) + # Force-evaluate weights so mx.get_active_memory() reflects # actual usage before KV pool sizing. mx.eval(self.model.parameters()) diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index a12d25240..8dddf97f2 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/python/sglang/srt/hardware_backend/mlx/tp_worker.py @@ -53,6 +53,7 @@ class MlxTpModelWorker(TpModelWorker): trust_remote_code=self.server_args.trust_remote_code, disable_radix_cache=self.server_args.disable_radix_cache, mem_fraction_static=self.server_args.mem_fraction_static, + quantization=self.server_args.quantization, ) if self.server_args.max_total_tokens is not None: init_kwargs["pool_size"] = self.server_args.max_total_tokens diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index b63db577f..095a36922 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -29,6 +29,7 @@ from sglang.srt.layers.quantization.fpgemm_fp8 import FBGEMMFp8Config from sglang.srt.layers.quantization.gguf import GGUFConfig from sglang.srt.layers.quantization.gptq import GPTQConfig, GPTQMarlinConfig from sglang.srt.layers.quantization.gptq_cpu import CPUGPTQConfig +from sglang.srt.layers.quantization.mlx import MlxQuantizationConfig from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ModelOptFp8Config, @@ -48,6 +49,7 @@ from sglang.srt.utils import ( cpu_has_amx_support, is_cuda, is_hip, + is_mps, is_npu, mxfp_supported, ) @@ -94,6 +96,15 @@ if is_cuda() or (_is_mxfp_supported and is_hip()): } ) + +if is_mps(): + BASE_QUANTIZATION_METHODS.update( + { + "mlx_q4": MlxQuantizationConfig, + "mlx_q8": MlxQuantizationConfig, + } + ) + # subset of above quant methods, supported on CPU CPU_QUANTIZATION_METHODS = { "fp8": Fp8Config, diff --git a/python/sglang/srt/layers/quantization/mlx.py b/python/sglang/srt/layers/quantization/mlx.py new file mode 100644 index 000000000..60d43cdd3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/mlx.py @@ -0,0 +1,74 @@ +"""Marker config for MLX backend on-the-fly quantization (mlx_q4 / mlx_q8). + +The MLX backend (``python/sglang/srt/hardware_backend/mlx/``) performs its own +quantization at model-load time via :func:`mlx_lm.utils.quantize_model`. The +standard PyTorch ``QuantizationConfig`` machinery is **never** invoked on that +path. + +This module exists purely so that the names ``mlx_q4`` and ``mlx_q8`` are +recognized by ``QUANTIZATION_METHODS`` — that way +:meth:`ModelConfig._verify_quantization` and downstream registry lookups treat +them as known methods without any backend-specific carve-outs in the generic +config code. + +If a user passes ``--quantization mlx_q4`` without ``SGLANG_USE_MLX=1`` they +will eventually reach a code path that tries to instantiate this Config class, +at which point we raise a clear error. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import torch + +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) + + +class MlxQuantizationConfig(QuantizationConfig): + """Marker config for MLX backend on-the-fly quantization presets. + + Not a real quantization config — the MLX backend handles quantization + itself. Any standard-PyTorch-path method that touches this class raises + a helpful error pointing the user at ``SGLANG_USE_MLX=1``. + """ + + _ERR = ( + "MLX on-the-fly quantization (--quantization mlx_q4 / mlx_q8) is " + "handled by the MLX backend at model-load time via mlx_lm.utils." + "quantize_model, not by this QuantizationConfig class. If you " + "reached this error, SGLANG_USE_MLX=1 is likely not set." + ) + + def __init__(self, preset: str): + super().__init__() + self.preset = preset + + @classmethod + def get_name(cls) -> str: + return "mlx" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [] + + @classmethod + def get_min_capability(cls) -> int: + # Capability check is for NVIDIA SM versions; not meaningful for MLX. + return 0 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "MlxQuantizationConfig": + raise NotImplementedError(cls._ERR) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + raise NotImplementedError(self._ERR) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d70f23191..0391dd2f0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -134,6 +134,10 @@ QUANTIZATION_CHOICES = [ "modelslim", # for NPU "quark", # AMD Quark quantizer (FP8 / MXFP4 / Int4FP8 etc.) "quark_int4fp8_moe", + # Apple Silicon MLX backend — on-the-fly quantization of fp16 weights at load + # time via mlx.nn.quantize. Only takes effect when SGLANG_USE_MLX=1. + "mlx_q4", # 4 bits, group_size=64 (mlx-community default) + "mlx_q8", # 8 bits, group_size=64 "unquant", ] diff --git a/test/registered/unit/hardware_backend/mlx/test_quantization.py b/test/registered/unit/hardware_backend/mlx/test_quantization.py new file mode 100644 index 000000000..76f427f2d --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_quantization.py @@ -0,0 +1,190 @@ +"""Unit tests for MLX backend on-the-fly quantization. + +Covers: + - mlx_q4 / mlx_q8 quantize fp16 weights to QuantizedLinear in-place + - active-memory drops after quantization + - smoke /generate still works post-quantize + - pre-quantized HF repos still load (regression guard for mlx_lm passthrough) + - mlx_q4 flag on an already-quantized model is a no-op (skip + log) + +Skips on non-Apple-Silicon platforms and when ``mlx`` / ``mlx_lm`` are missing. +""" + +from __future__ import annotations + +import gc +import importlib.util +import platform +import unittest + +from sglang.test.ci.ci_register import register_cpu_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. +register_cpu_ci(est_time=10, suite="stage-a-test-cpu") + +_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64" +_HAS_MLX = ( + importlib.util.find_spec("mlx") is not None + and importlib.util.find_spec("mlx_lm") is not None +) + +_SKIP_REASON = "Apple-Silicon-only test (requires Darwin/arm64 + mlx + mlx_lm)" + +# Tiny model used across tests; ~0.6B fp16 = ~1.1 GB on disk after first download. +_TEST_MODEL = "Qwen/Qwen3-0.6B" +_TEST_MODEL_PREQUANT = "mlx-community/Qwen3-0.6B-4bit" + + +@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON) +class TestMlxQuantization(unittest.TestCase): + """Smoke tests for --quantization mlx_q4 / mlx_q8 in MlxModelRunner.""" + + # ---------- helpers ---------- + + @staticmethod + def _module_counts(model) -> tuple[int, int]: + n_quant, n_linear = 0, 0 + for _, m in model.named_modules(): + cls = type(m).__name__ + if cls == "QuantizedLinear": + n_quant += 1 + elif cls == "Linear": + n_linear += 1 + return n_quant, n_linear + + @staticmethod + def _reset_mlx_memory() -> None: + import mlx.core as mx + + gc.collect() + mx.clear_cache() + + def _build_runner(self, model_path: str, quantization: str | None): + from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner + + return MlxModelRunner( + model_path=model_path, + quantization=quantization, + pool_size=1024, # small pool — these tests don't drive generation depth + ) + + # ---------- tests ---------- + + def test_mlx_q4_creates_quantized_linear_modules(self): + """All Linear modules should be QuantizedLinear after mlx_q4 load.""" + self._reset_mlx_memory() + runner = self._build_runner(_TEST_MODEL, "mlx_q4") + try: + n_quant, n_linear = self._module_counts(runner.model) + self.assertGreater( + n_quant, 0, "expected at least one QuantizedLinear module" + ) + self.assertEqual( + n_linear, + 0, + f"all Linear modules should have been quantized, got {n_linear} remaining", + ) + finally: + del runner + self._reset_mlx_memory() + + def test_mlx_q4_reduces_memory_vs_fp16(self): + """mlx_q4 should use meaningfully less memory than the fp16 baseline.""" + import mlx.core as mx + + self._reset_mlx_memory() + runner_fp = self._build_runner(_TEST_MODEL, None) + mx.eval(runner_fp.model.parameters()) + mem_fp = mx.get_active_memory() + del runner_fp + self._reset_mlx_memory() + + runner_q4 = self._build_runner(_TEST_MODEL, "mlx_q4") + mx.eval(runner_q4.model.parameters()) + mem_q4 = mx.get_active_memory() + del runner_q4 + self._reset_mlx_memory() + + # Conservative: expect at least 40% reduction. On Qwen3-0.6B we measured ~72%; + # 40% leaves headroom for different mlx_lm versions, model shapes, etc. + reduction = 1 - (mem_q4 / max(mem_fp, 1)) + self.assertGreater( + reduction, + 0.40, + f"expected >40% memory reduction with mlx_q4, got {reduction*100:.1f}% " + f"(fp16={mem_fp/1024**3:.2f} GB, q4={mem_q4/1024**3:.2f} GB)", + ) + + def test_mlx_q8_creates_quantized_linear_modules(self): + """Same check for the 8-bit variant.""" + self._reset_mlx_memory() + runner = self._build_runner(_TEST_MODEL, "mlx_q8") + try: + n_quant, n_linear = self._module_counts(runner.model) + self.assertGreater(n_quant, 0) + self.assertEqual(n_linear, 0) + finally: + del runner + self._reset_mlx_memory() + + def test_mlx_q4_generates_text(self): + """After on-the-fly quantization the model must still generate non-empty text.""" + from mlx_lm import generate + from transformers import AutoTokenizer + + self._reset_mlx_memory() + runner = self._build_runner(_TEST_MODEL, "mlx_q4") + try: + tok = AutoTokenizer.from_pretrained(_TEST_MODEL) + output = generate( + runner.model, + tok, + prompt="The capital of France is", + max_tokens=5, + verbose=False, + ) + self.assertIsInstance(output, str) + self.assertGreater( + len(output.strip()), 0, "generation returned empty string" + ) + finally: + del runner + self._reset_mlx_memory() + + def test_pre_quantized_hf_repo_passthrough(self): + """Loading mlx-community/-4bit must still work (mlx_lm passthrough, + regression guard for the no-quantization-flag path). + """ + self._reset_mlx_memory() + runner = self._build_runner(_TEST_MODEL_PREQUANT, quantization=None) + try: + n_quant, n_linear = self._module_counts(runner.model) + self.assertGreater( + n_quant, + 0, + "pre-quantized HF repo should load as QuantizedLinear without --quantization", + ) + finally: + del runner + self._reset_mlx_memory() + + def test_quantize_flag_on_already_quantized_model_is_noop(self): + """Passing --quantization mlx_q4 on a pre-quantized repo should NOT double-quantize.""" + self._reset_mlx_memory() + # Using mlx_q4 against an already-q4 repo. The runner logs the skip and leaves + # the existing QuantizedLinear modules untouched. + runner = self._build_runner(_TEST_MODEL_PREQUANT, "mlx_q4") + try: + n_quant, n_linear = self._module_counts(runner.model) + self.assertGreater(n_quant, 0) + self.assertEqual(n_linear, 0) + finally: + del runner + self._reset_mlx_memory() + + +if __name__ == "__main__": + unittest.main()