[MLX] Add quant regrouping to MLX lowering flow, and update quantize_and_save in Qwen to support MLX path (#21023)

For the MLX repacking and quantize_and_save work, we removed the
redundant hand-rolled MLX regrouping/packing code in both gemma4 and
qwen3.5 (deleting the old pack_mlx/export.py regroup paths and their
tests).

We also added MLX path in qwen's quantize_and_save.py

Author: metascroy

backends/mlx/builder/op_helpers.py

@@ -519,6 +519,7 @@ def to_mlx_qparams(
     zero_point: torch.Tensor,
     bits: int,
     compute_biases: bool = True,
+    prepacked: bool = False,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
     """
     Convert TorchAO quantization params to MLX format.
@@ -536,10 +537,31 @@ def to_mlx_qparams(
                        Returns (Q, None) in this case. This is valid when
                        zero_point is all zeros, as the C++ runtime will compute
                        biases = -scales * 2^(bits-1).
+        prepacked: If True, ``qdata`` is already nibble-packed uint8 holding the
+                       unsigned values ``q + offset`` (two 4-bit values per byte,
+                       even index -> low nibble) — the exact layout this function
+                       produces for ``bits == 4``. It is reinterpreted as uint32
+                       directly (``view``), skipping the int8 -> uint32 repack.
+                       Only supported for ``bits == 4``.
     """
-    assert qdata.dtype == torch.int8
     offset = 2 ** (bits - 1)
 
+    if prepacked:
+        assert bits == 4, "prepacked to_mlx_qparams only supports 4-bit"
+        assert (
+            qdata.dtype == torch.uint8
+        ), f"prepacked qdata must be uint8, got {qdata.dtype}"
+        # (rows, cols//2) uint8 -> (rows, cols//8) uint32. cols//2 must be a
+        # multiple of 4, i.e. in_features a multiple of 8 (holds: gs >= 32).
+        assert qdata.shape[-1] % 4 == 0, "packed cols must be a multiple of 4"
+        Q = qdata.contiguous().view(torch.uint32)
+        if compute_biases:
+            B = -scale * (zero_point.to(scale.dtype) + offset)
+            return Q, B
+        return Q, None
+
+    assert qdata.dtype == torch.int8
+
     # Pack data into a contiguous uint32 bitstream. cols*bits must be a
     # multiple of 32 (holds since in_features is a multiple of group_size>=32).
     rows, cols = qdata.shape
@@ -640,7 +662,9 @@ def parse_dequant_int4_node(
     """Parse a torchao.dequantize_int4_tensor node.
 
     Returns (qdata, scale, zero_point, group_size, output_dtype) or None if not a
-    dequantize_int4_tensor node or the custom op is not registered.
+    dequantize_int4_tensor node or the custom op is not registered. ``group_size``
+    is the MLX-legal group_size (16/32/64/128); coarser int4 groups (e.g. 256)
+    are regrouped to it via ``regroup_affine_scales`` at emit time.
     """
     target = get_aten_target(node.target)
     try:
@@ -652,6 +676,10 @@ def parse_dequant_int4_node(
         return None
 
     qdata, scale, zero_point, group_size = node.args[0:4]
+    mlx_group_size = mlx_affine_group_size(int(group_size))
+    if mlx_group_size is None:
+        return None
+    group_size = mlx_group_size
 
     output_dtype = None
     if len(node.args) > 4:
@@ -662,6 +690,60 @@ def parse_dequant_int4_node(
     return qdata, scale, zero_point, group_size, output_dtype
 
 
+_MLX_AFFINE_GROUP_SIZES = (128, 64, 32, 16)
+
+
+def mlx_affine_group_size(group_size: int) -> Optional[int]:
+    """Largest MLX-supported group_size (128/64/32/16) that divides ``group_size``.
+
+    MLX's affine quantized kernels only support group_size in {16, 32, 64, 128}.
+    torchao may quantize with a coarser or per-axis group_size (e.g. 256 or a
+    full row of 5376). A coarse group is a stack of finer groups that share one
+    scale, so any legal divisor is exact after repeating scale/zero_point (see
+    ``regroup_affine_scales``). Returns ``group_size`` unchanged when already
+    legal, the coarsest legal divisor otherwise, or ``None`` when none divides.
+    """
+    if group_size in _MLX_AFFINE_GROUP_SIZES:
+        return group_size
+    for candidate in _MLX_AFFINE_GROUP_SIZES:
+        if group_size % candidate == 0:
+            return candidate
+    return None
+
+
+def regroup_affine_scales(
+    scale: torch.Tensor,
+    zero_point: torch.Tensor,
+    in_features: int,
+    target_group_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor, bool]:
+    """Repeat per-group scale/zero_point so the effective group_size becomes
+    ``target_group_size`` (an MLX-legal size from ``mlx_affine_group_size``).
+
+    ``scale``/``zero_point`` are shaped ``[..., in_features // old_group_size]``.
+    A coarse group shares one scale across ``old_group_size // target_group_size``
+    finer groups, so repeat-interleaving along the last axis is numerically exact.
+
+    Returns ``(scale, zero_point, changed)``; a no-op (``changed=False``) when the
+    group_size is already ``target_group_size``.
+    """
+    old_groups = scale.shape[-1]
+    old_group_size = in_features // old_groups
+    assert (
+        old_group_size >= target_group_size and old_group_size % target_group_size == 0
+    ), (
+        f"cannot regroup: weight group_size={old_group_size} is finer than, or "
+        f"not a multiple of, target_group_size={target_group_size} — "
+        f"repeat-interleave can only refine groups"
+    )
+    repeat = old_group_size // target_group_size
+    if repeat == 1:
+        return scale, zero_point, False
+    scale = scale.repeat_interleave(repeat, dim=-1)
+    zero_point = zero_point.repeat_interleave(repeat, dim=-1)
+    return scale, zero_point, True
+
+
 def parse_dequant_node(
     node: Node,
 ) -> Optional[Tuple[Node, Node, Node, int, int, Optional[torch.dtype], int]]:
@@ -673,7 +755,9 @@ def parse_dequant_node(
       - Conv2d weights (4D):  block_size=[1, 32, 1, 1] → quantized_dim=1
 
     Returns (qdata, scale, zero_point, group_size, bits, out_dtype, quantized_dim)
-    or None if unsupported.
+    or None if unsupported. ``group_size`` is the MLX-legal group_size (16/32/64/
+    128); when the weight uses a coarser group, callers must regroup scale/
+    zero_point to it via ``regroup_affine_scales`` before packing.
     """
     qdata, block_size, scale, zero_point, dtype, qmin, qmax = node.args[0:7]
     out_dtype = (
@@ -687,8 +771,13 @@ def parse_dequant_node(
     if len(non_one) != 1:
         return None
     quantized_dim, group_size = non_one[0]
-    if group_size not in [16, 32, 64, 128]:
+    # MLX kernels only support group_size in {16,32,64,128}. Coarser/per-axis
+    # groups are lowered by repeating scale/zero_point to the coarsest legal
+    # divisor at emit time (regroup_affine_scales); reject if none divides.
+    mlx_group_size = mlx_affine_group_size(group_size)
+    if mlx_group_size is None:
         return None
+    group_size = mlx_group_size
 
     # MLX supports 2,3,4,5,6,8-bit affine quantization. to_mlx_qparams packs
     # 2/4/8 via fast paths and other widths (e.g. 5, 6) via a general

backends/mlx/custom_ops.py

@@ -354,6 +354,17 @@ def gather_qmm(
         s_sel = scales
         b_sel = biases
 
+    subbyte_packed = w_sel.dtype == torch.uint8
+    if subbyte_packed:
+        assert (
+            bits == 4
+        ), "Subbyte packing qdata (uint8) is only supported for bits=4 now."
+
+        offset = 2 ** (bits - 1)
+        lo = (w_sel & 0x0F).to(torch.int16)
+        hi = ((w_sel >> 4) & 0x0F).to(torch.int16)
+        w_sel = torch.stack([lo, hi], dim=-1).reshape(*w_sel.shape[:-1], -1) - offset
+
     # Dequantize
     w_float = w_sel.to(x.dtype)
     s_expanded = s_sel.repeat_interleave(group_size, dim=-1)

backends/mlx/llm/switch.py

@@ -93,31 +93,71 @@ def pack(self):
         - Quantized: extracts inner tensors (qdata, scale, zero_point),
           stacks into [E, out, in_packed] buffers. Weight layout matches
           mlx::gather_qmm's expectations (transpose=True handles transposition).
+          4-bit qdata is nibble-packed to uint8 [E, out, in//2] (two values per
+          byte, value = signed q + 8) so it stays packed end-to-end — the
+          gather_qmm lowering reinterprets it as uint32 directly. 8-bit qdata
+          keeps the int8 [E, out, in] layout.
         - Unquantized: stacks weight.data into [E, out, in], then pretransposes
           to [E, in, out] so gather_mm receives the correct layout directly
           (no runtime transpose needed).
         """
         if self._packed:
             return
 
+        from executorch.extension.llm.export.int4 import ExportableInt4Tensor
+        from torchao.quantization import IntxUnpackedToInt8Tensor
+
         w0 = self.experts[0].weight
         self._is_quantized = hasattr(w0, "qdata")
 
-        if self._is_quantized:
-            _, metadata = w0.__tensor_flatten__()
-            self.group_size = metadata["block_size"][-1]
-
+        if isinstance(w0, ExportableInt4Tensor):
+            # Per-expert ExportableInt4Tensor: qdata is already nibble-packed
+            # uint8 (out, in//2); scale/zero_point are (in//gs, out) unsigned.
+            # Stack into the gather layout the gather_qmm handler expects
+            # (packed uint8 qdata + signed zero_point).
+            self.group_size = w0.group_size
             self.register_buffer(
-                "qdata",
-                torch.stack([e.weight.qdata for e in self.experts]),
+                "qdata", torch.stack([e.weight.qdata for e in self.experts])
             )
             self.register_buffer(
                 "scale",
-                torch.stack([e.weight.scale for e in self.experts]),
+                torch.stack([e.weight.scale.t().contiguous() for e in self.experts]),
             )
             self.register_buffer(
                 "zero_point",
-                torch.stack([e.weight.zero_point for e in self.experts]),
+                torch.stack(
+                    [
+                        (e.weight.zero_point.to(torch.int16) - 8)
+                        .t()
+                        .contiguous()
+                        .to(torch.int8)
+                        for e in self.experts
+                    ]
+                ),
+            )
+        elif isinstance(w0, IntxUnpackedToInt8Tensor):
+            _, metadata = w0.__tensor_flatten__()
+            self.group_size = metadata["block_size"][-1]
+
+            qdata = torch.stack([e.weight.qdata for e in self.experts])
+            scale = torch.stack([e.weight.scale for e in self.experts])
+            zero_point = torch.stack([e.weight.zero_point for e in self.experts])
+
+            # Nibble-pack 4-bit qdata into uint8 [E, out, in//2] so experts stay
+            # packed through export (half the storage; gather_qmm views as uint32).
+            if metadata.get("target_dtype") == torch.int4:
+                E, out, in_features = qdata.shape
+                qu = (qdata.to(torch.int16) + 8).to(torch.uint8)
+                qu = qu.reshape(E, out, in_features // 2, 2)
+                qdata = (qu[..., 0] | (qu[..., 1] << 4)).contiguous()
+
+            self.register_buffer("qdata", qdata)
+            self.register_buffer("scale", scale)
+            self.register_buffer("zero_point", zero_point)
+        elif self._is_quantized:
+            raise TypeError(
+                f"Unsupported quantized expert weight type: {type(w0).__name__}; "
+                "expected ExportableInt4Tensor or IntxUnpackedToInt8Tensor"
             )
         else:
             # Stack [E, out, in] then pretranspose to [E, in, out]

backends/mlx/ops.py

@@ -25,6 +25,7 @@
     emit_quantized_biases,
     emit_shape,
     parse_dequant_node,
+    regroup_affine_scales,
     to_mlx_qparams,
     torch_dtype_to_scalar_type,
 )
@@ -1849,7 +1850,12 @@ def _gather_qmm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     if biases_node is not None:
         zp_target, zp_data = P.get_placeholder_target_and_tensor(biases_node)
 
-    # Reshape 3D [E, out, in] to 2D for to_mlx_qparams, then reshape back
+    # Reshape 3D [E, out, in] to 2D for to_mlx_qparams, then reshape back.
+    # Packed int4 experts store qdata as uint8 [E, out, in//2] (two 4-bit values
+    # per byte); unpacked experts store int8 [E, out, in]. Detect via dtype and
+    # let to_mlx_qparams take the prepacked (view -> uint32) fast path so packed
+    # weights never expand to int8 during lowering.
+    prepacked = w_data.dtype == torch.uint8
     orig_shape = w_data.shape
     E, out_dim = orig_shape[0], orig_shape[1]
     w_2d = w_data.reshape(E * out_dim, -1)
@@ -1860,7 +1866,7 @@ def _gather_qmm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
         else torch.zeros_like(s_2d, dtype=torch.int8)
     )
 
-    Q, B = to_mlx_qparams(w_2d, s_2d, zp_2d, bits)
+    Q, B = to_mlx_qparams(w_2d, s_2d, zp_2d, bits, prepacked=prepacked)
     Q = Q.reshape(E, out_dim, -1)
     B = B.reshape(E, out_dim, -1)
 
@@ -4437,6 +4443,12 @@ def _dequantize_affine_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     scale_2d = scale.reshape(-1, scale.shape[-1])
     zero_point_2d = zero_point.reshape(-1, zero_point.shape[-1])
 
+    # torchao may quantize with a coarser group_size than MLX supports; repeat
+    # scale/zero_point to the MLX-legal group_size (returned by parse_dequant_node).
+    scale_2d, zero_point_2d, _ = regroup_affine_scales(
+        scale_2d, zero_point_2d, qdata_2d.shape[-1], group_size
+    )
+
     Q, B = to_mlx_qparams(qdata_2d, scale_2d, zero_point_2d, bits)
 
     leading_dims = permuted_shape[:-1]