up

Author: metascroy

backends/mlx/builder/op_helpers.py

@@ -18,6 +18,11 @@
 if TYPE_CHECKING:
     from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder
 
+# When True, always serialize the biases tensor for quantized ops.
+# When False, use init-time computation when zero_point is all zeros,
+# computing biases = -scales * 2^(bits-1) during the init chain.
+QUANTIZED_SERIALIZE_BIASES = False
+
 
 def get_aten_target(target):
     """
@@ -168,6 +173,50 @@ def emit_lifted_constant(P: "MLXProgramBuilder", value, dtype: torch.dtype) -> S
     return slot
 
 
+def emit_quantized_biases(
+    P: "MLXProgramBuilder",
+    zero_point_key: str,
+    scale: torch.Tensor,
+    zero_point: torch.Tensor,
+    bits: int,
+    B: torch.Tensor,
+    scale_slot: "Slot",
+) -> "Slot":
+    """Emit biases for quantized ops, computing at init time when possible.
+
+    When zero_point is all zeros and QUANTIZED_SERIALIZE_BIASES is False,
+    avoids serializing the biases tensor by computing biases = scales * -offset
+    during the init chain instead.
+
+    Returns the biases Slot.
+    """
+    from executorch.backends.mlx.serialization.mlx_graph_schema import MultiplyNode
+    from torch._subclasses.fake_tensor import FakeTensor
+
+    is_scale_only = False
+    if not isinstance(zero_point, FakeTensor):
+        if torch.sum(torch.abs(zero_point)).item() == 0:
+            is_scale_only = True
+
+    if QUANTIZED_SERIALIZE_BIASES or not is_scale_only:
+        return P.make_or_get_constant(f"{zero_point_key}_to_biases", B)
+
+    scale_dtype = scale.dtype
+    offset = 1 << (bits - 1)
+    neg_offset = emit_lifted_constant(P, -offset, scale_dtype)
+    biases = P.make_or_get_constant(
+        f"{zero_point_key}_to_biases_dummy", torch.tensor(0.0, dtype=B.dtype)
+    )
+    P.emit_init(
+        MultiplyNode(
+            a=P.slot_to_tid(scale_slot),
+            b=P.slot_to_tid(neg_offset),
+            out=P.slot_to_tid(biases),
+        )
+    )
+    return biases
+
+
 def to_mlx_qparams(
     qdata: torch.Tensor,
     scale: torch.Tensor,
@@ -217,6 +266,34 @@ def to_mlx_qparams(
         return Q, None
 
 
+def parse_dequant_nvfp4_node(
+    node: Node,
+) -> Optional[Tuple[Node, Node, Node, torch.dtype]]:
+    """Parse a torchao.dequantize_nvfp4 node.
+
+    Returns (qdata, scale, per_tensor_scale, output_dtype) or None if not a
+    dequantize_nvfp4 node or the custom op is not registered.
+    """
+    target = get_aten_target(node.target)
+    try:
+        import executorch.extension.llm.export.nvfp4  # noqa: F401
+    except ImportError:
+        return None
+
+    if target is not torch.ops.torchao.dequantize_nvfp4.default:
+        return None
+
+    qdata, scale, per_tensor_scale = node.args[0:3]
+
+    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, per_tensor_scale, output_dtype
+
+
 def parse_dequant_node(
     node: Node,
 ) -> Optional[Tuple[Node, Node, Node, int, int, Optional[torch.dtype], int]]:

backends/mlx/ops.py

@@ -21,6 +21,7 @@
 import torch
 from executorch.backends.mlx.builder.op_helpers import (
     emit_lifted_constant,
+    emit_quantized_biases,
     parse_dequant_node,
     to_mlx_qparams,
     torch_dtype_to_scalar_type,
@@ -3646,8 +3647,10 @@ def _dequantize_affine_handler(P: MLXProgramBuilder, n: Node) -> Slot:
         B = B.reshape(*leading_dims, B.shape[-1])
 
     w = P.make_or_get_constant(f"{qdata_target}_to_packed", Q)
-    biases = P.make_or_get_constant(f"{zero_point_target}_to_biases", B)
     scale_const = P.make_or_get_constant(f"{scale_target}_scale", scale_nd)
+    biases = emit_quantized_biases(
+        P, zero_point_target, scale, zero_point, bits, B, scale_const
+    )
 
     if needs_permute:
         _, dequant_tmp = P.make_tmp_slot()