diff --git a/README.md b/README.md index f9f2320b..14bae260 100644 --- a/README.md +++ b/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/ \ diff --git a/docs/cli_reference.md b/docs/cli_reference.md index 98960809..8b3a7961 100644 --- a/docs/cli_reference.md +++ b/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`) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 821490e8..9debbb43 100644 --- a/src/mobius/__main__.py +++ b/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) diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index a02f6884..cfc7117f 100644 --- a/src/mobius/_builder.py +++ b/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(): diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index c1777181..01869191 100644 --- a/src/mobius/_optimizations.py +++ b/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) diff --git a/src/mobius/functions/matmul_nbits_test.py b/src/mobius/functions/matmul_nbits_test.py index 3eac8f27..55093d87 100644 --- a/src/mobius/functions/matmul_nbits_test.py +++ b/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 diff --git a/tests/cli_test.py b/tests/cli_test.py index 7a3e7cf6..b8b2ff1d 100644 --- a/tests/cli_test.py +++ b/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 (