Merge branch 'main' into perf/matmulnbits-accuracy-level

Author: justinchuby

src/mobius/__main__.py

@@ -351,7 +351,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"
@@ -594,7 +597,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

@@ -212,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
@@ -245,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):

src/mobius/components/_quantized_linear_test.py

@@ -354,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