support MXFP8 (#21146)

Differential Revision: D113113318

Pull Request resolved: #21146

Author: metascroy

backends/mlx/builder/op_helpers.py

@@ -656,6 +656,35 @@ def parse_dequant_nvfp4_node(
     return qdata, scale, per_tensor_scale, output_dtype
 
 
+def parse_dequant_mx_node(
+    node: Node,
+) -> Optional[Tuple[Node, Node, torch.dtype, int, torch.dtype]]:
+    """Parse a torchao.dequantize_mx node.
+
+    Returns (qdata, scale, elem_dtype, block_size, output_dtype) or None if not a
+    dequantize_mx node or the custom op is not registered. MX has no per-tensor
+    scale and no affine zero-point; the block scale is always E8M0.
+    """
+    target = get_aten_target(node.target)
+    try:
+        import executorch.extension.llm.export.mx  # noqa: F401
+    except ImportError:
+        return None
+
+    if target is not torch.ops.torchao.dequantize_mx.default:
+        return None
+
+    qdata, scale, elem_dtype, block_size = node.args[0:4]
+
+    output_dtype = torch.float32
+    if len(node.args) > 4:
+        output_dtype = node.args[4]
+    elif "output_dtype" in node.kwargs:
+        output_dtype = node.kwargs["output_dtype"]
+
+    return qdata, scale, elem_dtype, int(block_size), output_dtype
+
+
 def parse_dequant_int4_node(
     node: Node,
 ) -> Optional[Tuple[Node, Node, Node, int, Optional[torch.dtype]]]:

backends/mlx/builder/program_builder.py

@@ -567,6 +567,23 @@ def _process_nodes(self) -> None:  # noqa C901
             if n.op in ("placeholder", "output"):
                 dtype_error = _check_dtype(n)
                 if dtype_error is not None:
+                    # A placeholder consumed entirely by pattern handlers (which
+                    # repack it into new constants -- e.g. an MX fp8/e8m0 weight
+                    # reinterpreted to uint32/uint8) is never emitted as an MLX
+                    # tensor, so its own (unserializable) dtype is irrelevant. Its
+                    # users are the pattern's head/body nodes, which at this point
+                    # have a handler assigned (by _apply_patterns) but are not yet
+                    # marked handled, so check for a claimed handler too.
+                    if (
+                        n.op == "placeholder"
+                        and n.users
+                        and all(
+                            self._is_handled(u) or self.node_info[u].handler is not None
+                            for u in n.users
+                        )
+                    ):
+                        self._mark_supported(n)
+                        continue
                     self._mark_unsupported(n, f"{n.op} {dtype_error}")
                     continue
                 self._mark_supported(n)

backends/mlx/patterns.py

@@ -23,6 +23,7 @@
     emit_quantized_gather,
     emit_stop_position,
     parse_dequant_int4_node,
+    parse_dequant_mx_node,
     parse_dequant_node,
     parse_dequant_nvfp4_node,
     regroup_affine_scales,
@@ -1166,6 +1167,196 @@ def __call__(self, P, n):
         return out
 
 
+def _mx_mlx_mode_bits(elem_dtype: torch.dtype):
+    """Map an MX element dtype to the MLX (mode, bits). Rejects unsupported formats.
+
+    MLX's block-scaled kernels only implement E4M3 elements (mxfp8). Other MX
+    element encodings (e5m2, fp4/fp6, ...) are rejected rather than silently
+    mis-lowered.
+    """
+    if elem_dtype == torch.float8_e4m3fn:
+        return "mxfp8", 8
+    raise ValueError(
+        f"MLX backend does not support MX element dtype {elem_dtype}; "
+        "only float8_e4m3fn (mxfp8) is supported."
+    )
+
+
+def _mx_pack_for_mlx(P, qdata_node, scale_node):
+    """Reinterpret ExportableMXTensor's native FP8/E8M0 storage to MLX layout.
+
+    ExportableMXTensor stores ``qdata`` as native FP8 (one value per byte) and
+    ``scale`` as ``float8_e8m0fnu``; MLX's kernel wants ``qdata`` as uint32
+    (4 FP8 codes per word) and ``scale`` as uint8. Both conversions are pure bit
+    reinterpretations (``view``), registered as MLX constants.
+    """
+    qdata_target, qdata = P.get_placeholder_target_and_tensor(qdata_node)
+    scale_target, scale = P.get_placeholder_target_and_tensor(scale_node)
+    w = qdata.contiguous().view(torch.uint8).view(torch.uint32)
+    sc = scale.contiguous().view(torch.uint8)
+    w_slot = P.make_or_get_constant(f"{qdata_target}_mx_packed", w)
+    scale_slot = P.make_or_get_constant(f"{scale_target}_mx_u8", sc)
+    return w_slot, scale_slot
+
+
+@REGISTRY.register_pattern(name="MX_QUANTIZED_LINEAR")
+class MXQuantizedLinearHandler(PatternHandler):
+    """Fuse dequantize_mx + linear into QuantizedMatmulNode for the MX format.
+
+    Matches:
+        linear(x, dequantize_mx(qdata, scale, elem_dtype, block_size), bias)
+
+    Emits:
+        QuantizedMatmulNode [→ AddNode(bias)] [→ AsTypeNode]
+
+    The element dtype is mapped to the MLX ``mode``/``bits`` (rejecting formats
+    MLX can't lower). MX has no per-tensor scale, so no MultiplyNode is emitted.
+    """
+
+    def __init__(self, head, body, qdata, scale, elem_dtype, block_size, output_dtype):
+        super().__init__(head, body)
+        self.qdata = qdata
+        self.scale = scale
+        self.elem_dtype = elem_dtype
+        self.block_size = block_size
+        self.output_dtype = output_dtype
+
+    @classmethod
+    def maybe_create(cls, ep, head):
+        if not match_target(head, torch.ops.aten.linear.default):
+            return None
+        x, dequant = head.args[0:2]
+        if not isinstance(dequant, Node):
+            return None
+        if not has_single_user(dequant):
+            return None
+        parsed = parse_dequant_mx_node(dequant)
+        if parsed is None:
+            return None
+        qdata, scale, elem_dtype, block_size, output_dtype = parsed
+        return cls(head, [dequant], qdata, scale, elem_dtype, block_size, output_dtype)
+
+    def __call__(self, P, n):
+        assert n == self.head
+
+        x_node, w_node = n.args[0:2]
+        b_node = n.args[2] if len(n.args) > 2 else None
+
+        mode, bits = _mx_mlx_mode_bits(self.elem_dtype)
+        needs_cast = x_node.meta["val"].dtype != self.output_dtype
+        has_bias = b_node is not None
+
+        w, scales = _mx_pack_for_mlx(P, self.qdata, self.scale)
+        x, bias = P.slot_map([x_node, b_node])
+
+        out = P.make_or_get_slot(n)
+        P.emit(
+            QuantizedMatmulNode(
+                x=P.slot_to_tid(x),
+                w=P.slot_to_tid(w),
+                scales=P.slot_to_tid(scales),
+                out=P.slot_to_tid(out),
+                biases=None,
+                group_size=self.block_size,
+                bits=bits,
+                mode=mode,
+                transpose=True,
+            )
+        )
+
+        if has_bias:
+            P.emit(
+                AddNode(
+                    a=P.slot_to_tid(out),
+                    b=P.slot_to_tid(bias),
+                    out=P.slot_to_tid(out),
+                )
+            )
+
+        if needs_cast:
+            P.emit(
+                AsTypeNode(
+                    x=P.slot_to_tid(out),
+                    out=P.slot_to_tid(out),
+                    scalar_type=torch_dtype_to_scalar_type(self.output_dtype),
+                )
+            )
+
+        return out
+
+
+@REGISTRY.register_pattern(name="MX_QUANTIZED_EMBEDDING")
+class MXQuantizedEmbeddingHandler(PatternHandler):
+    """Fuse dequantize_mx + embedding into gather + DequantizeNode for the MX format.
+
+    Matches:
+        embedding(dequantize_mx(qdata, scale, elem_dtype, block_size), indices)
+
+    Emits:
+        TakeNode(qdata) → TakeNode(scales) → DequantizeNode [→ AsTypeNode]
+    """
+
+    def __init__(self, head, body, qdata, scale, elem_dtype, block_size, output_dtype):
+        super().__init__(head, body)
+        self.qdata = qdata
+        self.scale = scale
+        self.elem_dtype = elem_dtype
+        self.block_size = block_size
+        self.output_dtype = output_dtype
+
+    @classmethod
+    def maybe_create(cls, ep, head):
+        if not match_target(head, torch.ops.aten.embedding.default):
+            return None
+
+        w, x = head.args[0:2]
+        if not isinstance(w, Node):
+            return None
+        if not has_single_user(w):
+            return None
+        parsed = parse_dequant_mx_node(w)
+        if parsed is None:
+            return None
+        qdata, scale, elem_dtype, block_size, output_dtype = parsed
+        return cls(head, [w], qdata, scale, elem_dtype, block_size, output_dtype)
+
+    def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
+        assert n == self.head
+        w_node, x_node = n.args[0:2]
+
+        mode, bits = _mx_mlx_mode_bits(self.elem_dtype)
+        x_dtype = x_node.meta["val"].dtype
+        needs_cast = self.output_dtype != x_dtype
+
+        qdata_slot, scales_slot = _mx_pack_for_mlx(P, self.qdata, self.scale)
+        (x,) = P.slot_map([x_node])
+
+        out = P.make_or_get_slot(n)
+        emit_quantized_gather(
+            P,
+            out,
+            x,
+            qdata_slot,
+            scales_slot,
+            None,
+            group_size=self.block_size,
+            bits=bits,
+            mode=mode,
+            out_dtype=self.output_dtype,
+        )
+
+        if needs_cast:
+            P.emit(
+                AsTypeNode(
+                    x=P.slot_to_tid(out),
+                    out=P.slot_to_tid(out),
+                    scalar_type=torch_dtype_to_scalar_type(self.output_dtype),
+                )
+            )
+
+        return out
+
+
 @REGISTRY.register_pattern(name="INT4_QUANTIZED_LINEAR")
 class Int4QuantizedLinearHandler(PatternHandler):
     """Fuse dequantize_int4_tensor + linear into QuantizedMatmulNode(mode="affine").

backends/mlx/test/test_ops.py

@@ -7556,6 +7556,84 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (x,)
 
 
+class MXFP8QuantizedLinearModel(nn.Module):
+    """Simple linear layer that will be quantized with MXFP8."""
+
+    def __init__(
+        self, in_features: int = 64, out_features: int = 128, bias: bool = True
+    ):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+@register_test
+class MXFP8QuantizedLinearTest(OpTestCase):
+    """Test case for MXFP8 quantized nn.Linear (mode="mxfp8", group_size 32)."""
+
+    name = "mxfp8_quantized_linear"
+    rtol = 0.1
+    atol = 0.1
+
+    def __init__(
+        self,
+        in_features: int = 64,
+        out_features: int = 128,
+        batch_size: int = 2,
+        seq_len: int = 16,
+        bias: bool = True,
+        dtype: torch.dtype = torch.float32,
+    ):
+        self.in_features = in_features
+        self.out_features = out_features
+        self.batch_size = batch_size
+        self.seq_len = seq_len
+        self.bias = bias
+        self.dtype = dtype
+
+        parts = ["mxfp8_quantized_linear"]
+        if not bias:
+            parts.append("no_bias")
+        if dtype != torch.float32:
+            parts.append(str(dtype).split(".")[-1])
+        self.name = "_".join(parts)
+
+    @classmethod
+    def get_test_configs(cls) -> List["MXFP8QuantizedLinearTest"]:
+        return [
+            cls(),
+            cls(bias=False),
+            cls(dtype=torch.bfloat16),
+            cls(bias=False, dtype=torch.bfloat16),
+        ]
+
+    def get_edge_compile_config(self):
+        from executorch.exir import EdgeCompileConfig
+
+        return EdgeCompileConfig(_check_ir_validity=False)
+
+    def create_model(self) -> nn.Module:
+        model = MXFP8QuantizedLinearModel(
+            self.in_features, self.out_features, bias=self.bias
+        )
+        model = model.to(self.dtype)
+
+        from executorch.extension.llm.export.mx import ExportableMXConfig
+        from torchao.quantization import quantize_
+
+        quantize_(model, ExportableMXConfig())
+
+        return model
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        x = torch.randn(
+            self.batch_size, self.seq_len, self.in_features, dtype=self.dtype
+        )
+        return (x,)
+
+
 def _make_int4_quantized_weight(weight: torch.Tensor, group_size: int) -> torch.Tensor:
     """Groupwise affine 4-bit quantize a ``(N, K)`` weight into an
     ``ExportableInt4Tensor`` (torchao ``Int4Tensor`` packed layout)."""