MLX Q5K/Q6K upgrades (#20827)

This PR adds an MLX-native repack lowering path for GGUF Q5_K and Q6_K
weights, mirroring the existing Q4_K path: at export time the raw GGUF
blob is unpacked and repacked into MLX affine qparams (bits=5 / bits=6)
consumed by MLX's native QuantizedMatmul/quantized-gather kernels, gated
by ET_MLX_EMIT_DIRECT_GGUF (default = native, falling back to the fused
Metal kernels otherwise). It also introduces a lossless group-size
upgrade in to_intx_unpacked_to_int8_tensor (_max_uniform_group_size)
that merges adjacent sub-blocks with identical scale/min into the
largest MLX-supported group size (32→64→128) when bit-exact, which the
repack paths request via max_group_size=128; because Q6_K's native group
size (16) isn't MLX-compatible, its native path only activates when this
merge reaches ≥32 and otherwise transparently falls back to fused
kernels. Q6_K's symmetric zero-point lets biases be computed in the init
chain via emit_quantized_biases instead of being serialized, and a
FakeTensor guard keeps the data-dependent merge from tripping the
partitioner's support check. Coverage includes new unit tests for the
merge logic in test_gguf.py and native/fused + group-merged op tests
across Q4_K/Q5_K/Q6_K in test_linear.py and test_embedding.py.

Author: metascroy

backends/mlx/builder/op_helpers.py

@@ -24,6 +24,12 @@
 # computing biases = -scales * 2^(bits-1) during the init chain.
 QUANTIZED_SERIALIZE_BIASES = False
 
+# Row-chunk size (in elements) for the int64 bit-packer in ``to_mlx_qparams``.
+# Packing widths that don't divide 32 (5/6-bit) needs int64 scratch; chunking the
+# rows bounds that scratch to ~this many elements instead of the whole weight
+# (a full lm_head would otherwise be tens of GB of int64 -> OOM).
+_PACK_CHUNK_ELEMS = 1 << 24  # ~16M int64 elements (~128 MB per scratch tensor)
+
 
 def get_aten_target(target):
     """
@@ -555,27 +561,43 @@ def to_mlx_qparams(
         Q = q.contiguous().view(torch.uint32).reshape(rows, -1)
     else:
         # Contiguous LSB-first bit-packing for widths that don't divide 32
-        # (e.g. 6-bit), matching MLX's affine pack_and_quantize (ops.cpp).
+        # (e.g. 5/6-bit), matching MLX's affine pack_and_quantize (ops.cpp).
+        #
+        # We scatter each column's value directly into its uint32 word(s) rather
+        # than expanding to a per-bit stream. Column j occupies global bits
+        # [j*bits, (j+1)*bits): word j*bits//32 at shift j*bits%32, plus a carry
+        # into the next word when it straddles the boundary (at most one carry
+        # since bits <= 32). Column bit-ranges within a word are disjoint, so
+        # index_add_ (sum) is equivalent to OR.
         #
-        # We scatter each column's value directly into its uint32 word(s)
-        # rather than expanding to a per-bit stream, which would materialize an
-        # int64 (rows, cols, bits) tensor (tens of GB for lm_head -> OOM).
-        # Column j occupies global bits [j*bits, (j+1)*bits): word j*bits//32 at
-        # shift j*bits%32, plus a carry into the next word when it straddles the
-        # boundary (at most one carry since bits <= 32). Column bit-ranges within
-        # a word are disjoint, so index_add_ (sum) is equivalent to OR.
+        # The scatter needs int64 (a value shifted by up to 31 overflows int32),
+        # so we pack the rows in chunks to bound the peak int64 working set --
+        # packing a full lm_head in one shot would otherwise materialize a
+        # multi-GB int64 tensor.
         n_words = cols * bits // 32
-        q = qdata.to(torch.int64) + offset  # 0 .. 2**bits-1
         pos = torch.arange(cols, dtype=torch.int64) * bits
         word = pos // 32
         shift = pos % 32
-        packed = torch.zeros(rows, n_words, dtype=torch.int64)
-        packed.index_add_(1, word, q << shift)  # low bits (+ overflow)
         straddle = (shift + bits) > 32
-        if bool(straddle.any()):
-            carry = torch.where(straddle, q >> (32 - shift), torch.zeros_like(q))
-            packed.index_add_(1, (word + 1).clamp(max=n_words - 1), carry)
-        Q = (packed & 0xFFFFFFFF).to(torch.int32).contiguous().view(torch.uint32)
+        has_straddle = bool(straddle.any())
+        word_carry = (word + 1).clamp(max=n_words - 1) if has_straddle else None
+
+        rows_per_chunk = max(1, _PACK_CHUNK_ELEMS // cols)
+        chunks = []
+        for r0 in range(0, rows, rows_per_chunk):
+            q = qdata[r0 : r0 + rows_per_chunk].to(torch.int64) + offset
+            packed = torch.zeros(q.shape[0], n_words, dtype=torch.int64)
+            packed.index_add_(1, word, q << shift)  # low bits (+ overflow)
+            if has_straddle:
+                carry = torch.where(straddle, q >> (32 - shift), torch.zeros_like(q))
+                packed.index_add_(1, word_carry, carry)
+                del carry
+            del q
+            chunks.append((packed & 0xFFFFFFFF).to(torch.int32))
+            del packed
+        packed_i32 = chunks[0] if len(chunks) == 1 else torch.cat(chunks, dim=0)
+        del chunks
+        Q = packed_i32.contiguous().view(torch.uint32)
 
     if compute_biases:
         B = -scale * (zero_point.to(scale.dtype) + offset)

backends/mlx/custom_kernel_ops/gguf/patterns.py

@@ -17,10 +17,11 @@
 These handlers match that ``dequantize_gguf -> linear/embedding`` subgraph and
 lower it without materializing the dequantized weight:
 
-* **Q6_K** -> fused custom Metal kernels in :mod:`.q6k`.
-* **Q5_K** -> fused custom Metal kernels in :mod:`.q5k`.
-* **Q4_K** -> fused custom Metal kernels in :mod:`.q4k` (default), or the legacy
-  MLX-native repack path when ``ET_MLX_EMIT_DIRECT_GGUF=0``.
+* **Q4_K** / **Q5_K** -> MLX-native repack path by default; the fused custom
+  Metal kernels in :mod:`.q4k` / :mod:`.q5k` when ``ET_MLX_EMIT_DIRECT_GGUF=1``.
+* **Q6_K** -> MLX-native repack path when its sub-blocks merge to an
+  MLX-supported group size (>= 32), else the fused kernels in :mod:`.q6k`;
+  ``ET_MLX_EMIT_DIRECT_GGUF=1`` forces fused.
 
 All cover linear and embedding.
 
@@ -43,7 +44,8 @@
 from torch.export.exported_program import ExportedProgram
 from torch.fx.node import Node
 
-# Quant types each pattern can lower (both via fused custom Metal kernels).
+# Quant types each pattern can lower (MLX-native repack by default, or fused
+# custom Metal kernels via ``ET_MLX_EMIT_DIRECT_GGUF=1``).
 _LINEAR_TYPES = {"q4_k", "q5_k", "q6_k"}
 _EMBEDDING_TYPES = {"q4_k", "q5_k", "q6_k"}
 

backends/mlx/custom_kernel_ops/gguf/q4k/__init__.py

@@ -12,7 +12,7 @@
 ``custom_kernel_ops.gguf.patterns``); they are not imported here to keep the
 package import light.
 
-By default the legacy export-time repack path
+By default the export-time repack path
 (:mod:`.linear_mlx_native` / :mod:`.embedding_mlx_native`) is used. Set
 ``ET_MLX_EMIT_DIRECT_GGUF=1`` to emit the fused Metal kernels that read raw
 GGUF bytes instead.
@@ -26,7 +26,7 @@
 def emit_direct_gguf() -> bool:
     """Return True to emit fused kernels that read raw GGUF bytes.
 
-    Defaults to False (the legacy MLX-native repack path); set
+    Defaults to False (the MLX-native repack path); set
     ``ET_MLX_EMIT_DIRECT_GGUF=1`` to enable the fused kernels.
     """
     return os.environ.get("ET_MLX_EMIT_DIRECT_GGUF", "0") != "0"

backends/mlx/custom_kernel_ops/gguf/q4k/embedding.py

@@ -134,7 +134,7 @@ def emit_embedding(
     indices_node: Node,
     output_dtype: torch.dtype,
 ) -> Slot:
-    """Dispatch to fused Metal gather or the legacy MLX-native repack path."""
+    """Dispatch to fused Metal gather or the MLX-native repack path."""
     from executorch.backends.mlx.custom_kernel_ops.gguf.q4k import emit_direct_gguf
 
     if emit_direct_gguf():

backends/mlx/custom_kernel_ops/gguf/q4k/linear.py

@@ -487,7 +487,7 @@ def emit_linear(
     weight_node: Node,
     bias_node: Optional[Node],
 ) -> Slot:
-    """Dispatch to fused Metal kernels or the legacy MLX-native repack path."""
+    """Dispatch to fused Metal kernels or the MLX-native repack path."""
     from executorch.backends.mlx.custom_kernel_ops.gguf.q4k import emit_direct_gguf
 
     if emit_direct_gguf():

backends/mlx/custom_kernel_ops/gguf/q4k/linear_mlx_native.py

@@ -9,8 +9,9 @@
 """GGUF **Q4_K** linear lowering via MLX's native 4-bit quantized matmul.
 
 Lowers a ``dequantize_gguf -> linear`` pattern to a ``QuantizedMatmulNode``
-(mode "affine", group_size 32); the GGUF blob is repacked into MLX qparams at
-export time (see :mod:`.repack_mlx`).
+(mode "affine"); the GGUF blob is repacked into MLX qparams at export time (see
+:mod:`.repack_mlx`). The group size is 32 by default, or the merged 64/128 when
+adjacent sub-blocks are identical.
 """
 
 from __future__ import annotations

backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py

@@ -6,7 +6,7 @@
 # LICENSE file in the root directory of this source tree.
 #
 
-"""Q4_K -> MLX qparam repack for the legacy MLX-native lowering path.
+"""Q4_K -> MLX qparam repack for the MLX-native lowering path.
 
 Used when ``ET_MLX_EMIT_DIRECT_GGUF=0``: the GGUF blob is unpacked and repacked
 into MLX affine 4-bit qparams at export time instead of being consumed directly
@@ -28,16 +28,23 @@
 def repack_mlx(P: MLXProgramBuilder, weight_node: Node) -> Tuple[Slot, Slot, Slot, int]:
     """Unpack a raw Q4_K blob and repack into MLX qparam constants.
 
+    Adjacent sub-blocks with identical scale/min are merged into a larger group
+    size (up to 128) when lossless, so ``group_size`` may be 32, 64, or 128.
     Returns ``(packed_slot, scales_slot, biases_slot, group_size)``.
     """
     from executorch.extension.llm.export.gguf import ExportableGGUFTensor
 
     weight_target, raw = P.get_placeholder_target_and_tensor(weight_node)
-    intx = ExportableGGUFTensor.from_raw(raw, "q4_k").to_intx_unpacked_to_int8_tensor()
+    intx = ExportableGGUFTensor.from_raw(raw, "q4_k").to_intx_unpacked_to_int8_tensor(
+        max_group_size=128
+    )
     group_size = int(intx.block_size[-1])
-    packed, biases = to_mlx_qparams(intx.qdata, intx.scale, intx.zero_point, _BITS)
+    qdata, scale, zero_point = intx.qdata, intx.scale, intx.zero_point
+    del intx  # drop the tensor-subclass wrapper; keep only the fields we need
+    packed, biases = to_mlx_qparams(qdata, scale, zero_point, _BITS)
+    del qdata  # the (N, K) int8 is no longer needed once packed
 
     packed_slot = P.make_or_get_constant(f"{weight_target}_q4k_packed", packed)
-    scales_slot = P.make_or_get_constant(f"{weight_target}_q4k_scales", intx.scale)
+    scales_slot = P.make_or_get_constant(f"{weight_target}_q4k_scales", scale)
     biases_slot = P.make_or_get_constant(f"{weight_target}_q4k_biases", biases)
     return packed_slot, scales_slot, biases_slot, group_size

backends/mlx/custom_kernel_ops/gguf/q5k/__init__.py

@@ -6,16 +6,32 @@
 # LICENSE file in the root directory of this source tree.
 #
 
-"""GGUF Q5_K format implementation (fused custom Metal kernels).
+"""GGUF Q5_K format lowering for the MLX backend.
 
 Re-exports the lightweight constants/header from :mod:`.common` so they can be
 imported without pulling in the MLX builder. The ``emit_*`` lowerings live in
 :mod:`.linear` / :mod:`.embedding` (called by ``custom_kernel_ops.gguf.patterns``)
 and are not imported here.
+
+By default the export-time repack path
+(:mod:`.linear_mlx_native` / :mod:`.embedding_mlx_native`) is used. Set
+``ET_MLX_EMIT_DIRECT_GGUF=1`` to emit the fused Metal kernels that read raw
+GGUF bytes instead.
 """
 
+import os
+
 from executorch.backends.mlx.custom_kernel_ops.gguf.q5k.common import (  # noqa: F401
     _Q5K_HEADER,
     Q5K_BLOCK_BYTES,
     QK_K,
 )
+
+
+def emit_direct_gguf() -> bool:
+    """Return True to emit fused kernels that read raw GGUF bytes.
+
+    Defaults to False (the MLX-native repack path); set
+    ``ET_MLX_EMIT_DIRECT_GGUF=1`` to enable the fused kernels.
+    """
+    return os.environ.get("ET_MLX_EMIT_DIRECT_GGUF", "0") != "0"

backends/mlx/custom_kernel_ops/gguf/q5k/embedding.py

@@ -55,7 +55,7 @@
 """
 
 
-def emit_embedding(
+def _emit_embedding_fused(
     P: MLXProgramBuilder,
     head: Node,
     weight_node: Node,
@@ -122,3 +122,23 @@ def emit_embedding(
     )
 
     return out
+
+
+def emit_embedding(
+    P: MLXProgramBuilder,
+    head: Node,
+    weight_node: Node,
+    indices_node: Node,
+    output_dtype: torch.dtype,
+) -> Slot:
+    """Dispatch to fused Metal gather or the MLX-native repack path."""
+    from executorch.backends.mlx.custom_kernel_ops.gguf.q5k import emit_direct_gguf
+
+    if emit_direct_gguf():
+        return _emit_embedding_fused(P, head, weight_node, indices_node, output_dtype)
+
+    from executorch.backends.mlx.custom_kernel_ops.gguf.q5k.embedding_mlx_native import (
+        emit_embedding as emit_embedding_mlx_native,
+    )
+
+    return emit_embedding_mlx_native(P, head, weight_node, indices_node, output_dtype)

backends/mlx/custom_kernel_ops/gguf/q5k/embedding_mlx_native.py

@@ -0,0 +1,58 @@
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+"""GGUF **Q5_K** embedding lowering via MLX's native 5-bit quantized gather.
+
+Lowers a ``dequantize_gguf -> embedding`` pattern to a quantized gather: gather
+the packed quants / scales / biases by index, then dequantize the gathered rows
+(``DequantizeNode``, mode "affine"). The GGUF blob is repacked into MLX qparams
+at export time (see :mod:`.repack_mlx`).
+"""
+
+from __future__ import annotations
+
+from executorch.backends.mlx.builder.op_helpers import emit_quantized_gather
+from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder
+from executorch.backends.mlx.builder.slot_manager import Slot
+from executorch.backends.mlx.custom_kernel_ops.gguf.q5k.repack_mlx import (
+    _BITS,
+    repack_mlx,
+)
+from torch.fx.node import Node
+
+
+def emit_embedding(
+    P: MLXProgramBuilder,
+    head: Node,
+    weight_node: Node,
+    indices_node: Node,
+    output_dtype,
+) -> Slot:
+    """Lower a Q5_K ``dequantize_gguf -> embedding`` pattern to a quantized gather.
+
+    Gathers the packed quants / scales / biases by index, then dequantizes the
+    gathered rows (MLX affine 5-bit) -- the same shape as MLX's generic quantized
+    embedding.
+    """
+    w_slot, scales_slot, biases_slot, group_size = repack_mlx(P, weight_node)
+    (indices_slot,) = P.slot_map([indices_node])
+
+    out = P.make_or_get_slot(head)
+    emit_quantized_gather(
+        P,
+        out,
+        indices_slot,
+        w_slot,
+        scales_slot,
+        biases_slot,
+        group_size=group_size,
+        bits=_BITS,
+        mode="affine",
+        out_dtype=output_dtype,
+    )
+    return out