Add QDQ build feature

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

Author: Copilot

README.md

@@ -113,8 +113,8 @@ mobius build --model openai/whisper-tiny output_dir/
 ```
 
 Build-mode toggles use the cargo-style `--features` option. Available features
-are `static-cache`, `fp8-kv-cache`, `prune-lm-head`, and `text-only`. Pass them
-as a comma-separated list or repeat the option:
+are `static-cache`, `fp8-kv-cache`, `prune-lm-head`, `qdq`, and `text-only`.
+Pass them as a comma-separated list or repeat the option:
 
 ```sh
 mobius build --model meta-llama/Llama-3.2-1B output_dir/ \

docs/cli_reference.md

@@ -176,6 +176,7 @@ option. Pass a comma-separated list (and/or repeat the flag):
 ```
 --features fp8-kv-cache,static-cache
 --features prune-lm-head
+--features qdq
 --features text-only
 ```
 
@@ -186,6 +187,7 @@ Available features:
 | `static-cache` | Pre-allocate fixed-size KV cache buffers using `TensorScatter` (pair with `--max-seq-len N`). Requires `DecoderLayer` / `MoEDecoderLayer` models. Cannot combine with `--task`. |
 | `fp8-kv-cache` | Store the `GroupQueryAttention` KV cache as `FLOAT8E4M3FN` (per-tensor E4M3), halving KV-cache memory. Requires a GQA build (e.g. `--ep cuda --dtype f16`) and an ORT runtime with the FP8 KV-cache kernel (SM89+). Pair with `--kv-cache-scale-file` for calibrated scales. |
 | `prune-lm-head` | Select the final hidden-state position before the LM-head projection and emit logits shaped `[B, 1, vocab]`. Supported by models using the base `CausalLMModel.forward()` path; unsupported custom forwards fail explicitly. Use only when the downstream workflow does not need per-token logits. |
+| `qdq` | Lower quantized `com.microsoft::MatMulNBits` weights to standard ONNX QDQ form (`DequantizeLinear` + `MatMul`) even when the selected EP supports the native contrib op. |
 | `text-only` | Export the text backbone of a multimodal checkpoint as a standalone decoder-only LLM (see below). |
 
 The legacy boolean flags `--static-cache`, `--fp8-kv-cache`, and
@@ -200,6 +202,9 @@ mobius build --model Qwen/Qwen2.5-0.5B output/ \
 
 mobius build --model meta-llama/Llama-3.2-1B output/ \
     --features prune-lm-head
+
+mobius build --model meta-llama/Llama-3.2-1B output/ \
+    --features qdq
 ```
 
 ### Static Cache (`--features static-cache`)

src/mobius/__main__.py

@@ -41,6 +41,7 @@
     "fp8-kv-cache": "fp8_kv_cache",
     "prune-lm-head": "prune_lm_head",
     "text-only": "text_only",
+    "qdq": "qdq",
 }
 
 
@@ -224,6 +225,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
     # the --config and --model build paths can pass the same scales.
     fp8_kv_cache = getattr(args, "fp8_kv_cache", False)
     prune_lm_head = getattr(args, "prune_lm_head", False)
+    qdq = getattr(args, "qdq", False)
     kv_cache_scales: dict[int, tuple[float, float]] | None = None
     scale_file = getattr(args, "kv_cache_scale_file", None)
     if scale_file is not None and not fp8_kv_cache:
@@ -314,6 +316,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             fp8_kv_cache=fp8_kv_cache,
             kv_cache_scales=kv_cache_scales,
             prune_lm_head=prune_lm_head,
+            qdq=qdq,
         )
         for name, model in pkg.items():
             model.graph.name = f"{config_path}/{name}"
@@ -344,6 +347,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             fp8_kv_cache=fp8_kv_cache,
             kv_cache_scales=kv_cache_scales,
             prune_lm_head=prune_lm_head,
+            qdq=qdq,
         )
 
     _save_package(pkg, output_dir, args, optimize, component_filter)

src/mobius/_builder.py

@@ -160,6 +160,7 @@ def build_from_module(
     fp8_kv_cache: bool = False,
     kv_cache_scales: dict[int, tuple[float, float]] | None = None,
     prune_lm_head: bool = False,
+    qdq: bool = False,
 ) -> ModelPackage:
     """Build an ONNX :class:`ModelPackage` from a module instance and config.
 
@@ -204,6 +205,10 @@ def build_from_module(
             via the causal-LM task so runtimes can avoid full prefill LM-head
             projection. Only supported by ``text-generation`` and
             ``hybrid-text-generation`` tasks.
+        qdq: When ``True``, lower quantized ``com.microsoft::MatMulNBits``
+            weights to standard ONNX QDQ form (``DequantizeLinear`` +
+            ``MatMul``) even if the selected execution provider has a native
+            ``MatMulNBits`` kernel.
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).
@@ -256,6 +261,7 @@ def forward(self, op, input_ids, attention_mask,
             trace=trace_optimization,
             fp8_kv_cache=fp8_kv_cache,
             kv_cache_scales=kv_cache_scales,
+            qdq=qdq,
         )
 
     _maybe_apply_opset_lowering(pkg, execution_provider)
@@ -401,6 +407,7 @@ def build(
     fp8_kv_cache: bool = False,
     kv_cache_scales: dict[int, tuple[float, float]] | None = None,
     prune_lm_head: bool = False,
+    qdq: bool = False,
 ) -> ModelPackage:
     """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID.
 
@@ -481,6 +488,10 @@ def build(
             (``[B, 1, vocab]``). This is intended for single-token
             autoregressive generation and is incompatible with workflows that
             need per-token logits.
+        qdq: When ``True``, lower quantized ``com.microsoft::MatMulNBits``
+            weights to standard ONNX QDQ form (``DequantizeLinear`` +
+            ``MatMul``) even if the selected execution provider has a native
+            ``MatMulNBits`` kernel.
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).
@@ -660,6 +671,7 @@ def build(
         fp8_kv_cache=fp8_kv_cache,
         kv_cache_scales=kv_cache_scales,
         prune_lm_head=prune_lm_head,
+        qdq=qdq,
     )
 
     for name, model in pkg.items():

src/mobius/_optimizations.py

@@ -371,6 +371,7 @@ def optimize_model(
     trace: bool = False,
     fp8_kv_cache: bool = False,
     kv_cache_scales: dict[int, tuple[float, float]] | None = None,
+    qdq: bool = False,
 ) -> None:
     """Apply EP-aware optimization passes to *model* in-place.
 
@@ -406,6 +407,9 @@ def optimize_model(
             per-tensor FP8 scales (from offline calibration). Only used when
             ``fp8_kv_cache`` is ``True``; layers absent from the map use a unit
             scale of ``1.0``.
+        qdq: When ``True``, force ``com.microsoft::MatMulNBits`` lowering to
+            standard ONNX QDQ form (``DequantizeLinear`` + ``MatMul``),
+            regardless of whether *ep* has a native ``MatMulNBits`` kernel.
 
     Raises:
         ValueError: If *ep* is not a registered execution provider.
@@ -460,11 +464,13 @@ def _should_inline(func: ir.Function) -> bool:
         if func.domain == "com.microsoft" and func.name == "PackedMultiHeadAttention":
             return not caps.supports_packed_multi_head_attention
         # MatMulNBits (blockwise-INT4) → QDQ (DequantizeLinear + MatMul) for EPs
-        # without a MatMulNBits kernel (QNN HTP). Supported EPs (CPU/CUDA/…) keep
-        # the compact contrib op and its native kernel; the function body stays
-        # registered but uninlined (kernels take precedence over local functions).
+        # without a MatMulNBits kernel (QNN HTP), or when the user explicitly
+        # requests standard QDQ operators via the qdq build feature. Supported
+        # EPs (CPU/CUDA/…) keep the compact contrib op and its native kernel by
+        # default; the function body stays registered but uninlined (kernels
+        # take precedence over local functions).
         if func.domain == "com.microsoft" and func.name == "MatMulNBits":
-            return not caps.supports_matmul_nbits
+            return qdq or not caps.supports_matmul_nbits
         return False
 
     inline_pass = common_passes.InlinePass(criteria=_should_inline)

src/mobius/functions/matmul_nbits_test.py

@@ -133,7 +133,7 @@ def test_inline_matches_native_op_odd_blocks(self):
 class TestMatMulNBitsEpGating:
     """A qnn build lowers MatMulNBits to QDQ; a cpu build keeps the contrib op."""
 
-    def _build(self, ep: str):
+    def _build(self, ep: str, qdq: bool = False):
         import dataclasses
         from collections import Counter
 
@@ -166,7 +166,11 @@ def _build(self, ep: str):
         )
         module = registry.get("qwen2")(cfg)
         model = build_from_module(
-            module, cfg, task=_default_task_for_model("qwen2"), execution_provider=ep
+            module,
+            cfg,
+            task=_default_task_for_model("qwen2"),
+            execution_provider=ep,
+            qdq=qdq,
         )["model"]
         return Counter(n.op_type for n in model.graph)
 
@@ -179,3 +183,8 @@ def test_qnn_lowers_to_qdq(self):
         ops = self._build("qnn")
         assert ops.get("MatMulNBits", 0) == 0
         assert ops.get("DequantizeLinear", 0) > 0
+
+    def test_qdq_feature_lowers_to_qdq_on_cpu(self):
+        ops = self._build("cpu", qdq=True)
+        assert ops.get("MatMulNBits", 0) == 0
+        assert ops.get("DequantizeLinear", 0) > 0

tests/cli_test.py

@@ -268,6 +268,30 @@ def test_features_prune_lm_head_passed_through(self):
             )
         assert mock_build.call_args.kwargs.get("prune_lm_head") is True
 
+    def test_features_qdq_passed_through(self):
+        """--features qdq sets qdq on the build() call."""
+        with (
+            tempfile.TemporaryDirectory() as tmpdir,
+            mock.patch(
+                "mobius._diffusers_builder._load_diffusers_pipeline_index",
+                return_value=None,
+            ),
+            mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as mock_build,
+            mock.patch("mobius.__main__._save_package"),
+        ):
+            main(
+                [
+                    "build",
+                    "--model",
+                    "some/model",
+                    tmpdir,
+                    "--no-weights",
+                    "--features",
+                    "qdq",
+                ]
+            )
+        assert mock_build.call_args.kwargs.get("qdq") is True
+
     def test_features_comma_separated_multiple(self):
         """A single --features accepts a comma-separated list."""
         with (