Quantize GGUF output head (lm_head) → MatMulNBits (#401)

Follow-up to #400 (quantized embedding). Profiling showed the CPU model
shipped an untied `lm_head` as a plain fp32 MatMul (~544 MB) run every
token. This emits the output head as Q4 `MatMulNBits` (untied) / shares
the packed embedding table (tied).

- Result: 169 MatMulNBits, 1 GatherBlockQuantized, **0 plain fp32
MatMul**; no 544 MB fp32 tables. Model ~1.2 GB → ~399 MB.
- Coherent output verified ('Paris').
- Note: a quick 6-thread decode check showed ~38 vs ~40 tok/s (no clear
speedup) — a rigorous benchmark is pending; the win is model size +
correct all-quantized graph shape (matches llama.cpp). Stacks on #400's
embedding commit.

lintrunner clean; gguf pytest 157 passed.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

src/mobius/__main__.py

@@ -598,8 +598,8 @@ def main(argv: list[str] | None = None) -> None:
         "--keep-quantized",
         action="store_true",
         help=(
-            "Preserve supported projection and embedding quantization via "
-            "MatMulNBits/GatherBlockQuantized."
+            "Preserve supported projection, output-head, and embedding "
+            "quantization via MatMulNBits/GatherBlockQuantized."
         ),
     )
     gguf_parser.add_argument(

src/mobius/_configs/_quantization.py

@@ -31,7 +31,7 @@ class QuantizationConfig:
     # Olive RTN exports and quantized GGUF imports.
     quantize_embeddings: bool = False
     # When True, the LM head projection is block-wise quantized (MatMulNBits).
-    # Used by Olive RTN exports and tied quantized GGUF imports.
+    # 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

@@ -186,6 +186,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 +220,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

@@ -299,6 +299,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

src/mobius/integrations/gguf/_builder.py

@@ -8,10 +8,11 @@
 
 - **Dequantized** (default): All quantized tensors are dequantized to
   float.  Simple, but loses the compression benefit of quantization.
-- **Quantized** (``keep_quantized=True``): Linear-layer weights are
-  normalized to a common MatMulNBits layout, including mixed presets
-  such as Q4_K_M. Compatible quantized token embeddings are repacked
-  for GatherBlockQuantized. Other tensors are dequantized.
+- **Quantized** (``keep_quantized=True``): Linear-layer weights, including
+  a quantized output head, are repacked into MatMulNBits format and token
+  embeddings into GatherBlockQuantized format. Mixed presets such as
+  Q4_K_M are normalized to one quantization layout. Other tensors are
+  dequantized.
 """
 
 from __future__ import annotations
@@ -175,6 +176,11 @@ def build_from_gguf(
             bits=bits,
             block_size=block_size,
         )
+        quantize_lm_head = (
+            quantize_embeddings
+            if config.tie_word_embeddings
+            else _can_quantize_lm_head(gguf_model, gguf_arch)
+        )
         config = dataclasses.replace(
             config,
             quantization=QuantizationConfig(
@@ -184,17 +190,19 @@ def build_from_gguf(
                 sym=is_sym,
                 float_zero_point=float_zp,
                 quantize_embeddings=quantize_embeddings,
-                quantize_lm_head=quantize_embeddings and config.tie_word_embeddings,
-                tie_word_embeddings=quantize_embeddings and config.tie_word_embeddings,
+                quantize_lm_head=quantize_lm_head,
+                tie_word_embeddings=quantize_lm_head and config.tie_word_embeddings,
             ),
         )
         logger.info(
-            "Quantized mode: bits=%d, block_size=%d, symmetric=%s, float_zp=%s, embedding=%s",
+            "Quantized mode: bits=%d, block_size=%d, symmetric=%s, "
+            "float_zp=%s, embedding=%s, lm_head=%s",
             bits,
             block_size,
             is_sym,
             float_zp,
             quantize_embeddings,
+            quantize_lm_head,
         )
 
     # 4. Look up module class and resolve task
@@ -462,6 +470,45 @@ def _can_quantize_embedding(
     return False
 
 
+def _can_quantize_lm_head(gguf_model, gguf_arch: str) -> bool:
+    """Return whether an untied GGUF output head can be kept quantized."""
+    from mobius.integrations.gguf._tensor_mapping import map_gguf_to_hf_names
+
+    supported_types = {
+        "Q1_0",
+        "Q2_K",
+        "Q3_K",
+        "Q4_0",
+        "Q4_1",
+        "Q4_K",
+        "Q5_0",
+        "Q5_1",
+        "Q5_K",
+        "Q6_K",
+        "Q8_0",
+    }
+    for name, _raw, qtype, shape in gguf_model.tensor_items_raw():
+        if map_gguf_to_hf_names(name, gguf_arch) != "lm_head.weight":
+            continue
+        return len(shape) == 2 and getattr(qtype, "name", None) in supported_types
+    return False
+
+
+def _require_supported_requantization(
+    *,
+    bits: int,
+    block_size: int,
+    tensor_name: str,
+) -> None:
+    if bits != 4 or block_size != 32:
+        raise ValueError(
+            "keep_quantized MatMulNBits requantization currently supports only "
+            f"4-bit/block-32 targets; got bits={bits} block={block_size} "
+            f"for tensor {tensor_name}. Use keep_quantized=False or a "
+            "4-bit/block-32 target."
+        )
+
+
 def _load_dequantized_state_dict(
     gguf_model,
     gguf_arch: str,
@@ -500,10 +547,10 @@ def _load_quantized_state_dict(
 ) -> dict:
     """Load tensors, normalizing quantized projections to MatMulNBits.
 
-    Projection weights (Q/K/V/O and MLP) are converted to the graph's
-    common MatMulNBits format. Mixed source types are dequantized and
-    requantized when they do not already match that target. Compatible
-    quantized embeddings are converted for GatherBlockQuantized. Norms
+    Projection weights (Q/K/V/O, MLP, and a quantized output head) are
+    converted to the graph's common MatMulNBits format, and token embeddings
+    to GatherBlockQuantized format. Mixed or unsupported source types are
+    dequantized and requantized when they do not match that target. Norms
     and other non-linear tensors remain dequantized.
 
     For llama-family models, quantized Q/K weights receive the
@@ -603,6 +650,11 @@ def _load_quantized_state_dict(
                     shape_2d,
                 )
                 if repacked.bits != target_bits or repacked.block_size != target_block_size:
+                    _require_supported_requantization(
+                        bits=target_bits,
+                        block_size=target_block_size,
+                        tensor_name=hf_name,
+                    )
                     values = gguf_model.dequantize_raw_tensor(raw, qtype, np_shape)
                     repacked = repack_dequantized_tensor(
                         values,
@@ -612,6 +664,11 @@ def _load_quantized_state_dict(
                     )
                     n_requantized += 1
             else:
+                _require_supported_requantization(
+                    bits=target_bits,
+                    block_size=target_block_size,
+                    tensor_name=hf_name,
+                )
                 values = gguf_model.dequantize_raw_tensor(raw, qtype, np_shape)
                 repacked = repack_dequantized_tensor(
                     values,

src/mobius/integrations/gguf/_builder_test.py

@@ -21,6 +21,8 @@ def _write_quantized_gguf(
     intermediate_size: int = 128,
     vocab_size: int = 256,
     quantize_embedding: bool = False,
+    projection_quantization: str = "q4_0",
+    output_quantization: str | None = None,
     tie_embeddings: bool = False,
 ) -> None:
     """Write a GGUF file with Q4_0 quantized projection weights.
@@ -66,44 +68,72 @@ def _add_q4_0(name: str, n_out: int, k_in: int) -> None:
                 )
         writer.add_tensor(name, raw, raw_dtype=GGMLQuantizationType.Q4_0)
 
+    def _add_q8_0(name: str, n_out: int, k_in: int) -> None:
+        """Write a Q8_0-quantized weight tensor."""
+        block_size = 32
+        block_bytes = 34  # 2B scale + 32B quants
+        n_blocks = k_in // block_size
+        bytes_per_row = n_blocks * block_bytes
+        raw = np.zeros((n_out, bytes_per_row), dtype=np.uint8)
+        for row in range(n_out):
+            for b in range(n_blocks):
+                off = b * block_bytes
+                scale = np.random.uniform(0.01, 1.0)
+                raw[row, off : off + 2] = np.array([scale], dtype=np.float16).view(np.uint8)
+                raw[row, off + 2 : off + 34] = (
+                    np.random.randint(-127, 128, size=32, dtype=np.int16)
+                    .astype(np.int8, copy=False)
+                    .view(np.uint8)
+                )
+        writer.add_tensor(name, raw, raw_dtype=GGMLQuantizationType.Q8_0)
+
     if quantize_embedding:
         _add_q4_0("token_embd.weight", vocab_size, hidden_size)
     else:
         _add_f32("token_embd.weight", (vocab_size, hidden_size))
 
+    add_projection = {
+        "q4_0": _add_q4_0,
+        "q8_0": _add_q8_0,
+    }[projection_quantization]
+
     for i in range(num_layers):
-        # Projection weights (Q4_0)
-        _add_q4_0(
+        add_projection(
             f"blk.{i}.attn_q.weight",
             num_heads * head_dim,
             hidden_size,
         )
-        _add_q4_0(
+        add_projection(
             f"blk.{i}.attn_k.weight",
             num_kv_heads * head_dim,
             hidden_size,
         )
-        _add_q4_0(
+        add_projection(
             f"blk.{i}.attn_v.weight",
             num_kv_heads * head_dim,
             hidden_size,
         )
-        _add_q4_0(
+        add_projection(
             f"blk.{i}.attn_output.weight",
             hidden_size,
             num_heads * head_dim,
         )
-        _add_q4_0(f"blk.{i}.ffn_gate.weight", intermediate_size, hidden_size)
-        _add_q4_0(f"blk.{i}.ffn_up.weight", intermediate_size, hidden_size)
-        _add_q4_0(f"blk.{i}.ffn_down.weight", hidden_size, intermediate_size)
+        add_projection(f"blk.{i}.ffn_gate.weight", intermediate_size, hidden_size)
+        add_projection(f"blk.{i}.ffn_up.weight", intermediate_size, hidden_size)
+        add_projection(f"blk.{i}.ffn_down.weight", hidden_size, intermediate_size)
         # Norms (float32)
         _add_f32(f"blk.{i}.attn_norm.weight", (hidden_size,))
         _add_f32(f"blk.{i}.ffn_norm.weight", (hidden_size,))
 
-    # Output norm + optional untied lm_head (float32)
+    # Output norm + optional untied lm_head
     _add_f32("output_norm.weight", (hidden_size,))
     if not tie_embeddings:
-        _add_f32("output.weight", (vocab_size, hidden_size))
+        if output_quantization == "q4_0":
+            _add_q4_0("output.weight", vocab_size, hidden_size)
+        elif output_quantization == "q8_0":
+            _add_q8_0("output.weight", vocab_size, hidden_size)
+        else:
+            _add_f32("output.weight", (vocab_size, hidden_size))
 
     writer.write_header_to_file()
     writer.write_kv_data_to_file()
@@ -135,6 +165,30 @@ def q4_0_tied_embedding_gguf(tmp_path: Path) -> Path:
     return path
 
 
+@pytest.fixture
+def q4_0_embedding_q8_head_gguf(tmp_path: Path) -> Path:
+    """Create a GGUF with a Q4 embedding and an untied Q8 output head."""
+    path = tmp_path / "test_q4_0_embedding_q8_head.gguf"
+    _write_quantized_gguf(
+        path,
+        quantize_embedding=True,
+        output_quantization="q8_0",
+    )
+    return path
+
+
+@pytest.fixture
+def q8_0_projection_q4_head_gguf(tmp_path: Path) -> Path:
+    """Create a GGUF whose Q4 output head would require Q8 requantization."""
+    path = tmp_path / "test_q8_0_projection_q4_head.gguf"
+    _write_quantized_gguf(
+        path,
+        projection_quantization="q8_0",
+        output_quantization="q4_0",
+    )
+    return path
+
+
 class TestBuildQuantizedGguf:
     """Tests for build_from_gguf(keep_quantized=True)."""
 
@@ -191,6 +245,45 @@ def test_tied_quantized_embedding_drives_matmulnbits_head(
         assert "model.embed_tokens.qweight" in model.graph.initializers
         assert not any(name.startswith("lm_head.") for name in model.graph.initializers)
 
+    def test_untied_quantized_head_uses_q4_matmulnbits(
+        self, q4_0_embedding_q8_head_gguf: Path
+    ):
+        """An untied quantized output is requantized to the graph's Q4 layout."""
+        import onnx_ir as ir
+
+        from mobius.integrations.gguf import build_from_gguf
+
+        model = build_from_gguf(q4_0_embedding_q8_head_gguf, keep_quantized=True)["model"]
+        head_nodes = [
+            node
+            for node in model.graph
+            if node.op_type == "MatMulNBits" and node.outputs[0].name == "logits"
+        ]
+        assert len(head_nodes) == 1
+        assert head_nodes[0].attributes["bits"].value == 4
+        assert head_nodes[0].attributes["block_size"].value == 32
+
+        qweight = model.graph.initializers["lm_head.weight"]
+        assert qweight.dtype == ir.DataType.UINT8
+        assert list(qweight.shape) == [256, 2, 16]
+        assert list(model.graph.initializers["lm_head.scales"].shape) == [256, 2]
+        assert "lm_head.weight_t" not in model.graph.initializers
+
+    def test_unsupported_requantization_target_has_clear_error(
+        self, q8_0_projection_q4_head_gguf: Path
+    ):
+        """Mixed targets outside 4-bit/block-32 fail before the Q4 repacker."""
+        from mobius.integrations.gguf import build_from_gguf
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                "keep_quantized MatMulNBits requantization currently supports only "
+                r"4-bit/block-32 targets; got bits=8 block=32 for tensor lm_head\.weight"
+            ),
+        ):
+            build_from_gguf(q8_0_projection_q4_head_gguf, keep_quantized=True)
+
     def test_norms_are_float(self, q4_0_gguf: Path):
         """Norm weights remain float, not quantized."""
         import onnx_ir as ir