[MLX] Add on-the-fly --quantization mlx_q4 / mlx_q8 for Apple Silicon (#24907)

Co-authored-by: lezhang <lezhang@local>

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/<model>-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
 

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

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

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,

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)

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",
 ]
 

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/<model>-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()