Support FP16 GGUF export (#20968) (#20968)

metascroy · web-flow · commit e16c0484a1cd · 2026-07-15T16:06:20.000-07:00
Summary:

Adds fp16 support to the MLX GGUF export path. Previously activations,
the
KV cache, and unquantized weights were always bf16; this adds an
`--activation-dtype` option (bfloat16 | float16 | float) so a model can
be
exported for fp16 (or fp32) compute.

Key changes:
- The export threads an activation dtype through the pipeline:
unquantized
weights, buffers, and the KV cache are converted to the chosen dtype,
and the
  dtype is recorded in the .pte.
- Quantized-matmul scales now follow the activation dtype
(`scale_dtype`), so
  MLX does not promote to float32 (bf16 scales against fp16 activations
previously forced float32 and broke kernel compilation). Plumbed through
the
q4k/q5k/q6k MLX repack paths; CUDA and other callers pass bfloat16
explicitly.
- `iter_gguf` now returns F16 tensors in their native fp16 dtype instead
of
silently casting to bf16; call sites that require bf16 cast explicitly.

bf16 remains the default, so existing exports are unchanged.

Reviewed By: Gasoonjia

Differential Revision: D111968227
diff --git a/backends/cuda/dp4a_planar_int5_tensor.py b/backends/cuda/dp4a_planar_int5_tensor.py
@@ -286,7 +286,9 @@ def from_exportable_gguf(cls, gt) -> "CudaDp4aPlanarInt5Tensor":
                 "CudaDp4aPlanarInt5Tensor.from_exportable_gguf requires a q5_k "
                 f"ExportableGGUFTensor, got {gt.ggml_type!r}"
             )
-        return cls._from_intx_int8(gt.to_intx_unpacked_to_int8_tensor())
+        return cls._from_intx_int8(
+            gt.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16)
+        )
 
     def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor:
         """Dequantize to a dense tensor (asymmetric: ``w = scale * (u - zero)``).
diff --git a/backends/cuda/dp4a_planar_int6_tensor.py b/backends/cuda/dp4a_planar_int6_tensor.py
@@ -278,7 +278,9 @@ def from_exportable_gguf(cls, gt) -> "CudaDp4aPlanarInt6Tensor":
                 "CudaDp4aPlanarInt6Tensor.from_exportable_gguf requires a q6_k "
                 f"ExportableGGUFTensor, got {gt.ggml_type!r}"
             )
-        return cls._from_intx_int8(gt.to_intx_unpacked_to_int8_tensor())
+        return cls._from_intx_int8(
+            gt.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16)
+        )
 
     def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor:
         """Dequantize to a dense tensor (symmetric: ``w = q * scale``).
diff --git a/backends/cuda/tests/test_int6_dispatch.py b/backends/cuda/tests/test_int6_dispatch.py
@@ -258,7 +258,7 @@ def test_from_exportable_gguf(self):
         # (scale = code * step[:, g // (256 // gs)]); the GGUF sub-scales are
         # constant within this single super-block, so the int8 re-encoding is
         # exact here.
-        intx = gt.to_intx_unpacked_to_int8_tensor()
+        intx = gt.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16)
         q_rt = unpack_int6(t.ql, t.qh, N, nb * 256).to(torch.int8)
         self.assertTrue(torch.equal(q_rt, intx.qdata))
         n_groups = intx.scale.shape[1]
diff --git a/backends/mlx/custom_kernel_ops/gguf/q4k/embedding_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q4k/embedding_mlx_native.py
@@ -39,7 +39,9 @@ def emit_embedding(
     gathered rows (MLX affine 4-bit) -- the same shape as MLX's generic quantized
     embedding.
     """
-    w_slot, scales_slot, biases_slot, group_size = repack_mlx(P, weight_node)
+    w_slot, scales_slot, biases_slot, group_size = repack_mlx(
+        P, weight_node, scale_dtype=output_dtype
+    )
     (indices_slot,) = P.slot_map([indices_node])
 
     out = P.make_or_get_slot(head)
diff --git a/backends/mlx/custom_kernel_ops/gguf/q4k/linear_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q4k/linear_mlx_native.py
@@ -46,7 +46,9 @@ def emit_linear(
     node. The blob is repacked into MLX qparams at export time, so only the
     MLX-format constants are serialized.
     """
-    w_slot, scales_slot, biases_slot, group_size = repack_mlx(P, weight_node)
+    w_slot, scales_slot, biases_slot, group_size = repack_mlx(
+        P, weight_node, scale_dtype=x_node.meta["val"].dtype
+    )
     x_slot, bias_slot = P.slot_map([x_node, bias_node])
 
     out = P.make_or_get_slot(head)
diff --git a/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py
@@ -15,7 +15,9 @@
 
 from __future__ import annotations
 
-from typing import Tuple
+from typing import Optional, Tuple
+
+import torch
 
 from executorch.backends.mlx.builder.op_helpers import to_mlx_qparams
 from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder
@@ -25,18 +27,26 @@
 _BITS = 4
 
 
-def repack_mlx(P: MLXProgramBuilder, weight_node: Node) -> Tuple[Slot, Slot, Slot, int]:
+def repack_mlx(
+    P: MLXProgramBuilder,
+    weight_node: Node,
+    scale_dtype: Optional[torch.dtype] = None,
+) -> 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)``.
+
+    ``scale_dtype`` sets the dtype of the emitted scales/biases constants; pass
+    the activation dtype so MLX ``quantized_matmul`` does not promote (a bf16
+    activation with f16 scales, or vice versa, promotes to float32).
     """
     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(
-        max_group_size=128
+        max_group_size=128, scale_dtype=scale_dtype
     )
     group_size = int(intx.block_size[-1])
     qdata, scale, zero_point = intx.qdata, intx.scale, intx.zero_point
diff --git a/backends/mlx/custom_kernel_ops/gguf/q5k/embedding_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q5k/embedding_mlx_native.py
@@ -39,7 +39,9 @@ def emit_embedding(
     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)
+    w_slot, scales_slot, biases_slot, group_size = repack_mlx(
+        P, weight_node, scale_dtype=output_dtype
+    )
     (indices_slot,) = P.slot_map([indices_node])
 
     out = P.make_or_get_slot(head)
diff --git a/backends/mlx/custom_kernel_ops/gguf/q5k/linear_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q5k/linear_mlx_native.py
@@ -46,7 +46,9 @@ def emit_linear(
     node. The blob is repacked into MLX qparams at export time, so only the
     MLX-format constants are serialized.
     """
-    w_slot, scales_slot, biases_slot, group_size = repack_mlx(P, weight_node)
+    w_slot, scales_slot, biases_slot, group_size = repack_mlx(
+        P, weight_node, scale_dtype=x_node.meta["val"].dtype
+    )
     x_slot, bias_slot = P.slot_map([x_node, bias_node])
 
     out = P.make_or_get_slot(head)
diff --git a/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py
@@ -15,7 +15,9 @@
 
 from __future__ import annotations
 
-from typing import Tuple
+from typing import Optional, Tuple
+
+import torch
 
 from executorch.backends.mlx.builder.op_helpers import to_mlx_qparams
 from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder
@@ -25,18 +27,26 @@
 _BITS = 5
 
 
-def repack_mlx(P: MLXProgramBuilder, weight_node: Node) -> Tuple[Slot, Slot, Slot, int]:
+def repack_mlx(
+    P: MLXProgramBuilder,
+    weight_node: Node,
+    scale_dtype: Optional[torch.dtype] = None,
+) -> Tuple[Slot, Slot, Slot, int]:
     """Unpack a raw Q5_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)``.
+
+    ``scale_dtype`` sets the dtype of the emitted scales/biases constants; pass
+    the activation dtype so MLX ``quantized_matmul`` does not promote (a bf16
+    activation with f16 scales, or vice versa, promotes to float32).
     """
     from executorch.extension.llm.export.gguf import ExportableGGUFTensor
 
     weight_target, raw = P.get_placeholder_target_and_tensor(weight_node)
     intx = ExportableGGUFTensor.from_raw(raw, "q5_k").to_intx_unpacked_to_int8_tensor(
-        max_group_size=128
+        max_group_size=128, scale_dtype=scale_dtype
     )
     group_size = int(intx.block_size[-1])
     qdata, scale, zero_point = intx.qdata, intx.scale, intx.zero_point
diff --git a/backends/mlx/custom_kernel_ops/gguf/q6k/embedding_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q6k/embedding_mlx_native.py
@@ -42,7 +42,7 @@ def emit_embedding(
     Returns the output slot, or ``None`` when the weight does not merge to an
     MLX-supported group size (the caller should fall back to the fused gather).
     """
-    repacked = repack_mlx(P, weight_node)
+    repacked = repack_mlx(P, weight_node, scale_dtype=output_dtype)
     if repacked is None:
         return None
     w_slot, scales_slot, biases_slot, group_size = repacked
diff --git a/backends/mlx/custom_kernel_ops/gguf/q6k/linear_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q6k/linear_mlx_native.py
@@ -46,7 +46,7 @@ def emit_linear(
     Returns the output slot, or ``None`` when the weight does not merge to an
     MLX-supported group size (the caller should fall back to fused kernels).
     """
-    repacked = repack_mlx(P, weight_node)
+    repacked = repack_mlx(P, weight_node, scale_dtype=x_node.meta["val"].dtype)
     if repacked is None:
         return None
     w_slot, scales_slot, biases_slot, group_size = repacked
diff --git a/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py
@@ -24,6 +24,8 @@
 
 from typing import Optional, Tuple
 
+import torch
+
 from executorch.backends.mlx.builder.op_helpers import (
     emit_quantized_biases,
     to_mlx_qparams,
@@ -39,20 +41,26 @@
 
 
 def repack_mlx(
-    P: MLXProgramBuilder, weight_node: Node
+    P: MLXProgramBuilder,
+    weight_node: Node,
+    scale_dtype: Optional[torch.dtype] = None,
 ) -> Optional[Tuple[Slot, Slot, Slot, int]]:
     """Unpack a raw Q6_K blob and repack into MLX qparam constants.
 
     Adjacent sub-blocks with identical scale are merged into a larger group size
     (up to 128) when lossless. Returns ``(packed_slot, scales_slot, biases_slot,
     group_size)`` when the merged ``group_size`` is MLX-compatible (>= 32), or
     ``None`` when it is not (so the caller can fall back to fused kernels).
+
+    ``scale_dtype`` sets the dtype of the emitted scales/biases constants; pass
+    the activation dtype so MLX ``quantized_matmul`` does not promote (a bf16
+    activation with f16 scales, or vice versa, promotes to float32).
     """
     from executorch.extension.llm.export.gguf import ExportableGGUFTensor
 
     weight_target, raw = P.get_placeholder_target_and_tensor(weight_node)
     intx = ExportableGGUFTensor.from_raw(raw, "q6_k").to_intx_unpacked_to_int8_tensor(
-        max_group_size=128
+        max_group_size=128, scale_dtype=scale_dtype
     )
     group_size = int(intx.block_size[-1])
     if group_size < _MIN_MLX_GROUP_SIZE:
diff --git a/backends/mlx/custom_kernel_ops/gguf/test/test_embedding.py b/backends/mlx/custom_kernel_ops/gguf/test/test_embedding.py
@@ -199,9 +199,9 @@ def _native_fires(self) -> bool:
     def compute_expected_outputs(self, model, test_inputs):
         if self._native_fires():
             if model.weight.ggml_type == "q4_k":
-                weight = _q4k_mlx_native_dequant(model.weight)
+                weight = _q4k_mlx_native_dequant(model.weight, model.weight.orig_dtype)
             else:
-                weight = _mlx_native_dequant(model.weight)
+                weight = _mlx_native_dequant(model.weight, model.weight.orig_dtype)
             out = torch.nn.functional.embedding(test_inputs[0], weight)
             return [out.to(model.weight.orig_dtype)]
         return model(*test_inputs)
diff --git a/backends/mlx/custom_kernel_ops/gguf/test/test_linear.py b/backends/mlx/custom_kernel_ops/gguf/test/test_linear.py
@@ -207,10 +207,10 @@ def _fp32_linear_reference(model: "GGUFLinearModel", x: torch.Tensor):
     return [out.to(x.dtype)]
 
 
-def _q4k_mlx_native_dequant(weight) -> torch.Tensor:
+def _q4k_mlx_native_dequant(weight, scale_dtype: torch.dtype) -> torch.Tensor:
     from executorch.backends.mlx.builder.op_helpers import to_mlx_qparams
 
-    intx = weight.to_intx_unpacked_to_int8_tensor()
+    intx = weight.to_intx_unpacked_to_int8_tensor(scale_dtype=scale_dtype)
     group_size = int(intx.block_size[-1])
     packed, biases = to_mlx_qparams(intx.qdata, intx.scale, intx.zero_point, 4)
     packed_bytes = packed.view(torch.uint8)
@@ -226,15 +226,18 @@ def _q4k_mlx_native_dequant(weight) -> torch.Tensor:
 _GGML_BITS = {"q4_k": 4, "q5_k": 5, "q6_k": 6}
 
 
-def _mlx_native_dequant(weight) -> torch.Tensor:
+def _mlx_native_dequant(weight, scale_dtype: torch.dtype) -> torch.Tensor:
     """Reference for the MLX-native repack path (q4_k / q5_k / q6_k).
 
     Mirrors the export-time repack (``to_intx_unpacked_to_int8_tensor`` with
-    ``max_group_size=128``) and the MLX affine kernel, which computes
-    ``scale*Q + bias == scale*(qdata - zero_point)``. Invariant to lossless
-    group merging, so it validates the group-size upgrade too.
+    ``max_group_size=128`` and the activation ``scale_dtype``) and the MLX affine
+    kernel, which computes ``scale*Q + bias == scale*(qdata - zero_point)``.
+    Invariant to lossless group merging, so it validates the group-size upgrade
+    too.
     """
-    intx = weight.to_intx_unpacked_to_int8_tensor(max_group_size=128)
+    intx = weight.to_intx_unpacked_to_int8_tensor(
+        max_group_size=128, scale_dtype=scale_dtype
+    )
     gs = int(intx.block_size[-1])
     q = intx.qdata.float()
     scale = intx.scale.float().repeat_interleave(gs, dim=1)
@@ -246,9 +249,9 @@ def _fp32_linear_mlx_native_reference(model: "GGUFLinearModel", x: torch.Tensor)
     lin = model.linear
     # q4_k keeps its packed-nibble oracle; q5_k/q6_k use the generic one.
     if lin.weight.ggml_type == "q4_k":
-        w = _q4k_mlx_native_dequant(lin.weight)
+        w = _q4k_mlx_native_dequant(lin.weight, x.dtype)
     else:
-        w = _mlx_native_dequant(lin.weight)
+        w = _mlx_native_dequant(lin.weight, x.dtype)
     bias = lin.bias.float() if lin.bias is not None else None
     out = torch.nn.functional.linear(x.float(), w, bias)
     return [out.to(x.dtype)]
diff --git a/examples/models/gemma4_31b/gguf_loader.py b/examples/models/gemma4_31b/gguf_loader.py
@@ -148,7 +148,7 @@ def _untie_embed_lm_head(model, gtensor, weight, backend):
         return weight, gtensor
 
     if gtensor.ggml_type in ("q6_k", "q4_k"):
-        intx = gtensor.to_intx_unpacked_to_int8_tensor()
+        intx = gtensor.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16)
         if gtensor.ggml_type == "q6_k":
             import torch.nn as nn
             from executorch.backends.cuda.dp4a_planar_int6_tensor import (
@@ -229,7 +229,7 @@ def load_gguf_model(
                     model, value, weight, backend
                 )
             value = weight
-        elif value.dtype == torch.float32:
+        elif value.dtype in (torch.float32, torch.float16):
             value = value.to(torch.bfloat16)
 
         pack_one(model, model_key, value, packers)
diff --git a/examples/models/gemma4_31b/quant/tests/test_pack_cuda.py b/examples/models/gemma4_31b/quant/tests/test_pack_cuda.py
@@ -228,7 +228,7 @@ def test_pack_linear_real_int8_passthrough(self):
     def test_matmul_correct(self):
         _require_cuda(self)
         gt = self._make_q6k_gguf(256, nb=1)  # (256, 256)
-        intx = gt.to_intx_unpacked_to_int8_tensor()
+        intx = gt.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16)
         q, scale = intx.qdata, intx.scale
         module = nn.Linear(256, 256, bias=False)
         pack_linear_for_cuda(module, {"weight": gt})
@@ -267,7 +267,7 @@ def test_e2e_q6k_export_lower_decode(self):
 
         gt = self._make_q6k_gguf(64, nb=1)  # (64, 256)
         # Reference: the shared Q6_K int8 decode, dequantized (w = q * scale).
-        intx = gt.to_intx_unpacked_to_int8_tensor()
+        intx = gt.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16)
         w_ref = (
             intx.qdata.to(torch.bfloat16)
             * intx.scale.to(torch.bfloat16).repeat_interleave(16, dim=-1)
diff --git a/extension/llm/export/gguf.py b/extension/llm/export/gguf.py
@@ -384,7 +384,9 @@ def to_int4_tensor(self) -> Tensor:
         )
 
     def to_intx_unpacked_to_int8_tensor(
-        self, max_group_size: Optional[int] = None
+        self,
+        max_group_size: Optional[int] = None,
+        scale_dtype: Optional[torch.dtype] = None,
     ) -> Tensor:
         """Convert to a torchao ``IntxUnpackedToInt8Tensor`` (Q4_K, Q5_K or Q6_K).
 
@@ -403,6 +405,13 @@ def to_intx_unpacked_to_int8_tensor(
         """
         from torchao.quantization import IntxUnpackedToInt8Tensor
 
+        # Scale/zero-point dtype. GGUF stores the super-block scale as float16,
+        # so float16 is the faithful default. Callers pass their own compute
+        # dtype: the MLX repack passes the activation dtype so quantized_matmul
+        # scales match activations and MLX does not promote (e.g. bf16 + f16 ->
+        # f32); the CUDA integrations pass bfloat16.
+        sdt = scale_dtype if scale_dtype is not None else torch.float16
+
         N, K = int(self.shape[0]), int(self.shape[1])
         if self.ggml_type == "q6_k":
             q, eff_scale = _q6_k_fields(self.raw, N, K)
@@ -418,11 +427,11 @@ def to_intx_unpacked_to_int8_tensor(
                 )
             return IntxUnpackedToInt8Tensor(
                 qdata=q,
-                scale=eff_scale.to(torch.bfloat16),
+                scale=eff_scale.to(sdt),
                 zero_point=torch.zeros_like(eff_scale, dtype=torch.int8),
                 target_dtype=torch.int8,
                 block_size=(1, group_size),
-                dtype=torch.bfloat16,
+                dtype=sdt,
                 activation_quantization=None,
             )
         if self.ggml_type == "q5_k":
@@ -439,11 +448,11 @@ def to_intx_unpacked_to_int8_tensor(
             # match (dequant = scale * (q - zp) is preserved).
             return IntxUnpackedToInt8Tensor(
                 qdata=q.to(torch.int8) - 16,
-                scale=eff_scale.to(torch.bfloat16),
-                zero_point=(zero - 16).to(torch.bfloat16),
+                scale=eff_scale.to(sdt),
+                zero_point=(zero - 16).to(sdt),
                 target_dtype=torch.int5,
                 block_size=(1, group_size),
-                dtype=torch.bfloat16,
+                dtype=sdt,
                 activation_quantization=None,
             )
         if self.ggml_type == "q4_k":
@@ -460,11 +469,11 @@ def to_intx_unpacked_to_int8_tensor(
             # (dequant = scale * (q - zp) is preserved).
             return IntxUnpackedToInt8Tensor(
                 qdata=q.to(torch.int8) - 8,
-                scale=eff_scale.to(torch.bfloat16),
-                zero_point=(zero - 8).to(torch.bfloat16),
+                scale=eff_scale.to(sdt),
+                zero_point=(zero - 8).to(sdt),
                 target_dtype=torch.int4,
                 block_size=(1, group_size),
-                dtype=torch.bfloat16,
+                dtype=sdt,
                 activation_quantization=None,
             )
         raise NotImplementedError(
@@ -517,8 +526,8 @@ def iter_gguf(
     """Stream ``(name, value)`` for every tensor in a GGUF file (low peak mem).
 
     Quantized tensors (Q4_K, Q5_K, Q6_K) are wrapped as ``ExportableGGUFTensor``
-    with the raw block bytes; F32/F16 are returned as plain float tensors (bf16
-    for F16). GGUF shapes are reversed to PyTorch ``(N, K)`` convention.
+    with the raw block bytes; F32/F16/BF16 are returned as plain float tensors in
+    their native dtype. GGUF shapes are reversed to PyTorch ``(N, K)`` convention.
     """
     from gguf import GGMLQuantizationType, GGUFReader
 
@@ -538,10 +547,7 @@ def iter_gguf(
         elif tensor.tensor_type == GGMLQuantizationType.F32:
             yield tensor.name, flat.view(torch.float32).reshape(shape).clone()
         elif tensor.tensor_type == GGMLQuantizationType.F16:
-            yield (
-                tensor.name,
-                flat.view(torch.float16).reshape(shape).to(torch.bfloat16),
-            )
+            yield tensor.name, flat.view(torch.float16).reshape(shape).clone()
         elif tensor.tensor_type == GGMLQuantizationType.BF16:
             yield tensor.name, flat.view(torch.bfloat16).reshape(shape).clone()
         else:
diff --git a/extension/llm/export/test/test_gguf.py b/extension/llm/export/test/test_gguf.py
