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

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

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