Trim ONNX attention heads + MLP slots per layer (not just sparsify)

The streaming ONNX exporter gated both head trimming (compile.py:1034)
and MLP-slot trimming (1079) behind `on_layer_compiled is None`, but the
ONNX path always installs a streaming callback — so every ONNX compile
skipped both trims and emitted full-width (128-head) layers, merely
zeroing the unused heads/slots. The KV cache and MatMuls stayed full.

Fix: trim each layer inside `_make_stream_layer_weights_cb` — the one
chokepoint all 3 callsites flow through — before reading `attn.n_heads`
and before the tensors are nulled. `used_heads` and the per-slot weights
are already populated by write_attn_sublayer/write_mlp_sublayer when the
callback fires, so both trims see final counts. Gated by a new
`trim_heads` param on the factory, threaded through both exporters;
`trim_heads=False` keeps today's full-width sparsified model.

Downstream is already per-layer-ready (per_layer_n_heads sizes every
emitted tensor + ValueInfo; OnnxTokenRuntime discovers heads from
past_K shapes), so no runtime change. Post-loop in-process trims and
their guards are unchanged (corrected the false comments that claimed
the streaming sparse export already dropped zeros). Added a verbose
"Head pruning (ONNX)" summary to both exporters.

Tests: rewrote test_streaming_trim docstring (the ordering regression
stays); added 4 tests to test_headless_onnx — KV cache shrinks per-layer
+ non-uniform, full width preserved under trim_heads=False, MLP W1 slots
shrink, and a trim-vs-no-trim numerical no-op (only all-zero
heads/slots removed). Full torchwright + torchwright_doom suites green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: physicsrob

tests/compile/forward/test_headless_onnx.py

@@ -53,7 +53,7 @@ def _build_sample_graph():
     return out, create_pos_encoding()
 
 
-def _export(output_node, pos_encoding, tmpdir, name="model.onnx"):
+def _export(output_node, pos_encoding, tmpdir, name="model.onnx", trim_heads=True):
     onnx_path = os.path.join(tmpdir, name)
     compile_headless_to_onnx(
         output_node,
@@ -63,10 +63,39 @@ def _export(output_node, pos_encoding, tmpdir, name="model.onnx"):
         d_head=D_HEAD,
         max_seq_len=32,
         verbose=False,
+        trim_heads=trim_heads,
     )
     return onnx_path
 
 
+def _l_w1_d_hidden(onnx_path):
+    """Per-layer MLP hidden widths read straight off the ONNX initializers.
+
+    ``l{i}_W1`` is the first MLP weight, emitted with shape ``(d, d_hidden_i)``
+    (it is the LHS of ``MatMul(res, W1)`` where ``res`` is ``(t, d)``), so its
+    second dim is that layer's hidden width after trimming.  A mostly-zero W1 is
+    emitted as a SparseTensorProto, so look in both the dense and sparse init
+    lists (the sparse proto's ``dims`` still carries the full dense shape and its
+    name lives on ``values.name``).
+    """
+    import onnx
+
+    model = onnx.load(onnx_path)
+    dims_by_name = {init.name: list(init.dims) for init in model.graph.initializer}
+    for sp in model.graph.sparse_initializer:
+        dims_by_name[sp.values.name] = list(sp.dims)
+
+    widths = []
+    i = 0
+    while f"l{i}_W1" in dims_by_name:
+        dims = dims_by_name[f"l{i}_W1"]
+        assert dims[0] == D, f"l{i}_W1 first dim {dims[0]} != d {D}"
+        widths.append(int(dims[1]))
+        i += 1
+    assert widths, "no l{i}_W1 initializers found"
+    return widths
+
+
 # ---------------------------------------------------------------------------
 # Test 1: Prefill on full sequence matches compute() reference
 # ---------------------------------------------------------------------------
@@ -401,3 +430,107 @@ def test_compiled_headless_step_prefill_decode_matches_full():
     # past_K should have grown to n_total = 5
     past_K, _ = past
     assert past_K[0].shape[1] == 5
+
+
+# ---------------------------------------------------------------------------
+# Test 7: trim_heads actually trims the exported ONNX (heads + MLP slots)
+# ---------------------------------------------------------------------------
+#
+# The streaming exporter used to leave every layer full-width (d/d_head heads,
+# full d_hidden MLP), merely sparsifying the unused heads/slots to zero.  These
+# tests pin that ``trim_heads=True`` genuinely shrinks the per-layer KV cache
+# (past_K_i widths) and MLP MatMuls (l{i}_W1 widths), that ``trim_heads=False``
+# preserves the full width, and that trimming is a numerical no-op (only all-zero
+# heads/slots are removed).
+
+
+def test_headless_onnx_trim_heads_shrinks_kv_cache():
+    """trim_heads=True: per-layer past_K widths are below the full head count."""
+    out, pos = _build_sample_graph()
+    max_heads = D // D_HEAD  # 16
+    with tempfile.TemporaryDirectory() as tmpdir:
+        onnx_path = _export(out, pos, tmpdir, trim_heads=True)
+        session = onnxruntime.InferenceSession(onnx_path)
+        _, per_layer_n_heads, _ = _discover_meta(session)
+
+    assert per_layer_n_heads, "no layers discovered"
+    assert all(1 <= nh <= max_heads for nh in per_layer_n_heads)
+    assert min(per_layer_n_heads) < max_heads, (
+        f"no layer trimmed below full width {max_heads}: {per_layer_n_heads}"
+    )
+    # A real per-layer trim leaves the layers non-uniform (different layers use
+    # different head counts); a coincidental uniform value would still satisfy
+    # the bound above, so assert the per-layer variation explicitly.
+    assert len(set(per_layer_n_heads)) > 1, (
+        f"head counts are uniform across layers: {per_layer_n_heads}"
+    )
+
+
+def test_headless_onnx_no_trim_preserves_full_width():
+    """trim_heads=False: every layer keeps the full d/d_head head count."""
+    out, pos = _build_sample_graph()
+    max_heads = D // D_HEAD  # 16
+    with tempfile.TemporaryDirectory() as tmpdir:
+        onnx_path = _export(out, pos, tmpdir, trim_heads=False)
+        session = onnxruntime.InferenceSession(onnx_path)
+        _, per_layer_n_heads, _ = _discover_meta(session)
+
+    assert per_layer_n_heads, "no layers discovered"
+    assert all(nh == max_heads for nh in per_layer_n_heads), (
+        f"expected all layers at full width {max_heads}: {per_layer_n_heads}"
+    )
+
+
+def test_headless_onnx_trim_shrinks_mlp_slots():
+    """trim_heads=True shrinks some l{i}_W1 below full d_hidden; False keeps all."""
+    out, pos = _build_sample_graph()
+    full_d_hidden = D  # d_hidden defaults to d when omitted
+    with tempfile.TemporaryDirectory() as tmpdir:
+        trim_path = _export(out, pos, tmpdir, name="trim.onnx", trim_heads=True)
+        notrim_path = _export(out, pos, tmpdir, name="notrim.onnx", trim_heads=False)
+        trim_widths = _l_w1_d_hidden(trim_path)
+        notrim_widths = _l_w1_d_hidden(notrim_path)
+
+    assert all(1 <= w <= full_d_hidden for w in trim_widths)
+    assert min(trim_widths) < full_d_hidden, (
+        f"no layer's MLP trimmed below full d_hidden {full_d_hidden}: {trim_widths}"
+    )
+    assert all(w == full_d_hidden for w in notrim_widths), (
+        f"expected all layers at full d_hidden {full_d_hidden}: {notrim_widths}"
+    )
+
+
+def test_headless_onnx_trim_is_numerical_noop():
+    """Trimmed and full-width ONNX produce identical output on the same prefill.
+
+    Trimming removes only all-zero heads/slots, so prefilling the same rows
+    through both sessions must agree — any divergence means a non-zero head or
+    slot was trimmed (a real bug).
+    """
+    out, pos = _build_sample_graph()
+    a_vals = torch.tensor([[3.0], [5.0], [-2.0], [0.0], [4.0]])
+    b_vals = torch.tensor([[4.0], [-1.0], [3.0], [7.0], [2.0]])
+    inputs_np = torch.cat([a_vals, b_vals], dim=1).numpy().astype(np.float32)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        trim_path = _export(out, pos, tmpdir, name="trim.onnx", trim_heads=True)
+        notrim_path = _export(out, pos, tmpdir, name="notrim.onnx", trim_heads=False)
+
+        sess_trim = onnxruntime.InferenceSession(trim_path)
+        sess_notrim = onnxruntime.InferenceSession(notrim_path)
+
+        _, heads_trim, d_head = _discover_meta(sess_trim)
+        _, heads_notrim, _ = _discover_meta(sess_notrim)
+
+        feeds_trim = {"inputs": inputs_np}
+        feeds_trim.update(_empty_past_feeds(heads_trim, d_head))
+        out_trim = sess_trim.run(["outputs"], feeds_trim)[0]
+
+        feeds_notrim = {"inputs": inputs_np}
+        feeds_notrim.update(_empty_past_feeds(heads_notrim, d_head))
+        out_notrim = sess_notrim.run(["outputs"], feeds_notrim)[0]
+
+    assert np.allclose(out_trim, out_notrim, atol=1e-4), (
+        f"trim changed the output: max diff "
+        f"{np.abs(out_trim - out_notrim).max():.6e}"
+    )

tests/compile/forward/test_streaming_trim.py

@@ -10,9 +10,17 @@
 
 The onnxruntime-based ``compile_to_onnx`` tests are skipped when onnxruntime is
 absent, so this no-onnxruntime test pins the trim/null ordering directly: with a
-streaming callback present, ``forward_compile`` must complete (trims skipped).
-The streamed sparse export already drops the zero heads/slots, so the in-memory
-trim is both impossible and redundant in that mode.
+streaming callback present, ``forward_compile`` must complete (post-loop trims
+skipped because the callback already nulled the weights).
+
+NOTE: the *real* per-layer head/slot trim for the ONNX export now happens
+*inside* ``_make_stream_layer_weights_cb`` (it calls ``trim_unused_heads`` /
+``trim_unused_slots`` on each layer before reading and nulling its tensors), so
+the exported ONNX genuinely shrinks the KV cache and MatMuls.  That behaviour is
+covered by ``test_headless_onnx.py`` (per-layer ``past_K_i`` widths, ``l{i}_W1``
+shapes, and a trim-vs-no-trim numerical no-op).  The stub callback below does
+NOT trim — it only nulls — so this test remains a pure regression for the
+ordering guard: the post-loop trim must not slice already-nulled weights.
 """
 
 import torch

torchwright/compiler/export.py

@@ -258,14 +258,27 @@ def _make_stream_layer_weights_cb(
     dense_inits: list,
     sparse_inits: list,
     per_layer_n_heads: list,
+    trim_heads: bool = True,
 ) -> Callable[[int, object], None]:
     """Factory for the forward_compile on_layer_compiled callback.
 
     Emits each freshly-compiled layer's dense weights into the
     dense/sparse init lists (sparsified on the fly when mostly zero) and
     nulls out the layer's tensor attributes so GC can reclaim them.
 
-    After head trimming, each layer's attention matrices have shape
+    When ``trim_heads`` is set (the default), each layer is trimmed *in
+    place here* — before its tensors are read — so the unused (all-zero)
+    attention heads and MLP hidden slots are physically removed: the
+    emitted WQ/WK/WV/WO matrices, the per-layer reshape constants, and
+    the ``past_K_i``/``past_V_i`` ValueInfo all shrink to the used count,
+    cutting both the KV cache and the inference MatMuls.  This is the
+    *only* place the ONNX export trims; the post-loop trims in
+    ``forward_compile`` are skipped whenever a streaming callback is
+    installed (their tensors are already nulled here).  With
+    ``trim_heads=False`` the full-width (sparsified-but-not-trimmed)
+    model is emitted instead.
+
+    After trimming, each layer's attention matrices have shape
     ``(nh, d, d_head)`` where ``nh <= d // d_head``.  Weight matrices
     are emitted at the trimmed size and per-layer reshape constants
     are added so the ONNX graph can reshape to the correct head count.
@@ -277,6 +290,14 @@ def on_layer_compiled(i: int, layer) -> None:
         attn = layer.attn.attn
         mlp = layer.mlp
 
+        # Trim each layer here, before its tensors are read and nulled
+        # below.  ``used_heads`` (attn) and the per-slot weight columns
+        # (mlp) were populated by write_attn_sublayer/write_mlp_sublayer
+        # during this layer's compile, so both trims see final counts.
+        if trim_heads:
+            attn.trim_unused_heads()  # n_heads -> used heads (KV cache + attn MatMuls)
+            mlp.trim_unused_slots()  # d_hidden -> used slots (MLP MatMuls)
+
         nh = attn.n_heads  # post-trim head count
         hd = nh * d_head
         per_layer_n_heads.append(nh)
@@ -621,6 +642,7 @@ def compile_headless_to_onnx(
         dense_inits,
         sparse_inits,
         per_layer_n_heads,
+        trim_heads=trim_heads,
     )
 
     # --- Phase 1: streaming compile ---------------------------------------
@@ -639,6 +661,15 @@ def compile_headless_to_onnx(
         assume_zero_init=assume_zero_init,
     )
     t_compile = time.perf_counter() - t0
+    if verbose and per_layer_n_heads:
+        _max_heads = d // d_head
+        _total = _max_heads * len(per_layer_n_heads)
+        _kept = sum(per_layer_n_heads)
+        print(
+            f"  Head pruning (ONNX): {_total - _kept}/{_total} heads pruned; "
+            f"per-layer heads range [{min(per_layer_n_heads)}, "
+            f"{max(per_layer_n_heads)}] of {_max_heads}"
+        )
 
     # --- Phase 2: metadata + graph assembly -------------------------------
     assert compiled.residual_assignment is not None
@@ -833,6 +864,7 @@ def compile_to_onnx(
         dense_inits,
         sparse_inits,
         per_layer_n_heads,
+        trim_heads=trim_heads,
     )
 
     # --- Phase 1: streaming compile ---------------------------------------
@@ -852,6 +884,15 @@ def compile_to_onnx(
         d_hidden=d_hidden,
     )
     t_compile = time.perf_counter() - t0
+    if verbose and per_layer_n_heads:
+        _max_heads = d // d_head
+        _total = _max_heads * len(per_layer_n_heads)
+        _kept = sum(per_layer_n_heads)
+        print(
+            f"  Head pruning (ONNX): {_total - _kept}/{_total} heads pruned; "
+            f"per-layer heads range [{min(per_layer_n_heads)}, "
+            f"{max(per_layer_n_heads)}] of {_max_heads}"
+        )
 
     # --- Phase 2: metadata + graph assembly -------------------------------
     assert compiled.residual_assignment is not None

torchwright/compiler/forward/compile.py

@@ -1026,11 +1026,12 @@ def forward_compile(
     net.residual_assignment = ra
     net.assert_aliases = graph.get_assert_aliases()
 
-    # Skip in-place weight trimming when streaming: ``on_layer_compiled`` (the
-    # ONNX streaming path) extracts each layer's weights and NULLs the tensors
-    # per layer, so the post-loop trims here would slice ``None``. The streamed
-    # sparse export already drops the unused (zero) heads/slots, making the
-    # in-memory trim both impossible and redundant in that mode.
+    # This post-loop trim is the in-process (no-callback) path's trim.  The ONNX
+    # streaming path trims each layer *inside* ``on_layer_compiled`` (see
+    # ``_make_stream_layer_weights_cb``) before reading and NULLing its tensors,
+    # so by the time we get here those tensors are already trimmed and nulled —
+    # re-trimming would slice ``None``.  Hence the ``on_layer_compiled is None``
+    # guard: only the in-process return value still carries live weights to trim.
     if trim_heads and on_layer_compiled is None:
         max_heads = d // d_head
         for layer in net.layers:
@@ -1073,9 +1074,10 @@ def forward_compile(
                     f"{kv_depth:>10}  {cp:>10}  {sa:>10}"
                 )
 
-    # Trim trailing unused MLP slots (same idea as head trimming). Skipped when
-    # streaming for the same reason as head trimming above (weights already
-    # extracted + nulled per layer by ``on_layer_compiled``).
+    # Trim trailing unused MLP slots (same idea as head trimming, and the
+    # in-process counterpart to the per-layer ``mlp.trim_unused_slots()`` the
+    # streaming callback runs).  Skipped when streaming because the callback
+    # already trimmed and nulled each layer's MLP weights above.
     if on_layer_compiled is None:
         for layer in net.layers:
             layer.mlp.trim_unused_slots()