Merge branch 'main' into feat/onnx-genai-metadata-export

Author: justinchuby

src/mobius/__main__.py

@@ -363,7 +363,10 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None:
         raise SystemExit(1)
 
     if args.keep_quantized:
-        print("Quantized mode: preserving GGUF quantization as MatMulNBits...")
+        print(
+            "Quantized mode: preserving GGUF quantization as "
+            "MatMulNBits/GatherBlockQuantized..."
+        )
 
     gguf_path = args.gguf_path
     output_dir = args.output or os.path.splitext(gguf_path)[0] + "_onnx"
@@ -617,7 +620,10 @@ def main(argv: list[str] | None = None) -> None:
     gguf_parser.add_argument(
         "--keep-quantized",
         action="store_true",
-        help="Preserve quantization via MatMulNBits (Q4_0/Q4_1/Q8_0).",
+        help=(
+            "Preserve supported projection, output-head, and embedding "
+            "quantization via MatMulNBits/GatherBlockQuantized."
+        ),
     )
     gguf_parser.add_argument(
         "--dtype",

src/mobius/_configs/_quantization.py

@@ -27,11 +27,11 @@ class QuantizationConfig:
     # (e.g. Tencent SEQ uses 1.5).
     float_zero_point: bool = False
     # When True, the input embedding table is block-wise quantized and is
-    # looked up with GatherBlockQuantized instead of a plain Gather. Set by
-    # Olive RTN exports that pass ``embeds: true``.
+    # looked up with GatherBlockQuantized instead of a plain Gather. Used by
+    # Olive RTN exports and quantized GGUF imports.
     quantize_embeddings: bool = False
     # When True, the LM head projection is block-wise quantized (MatMulNBits).
-    # Set by Olive RTN exports that pass ``lm_head: true``.
+    # Used by Olive RTN exports and quantized GGUF imports.
     quantize_lm_head: bool = False
     # When True, the input embedding and LM head share one weight table. Olive
     # RTN records this in its own config (``tie_word_embeddings``) and may clear

src/mobius/components/_quantized_linear.py

@@ -15,6 +15,8 @@
 import onnx_ir as ir
 from onnxscript import OpBuilder, nn
 
+from mobius._build_context import ep_capabilities
+
 # MatMulNBits packs weights into uint8 blobs.  The packed shape depends
 # on bits and block_size:
 #   packed_weights: [N, n_blocks, blob_size]  (uint8)
@@ -26,6 +28,29 @@
 _MICROSOFT_DOMAIN = "com.microsoft"
 
 
+def _accuracy_level_attrs(bits: int) -> dict[str, int]:
+    """Return the ``accuracy_level`` attribute for ``MatMulNBits``, if any.
+
+    Only emitted for 4-bit weights: ``accuracy_level`` is sourced from
+    ``EpCapabilities.default_int4_accuracy_level`` and its int8-accumulation
+    semantics are defined for INT4 ``MatMulNBits``. For 2-bit / 8-bit weights the
+    attribute is omitted so those paths keep ORT's default behavior.
+
+    ORT's MLAS CPU ``MatMulNBits`` kernel selects its compute path from the
+    ``accuracy_level`` attribute: unset/0 keeps the highest-precision fp32
+    dequant + fp32 GEMM path, while ``4`` dynamically quantizes activations to
+    int8 and uses int8 dot-products (SDOT/NEON on ARM, AVX-VNNI on x86) — the
+    same class of kernel llama.cpp uses, and typically 2-4x faster on CPU with
+    no observable quality loss for Q4 weights. The value is sourced from the
+    active EP's :attr:`EpCapabilities.default_int4_accuracy_level` (4 for CPU /
+    WebGPU). When it is 0 the attribute is omitted so ORT keeps its default.
+    """
+    if bits != 4:
+        return {}
+    level = ep_capabilities().default_int4_accuracy_level
+    return {"accuracy_level": level} if level else {}
+
+
 class QuantizedLinear(nn.Module):
     """Linear layer backed by the MatMulNBits custom op.
 
@@ -126,6 +151,7 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
             N=self._n,
             bits=self._bits,
             block_size=self._block_size,
+            **_accuracy_level_attrs(self._bits),
             _domain=_MICROSOFT_DOMAIN,
         )
         if self.bias is not None:
@@ -186,6 +212,7 @@ def __init__(
 
         self._bits = bits
         self._block_size = block_size
+        self._embedding_dim = embedding_dim
         self.padding_idx = padding_idx
 
         n_blocks = embedding_dim // block_size
@@ -219,14 +246,18 @@ def forward(self, op: OpBuilder, input_ids: ir.Value) -> ir.Value:
         if self.zero_points is not None:
             inputs.append(self.zero_points)
 
-        return op.GatherBlockQuantized(
+        result = op.GatherBlockQuantized(
             *inputs,
             bits=self._bits,
             block_size=self._block_size,
             gather_axis=0,
             quantize_axis=1,
             _domain=_MICROSOFT_DOMAIN,
         )
+        result.dtype = self.scales.dtype
+        if input_ids.shape is not None:
+            result.shape = ir.Shape([*input_ids.shape, self._embedding_dim])
+        return result
 
 
 class TiedQuantizedLMHead(nn.Module):
@@ -304,6 +335,7 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
             N=self._n,
             bits=self._bits,
             block_size=self._block_size,
+            **_accuracy_level_attrs(self._bits),
             _domain=_MICROSOFT_DOMAIN,
         )
 

src/mobius/components/_quantized_linear_test.py

@@ -9,6 +9,8 @@
 
 import pytest
 
+from mobius._build_context import build_context
+from mobius._execution_providers import ep_registry
 from mobius._testing import (
     count_op_type,
     create_test_builder,
@@ -125,7 +127,60 @@ def test_matmulnbits_attributes(self):
                 assert attrs["N"] == OUT_FEATURES
                 assert attrs["bits"] == 4
                 assert attrs["block_size"] == 32
+                # No active build context -> default EP (accuracy_level 0),
+                # so the attribute is omitted and ORT keeps its default path.
+                assert "accuracy_level" not in attrs
                 break
+        else:
+            pytest.fail("MatMulNBits node not found")
+
+    def test_no_accuracy_level_without_context(self):
+        """Default EP omits accuracy_level (ORT default / highest precision)."""
+        ql = QuantizedLinear(IN_FEATURES, OUT_FEATURES, bits=4, block_size=32)
+        b, op, graph = create_test_builder()
+        x = create_test_input(b, "x", [1, 4, IN_FEATURES])
+        result = ql(op, x)
+        b._adapt_outputs([result], "")
+        for node in graph:
+            if node.op_type == "MatMulNBits":
+                attrs = {a.name: a.value for a in node.attributes.values()}
+                assert "accuracy_level" not in attrs
+                break
+        else:
+            pytest.fail("MatMulNBits node not found")
+
+    def test_cpu_ep_emits_accuracy_level_4(self):
+        """CPU EP context stamps accuracy_level=4 (int8 MLAS path)."""
+        with build_context(ep_registry.require("cpu")):
+            ql = QuantizedLinear(IN_FEATURES, OUT_FEATURES, bits=4, block_size=32)
+            b, op, graph = create_test_builder()
+            x = create_test_input(b, "x", [1, 4, IN_FEATURES])
+            result = ql(op, x)
+            b._adapt_outputs([result], "")
+        for node in graph:
+            if node.op_type == "MatMulNBits":
+                attrs = {a.name: a.value for a in node.attributes.values()}
+                assert attrs["accuracy_level"] == 4
+                break
+        else:
+            pytest.fail("MatMulNBits node not found")
+
+    def test_cpu_ep_omits_accuracy_level_for_non_int4(self):
+        """accuracy_level is INT4-specific: 8-bit weights keep ORT's default."""
+        with build_context(ep_registry.require("cpu")):
+            ql = QuantizedLinear(IN_FEATURES, OUT_FEATURES, bits=8, block_size=32)
+            b, op, graph = create_test_builder()
+            x = create_test_input(b, "x", [1, 4, IN_FEATURES])
+            result = ql(op, x)
+            b._adapt_outputs([result], "")
+        for node in graph:
+            if node.op_type == "MatMulNBits":
+                attrs = {a.name: a.value for a in node.attributes.values()}
+                assert attrs["bits"] == 8
+                assert "accuracy_level" not in attrs
+                break
+        else:
+            pytest.fail("MatMulNBits node not found")
 
     def test_3_inputs_without_zero_points(self):
         ql = QuantizedLinear(IN_FEATURES, OUT_FEATURES)
@@ -299,6 +354,8 @@ def test_graph_has_gather_block_quantized_node(self):
         result = qe(op, ids)
         b._adapt_outputs([result], "")
         assert count_op_type(graph, "GatherBlockQuantized") == 1
+        assert result.dtype == ir.DataType.FLOAT
+        assert result.shape == ir.Shape([1, 4, self.DIM])
 
     def test_node_domain_and_attributes(self):
         import onnx_ir as ir