Enable weight-free MXFP6 DQ export

QEfficient/base/modeling_qeff.py

@@ -38,6 +38,7 @@
 )
 from QEfficient.compile.qnn_compiler import compile as qnn_compile
 from QEfficient.exporter.weight_free.export import embed_weight_spec_as_metadata, link_prepared_checkpoint_dir
+from QEfficient.exporter.weight_free.mxfp6 import MXFP6_ONNX_OPSET
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.transformers.models.pytorch_transforms import (
     BlockingAttentionTransform,
@@ -250,6 +251,7 @@ def __init__(self, model: torch.nn.Module, **kwargs) -> None:
         # Flag for checking if weights are offloaded
         self._is_weights_offloaded: bool = False
         self._weight_free: bool = kwargs.get("weight_free", False)
+        self._mxfp6_config = kwargs.get("mxfp6_config", None)
         # Flag for checking if model has been transformed yet
         self.is_transformed: bool = False
 
@@ -609,14 +611,24 @@ def _resolve_pkv_names(layer_idx, layer_state):
             active_transforms = [
                 transform for transform in self._onnx_transforms if transform not in excluded_transforms
             ]
+            if (
+                getattr(getattr(self, "_mxfp6_config", None), "enabled", False)
+                and CustomOpTransform not in active_transforms
+                and CustomOpTransform not in excluded_transforms
+            ):
+                active_transforms.append(CustomOpTransform)
             needs_external_tensor_data = any(
                 transform in active_transforms for transform in (FP16ClipTransform, SplitTensorsTransform)
             )
             transform_kwargs = {
                 "onnx_base_dir": str(export_dir) if needs_external_tensor_data else None,
                 "model_name": self.model_name,
                 "dynamic_axes": None if dynamo else dynamic_axes,
-                "onnx_export_opset": constants.get_onnx_export_opset(dynamo),
+                "onnx_export_opset": (
+                    MXFP6_ONNX_OPSET
+                    if getattr(getattr(self, "_mxfp6_config", None), "enabled", False)
+                    else constants.get_onnx_export_opset(dynamo)
+                ),
             }
             if onnx_transform_kwargs is not None:
                 transform_kwargs.update(onnx_transform_kwargs)
@@ -1073,6 +1085,14 @@ def _compile(
         """
 
         layerwise_cache_probe = compiler_options.pop("_layerwise_cache_probe", False)
+        if getattr(getattr(self, "_mxfp6_config", None), "enabled", False):
+            compiler_mxfp6_keys = ("mxfp6_matmul", "mxfp6-matmul", "mxfp6")
+            if any(compiler_options.get(key, False) for key in compiler_mxfp6_keys):
+                raise ValueError(
+                    "`mxfp6_matmul=True`/`mxfp6=True` cannot be used when QEff-owned `mxfp6=True` is active."
+                )
+            if enable_qnn:
+                raise ValueError("QNN compilation is not supported when QEff-owned `mxfp6=True` is active.")
 
         for removed_option in ("compile_only", "compile-only"):
             if removed_option in compiler_options:

QEfficient/base/onnx_transforms.py

@@ -49,7 +49,7 @@
     CtxScatterFuncCB3D,
 )
 from QEfficient.customop.onnxscript_utils import get_onnxscript_func
-from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func
+from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func, UnpackMxfp6, UnpackMxfp6Func
 from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc
 from QEfficient.utils import constants
 from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, SIZE_THRESHOLD_DEFAULT
@@ -116,6 +116,7 @@ class CustomOpTransform(BaseOnnxTransform):
         "CtxScatterFuncCB": (CtxScatterFuncCB, CtxScatterCB),
         "CtxGatherFuncCB": (CtxGatherFuncCB, CtxGatherCB),
         "CastToUInt4": (CastToUInt4Func, CastToUInt4),
+        "UnpackMxfp6": (UnpackMxfp6Func, UnpackMxfp6),
         "CtxChunkScatterBatchFunc": (CtxChunkScatterBatchFunc, CtxChunkScatterBatch),
         "CtxGatherFuncBlockedKVBatch": (CtxGatherFuncBlockedKVBatch, CtxGatherBlockedKVBatch),
     }
@@ -136,8 +137,16 @@ def apply(cls, model: ModelProto, onnx_export_opset: int = constants.ONNX_LEGACY
         # Add function prototypes to model
         existing = {f.name for f in model.functions}
 
-        for func_name, onnxscript_func in cls._custom_ops.values():
-            proto = get_onnxscript_func(onnxscript_func, onnx_export_opset).to_function_proto()
+        for op_name, (_, onnxscript_func) in cls._custom_ops.items():
+            if onnxscript_func is None:
+                if op_name in used_op_types:
+                    cls._ensure_opset_imports(model, "com.qti.aisw.onnx", 1)
+                    op_applied = True
+                continue
+            try:
+                proto = get_onnxscript_func(onnxscript_func, onnx_export_opset).to_function_proto()
+            except AttributeError:
+                proto = onnxscript_func.to_function_proto()
             if proto.name not in used_op_types:
                 continue
             if proto.name not in existing:

QEfficient/customop/onnxscript_utils.py

@@ -68,6 +68,6 @@ def get_dynamo_onnxscript_func(onnxscript_func):
 
 def get_onnxscript_func(onnxscript_func, onnx_export_opset: int):
     """Return the ONNXScript variant matching the requested export opset."""
-    if onnx_export_opset == constants.ONNX_DYNAMO_EXPORT_OPSET:
+    if onnx_export_opset >= constants.ONNX_DYNAMO_EXPORT_OPSET:
         return get_dynamo_onnxscript_func(onnxscript_func)
     return onnxscript_func

QEfficient/customop/quantization_ops.py

@@ -67,6 +67,57 @@ def CastToUInt4(weight_packed: onnxscript.UINT8) -> onnxscript.UINT8:
     return ops.Cast(reshaped, to=int(TensorProto.UINT4))
 
 
+@onnxscript.script(onnxscript.values.Opset("com.qti.aisw.onnx", 1))
+def UnpackMxfp6(weight_packed: onnxscript.UINT8) -> onnxscript.UINT8:
+    """
+    Unpack packed MXFP6 bytes into logical FLOAT6E2M3 values.
+
+    Supports N-D input: all leading dimensions are preserved; the last
+    dimension is expanded by 4/3. Each 3-byte group stores four 6-bit codes:
+
+      code0 = byte0[5:0]
+      code1 = byte0[7:6] | byte1[3:0] << 2
+      code2 = byte1[7:4] | byte2[1:0] << 4
+      code3 = byte2[7:2]
+    """
+    three = ops.Constant(value_ints=[3])
+    four = ops.Constant(value_ints=[4])
+    sixteen = ops.CastLike(ops.Constant(value_ints=[16]), weight_packed)
+    sixty_four = ops.CastLike(ops.Constant(value_ints=[64]), weight_packed)
+
+    packed_shape = ops.Shape(weight_packed)
+    leading_dims = ops.Slice(packed_shape, starts=[0], ends=[-1], axes=[0])
+    last_dim = ops.Slice(packed_shape, starts=[-1], ends=[2147483647], axes=[0])
+    group_count = ops.Div(last_dim, three)
+    grouped_shape = ops.Concat(leading_dims, group_count, three, axis=0)
+    grouped = ops.Reshape(weight_packed, grouped_shape)
+
+    byte0 = ops.Slice(grouped, starts=[0], ends=[1], axes=[-1])
+    byte1 = ops.Slice(grouped, starts=[1], ends=[2], axes=[-1])
+    byte2 = ops.Slice(grouped, starts=[2], ends=[3], axes=[-1])
+
+    shift2 = ops.CastLike(ops.Constant(value_ints=[2]), weight_packed)
+    shift4 = ops.CastLike(ops.Constant(value_ints=[4]), weight_packed)
+    shift6 = ops.CastLike(ops.Constant(value_ints=[6]), weight_packed)
+
+    code0 = ops.Mod(byte0, sixty_four)
+    code1_low = ops.BitShift(byte0, shift6, direction="RIGHT")
+    code1_high = ops.BitShift(ops.Mod(byte1, sixteen), shift2, direction="LEFT")
+    code1 = ops.Mod(ops.Add(code1_low, code1_high), sixty_four)
+    code2_low = ops.BitShift(byte1, shift4, direction="RIGHT")
+    code2_masked_high = ops.Mod(byte2, ops.CastLike(ops.Constant(value_ints=[4]), weight_packed))
+    code2_high = ops.BitShift(code2_masked_high, shift4, direction="LEFT")
+    code2 = ops.Mod(ops.Add(code2_low, code2_high), sixty_four)
+    code3 = ops.Mod(ops.BitShift(byte2, shift2, direction="RIGHT"), sixty_four)
+
+    stacked = ops.Concat(code0, code1, code2, code3, axis=-1)
+    last_dim_unpacked = ops.Div(ops.Mul(last_dim, four), three)
+    new_shape = ops.Concat(leading_dims, last_dim_unpacked, axis=0)
+    reshaped = ops.Reshape(stacked, new_shape)
+
+    return ops.Cast(reshaped, to=int(TensorProto.FLOAT6E2M3))
+
+
 class CastToUInt4Func(torch.autograd.Function):
     """
     Custom op: unpacks packed uint8 → uint8 (values 0-15) in PyTorch.
@@ -100,6 +151,39 @@ def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value:
         return output
 
 
+class UnpackMxfp6Func(torch.autograd.Function):
+    """
+    Custom op: unpacks packed MXFP6 bytes into logical 6-bit code positions.
+
+    PyTorch forward returns uint8 code values for unit tests:
+      (..., in_features * 3 // 4) UINT8 -> (..., in_features) UINT8
+
+    ONNX symbolic emits an ONNXScript custom op. The subgraph ends with
+    Cast -> FLOAT6E2M3.
+    """
+
+    @staticmethod
+    def forward(weight_packed: torch.Tensor) -> torch.Tensor:
+        if weight_packed.shape[-1] % 3 != 0:
+            raise ValueError("Packed MXFP6 final dimension must be divisible by 3")
+        grouped = weight_packed.to(torch.uint8).reshape(*weight_packed.shape[:-1], weight_packed.shape[-1] // 3, 3)
+        code0 = grouped[..., 0] & 0x3F
+        code1 = ((grouped[..., 0] >> 6) | ((grouped[..., 1] & 0x0F) << 2)) & 0x3F
+        code2 = ((grouped[..., 1] >> 4) | ((grouped[..., 2] & 0x03) << 4)) & 0x3F
+        code3 = (grouped[..., 2] >> 2) & 0x3F
+        return torch.stack((code0, code1, code2, code3), dim=-1).reshape(
+            *weight_packed.shape[:-1], weight_packed.shape[-1] * 4 // 3
+        )
+
+    @staticmethod
+    def setup_context(ctx, inputs, outputs):
+        pass
+
+    @staticmethod
+    def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value:
+        return g.onnxscript_op(UnpackMxfp6, weight_packed)
+
+
 class DequantizeLinearFunc(torch.autograd.Function):
     """
     Emits a standard ONNX DequantizeLinear node (ai.onnx domain, not custom).

QEfficient/exporter/weight_free/checkpoint_key_resolver.py

@@ -180,6 +180,7 @@ def promote_initializers_and_build_spec(onnx_program, model_ref: str, model_name
     ]
     backbone = qeff_model.model.base_model if isinstance(qeff_model.model, PooledModel) else qeff_model.model
     promoted_inputs: List[WeightSpecInput] = []
+    spec_input_names = set()
 
     for name, init_value in list(model_ir.graph.initializers.items()):
         if name not in model_names:
@@ -211,6 +212,31 @@ def promote_initializers_and_build_spec(onnx_program, model_ref: str, model_name
                 location=WeightSpecLocation(file=checkpoint_files.index(checkpoint_file), key=checkpoint_key),
             )
         )
+        spec_input_names.add(name)
+
+    for value in model_ir.graph.inputs:
+        name = value.name
+        if name in spec_input_names or name not in model_names:
+            continue
+
+        onnx_name = tied_weight_map.get(name, name)
+        checkpoint_key = find_checkpoint_key(onnx_name, checkpoint_index, backbone)
+        if checkpoint_key is None:
+            if _is_computed_initializer(onnx_name):
+                continue
+            raise ValueError(
+                f"Could not resolve model graph input '{name}' to a safetensors checkpoint key "
+                f"(resolved name: '{onnx_name}', model: '{model_ref}')."
+            )
+
+        checkpoint_file = checkpoint_index[checkpoint_key]
+        promoted_inputs.append(
+            WeightSpecInput(
+                name=name,
+                location=WeightSpecLocation(file=checkpoint_files.index(checkpoint_file), key=checkpoint_key),
+            )
+        )
+        spec_input_names.add(name)
 
     return WeightSpec(
         model_name=model_name,

QEfficient/exporter/weight_free/export.py

@@ -15,6 +15,7 @@
 from accelerate import init_empty_weights
 
 from QEfficient.exporter.weight_free.checkpoint_key_resolver import promote_initializers_and_build_spec
+from QEfficient.exporter.weight_free.mxfp6 import finalize_mxfp6_export
 from QEfficient.exporter.weight_free.weight_spec import load_weight_spec, resolve_weight_spec_path, save_weight_spec
 from QEfficient.utils import load_json
 from QEfficient.utils.checkpoint_utils import resolve_checkpoint_dir
@@ -47,6 +48,8 @@ def _run_quantizer_for_wf(qeff_model, target_dtype: torch.dtype):
     quant_config = getattr(qeff_model.model.config, "quantization_config", None)
 
     if quant_config is not None:
+        if getattr(getattr(qeff_model, "_mxfp6_config", None), "enabled", False):
+            raise ValueError("`mxfp6=True` does not currently support source HF quantization configs.")
         # For quantized models the meta model must use the same quantized layer types as the
         # checkpoint so that ONNX initializer names match the checkpoint's storage keys.
         # qeff_model.model was built via from_config (no real quantized checkpoint load), so
@@ -133,6 +136,9 @@ def _prepare_checkpoint_for_weight_free_export(
 
     source_dir = resolve_checkpoint_dir(model_ref)
     dtype_suffix = str(target_dtype).replace("torch.", "")
+    mxfp6_config = getattr(qeff_model, "_mxfp6_config", None)
+    if getattr(mxfp6_config, "enabled", False):
+        dtype_suffix += f"-mxfp6-{mxfp6_config.scale_dtype}"
     # TODO(wf): For different flavours of the model that expect different checkpoint weight layouts,
     # we end up overriding old one. We need to add support of hashing/caching here.
     prepared_name = source_dir.name + f"-qeff-prepared-{dtype_suffix}"
@@ -236,7 +242,13 @@ def export_weight_free_onnx(
     )
     _prune_unused_fake_initializers(onnx_program)
     onnx_program.save(str(onnx_path))
-    save_weight_spec(resolve_weight_spec_path(onnx_path), spec)
+    weight_spec_path = save_weight_spec(resolve_weight_spec_path(onnx_path), spec)
+    finalize_mxfp6_export(
+        onnx_path=onnx_path,
+        weight_spec_path=weight_spec_path,
+        prepared_model_ref=prepared_model_ref,
+        config=getattr(qeff_model, "_mxfp6_config", None),
+    )
 
     return meta_qeff_model, onnx_transform_kwargs