From 747579b9b066e7bc709ef09e84b95172a4c81c1f Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Mon, 1 Jun 2026 22:24:10 +0000
Subject: [PATCH 01/38] docs(skills): add mobius ONNX export gotchas

CLI syntax (--model + positional out dir, f16 not fp16), GQA fusion vs
--static-cache incompatibility, and the fp16 GQA packed-QKV FLOAT32 bug
that makes exports fail to load in onnxruntime (with detect + fix).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../mobius-onnx-export-gotchas/SKILL.md       | 84 +++++++++++++++++++
 1 file changed, 84 insertions(+)
 create mode 100644 .agents/skills/mobius-onnx-export-gotchas/SKILL.md

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
new file mode 100644
index 00000000..65f3b640
--- /dev/null
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -0,0 +1,84 @@
+---
+name: mobius-onnx-export-gotchas
+description: Use when building/exporting ONNX models with the `mobius build` CLI (especially Phi-3 / Phi-3.5 or any model with `--execution-provider cuda` GQA fusion and/or `--static-cache`). Covers the current CLI syntax, the dtype flag values, the GQA-vs-static-cache interaction, and a known fp16 packed-QKV FLOAT32 bug that makes GQA fp16 exports fail to load in onnxruntime.
+---
+
+# mobius ONNX export gotchas
+
+## 1. CLI syntax (editable repo differs from older docs)
+`mobius build` requires `--model <hf_id>` and takes the **output dir as a POSITIONAL** arg.
+There is **no `-o` flag for `build`** (`-o` exists only on `build-gguf`).
+
+```bash
+mobius build --model microsoft/Phi-3.5-mini-instruct \
+  --dtype f16 --execution-provider cuda \
+  --external-data onnx --trust-remote-code \
+  /path/to/output_dir
+```
+
+- `--dtype` choices: `f16`/`float16`, `bf16`/`bfloat16`, `f32`/`float32`. **`fp16` is INVALID.**
+- `--execution-provider` is an alias of `--ep`. `cuda` + fp16/bf16 triggers GQA fusion;
+  `default` keeps plain ONNX `Attention`.
+
+## 2. `--static-cache` is incompatible with GQA fusion
+`--static-cache` wraps each attention with `TensorScatter` (in-place KV cache for the **ONNX Attention**
+op). That breaks the pattern the GQA rewrite matches, so combining
+`--execution-provider cuda --static-cache` yields **0 GroupQueryAttention + N Attention + 2N TensorScatter**
+(mobius prints: "GQA fusion expected … but found 0 GroupQueryAttention and N Attention nodes").
+
+- **GQA model:** `--execution-provider cuda` **alone**. GQA's shared KV buffer
+  (`past_present_share_buffer`) is enabled at **runtime** via IO-binding past & present to the same
+  OrtValue — NOT via `--static-cache`.
+- **ONNX-Attention + in-place cache:** `--execution-provider default --static-cache --max-seq-len N`.
+
+## 3. BUG: GQA fp16 export leaves packed-QKV weights as FLOAT32 → model won't load
+For an fp16 GQA export, the per-layer packed QKV weight
+(`..q_proj.weight__k_proj.weight__v_proj.weight__axis_0__concat`) is emitted as **FLOAT32**, while its
+MatMul's other input is fp16. onnxruntime then rejects the model at load:
+
+```
+Type Error: Type parameter (T) of Optype (MatMul) bound to different types
+(tensor(float16) and tensor(float)) in node (node_MatMul_*)
+```
+
+You'll also see at save time: `The value type for shape [H, 3H] is not known. Skipping serialization`.
+
+**Root cause:** `_cast_module_dtype` (`src/mobius/_builder.py:84`) casts module params to fp16 *before*
+graph build. The GQA `PackQKVWithBias` rewrite (`src/mobius/rewrite_rules/_group_query_attention.py`)
+then emits the packed weight as a graph-level `op.Concat(q_w,k_w,v_w)` that a constant-fold collapses
+into a NEW initializer whose dtype is FLOAT32/untyped — the fp16 cast never reaches it.
+
+### Detect
+```python
+import onnx
+m = onnx.load("model.onnx", load_external_data=False)
+fp32 = [i.name for i in m.graph.initializer if i.data_type == onnx.TensorProto.FLOAT]
+print(len(fp32), "FLOAT32 initializers (should be 0 for fp16)")
+```
+
+### Fix (post-export, numerically == intended fp16)
+Cast the FLOAT32 initializers to fp16, optionally strip dead pre-pack q/k/v initializers, re-save.
+**Gotcha when re-saving with external data:** if you save with `location="X.data"` and then rename the
+file, the references inside `model.onnx` still point to `X.data`. Either save directly with
+`location="model.onnx.data"`, or rewrite each initializer's `external_data` `location` entry.
+
+```python
+import onnx, numpy as np
+from onnx import numpy_helper, TensorProto
+m = onnx.load("model.onnx", load_external_data=True)
+for init in m.graph.initializer:
+    if init.data_type == TensorProto.FLOAT:
+        arr = numpy_helper.to_array(init).astype(np.float16)
+        init.CopyFrom(numpy_helper.from_array(arr, init.name))
+onnx.save(m, "model.onnx", save_as_external_data=True, all_tensors_to_one_file=True,
+          location="model.onnx.data", size_threshold=1024, convert_attribute=False)
+```
+
+**Proper upstream fix:** set the packed-Concat output type to the model dtype in the GQA pack rewrite,
+or cast ALL float initializers at save time regardless of registered value_info; add an e2e test that
+loads the fp16 GQA export in onnxruntime.
+
+## 4. Always validate the export in ORT before profiling
+Load the model on `CUDAExecutionProvider` and run one prefill + one decode `session.run`. Confirm:
+(a) the expected attention op (`com.microsoft::GroupQueryAttention` vs `ai.onnx::Attention`),
+(b) finite fp16 logits, (c) no FLOAT32 initializers for an fp16 build.

From 140afae7a7e34a11d3f6db2ecece05c51fa02f53 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:26:37 +0000
Subject: [PATCH 02/38] fix(static-cache): use is_causal=0 + explicit causal
 mask for opset-24 Attention

mobius emitted static-cache Attention with is_causal=1 + nonpad_kv_seqlen, which
the opset-24 ONNX Attention CUDA kernel rejects when S_q != total_kv with no
past_key (causal_cross_no_past guard). With a pre-allocated max_seq_len cache
this fires in BOTH prefill and decode -> NOT_IMPLEMENTED at runtime.

Fix per ORT guidance: set is_causal=0 and pass an explicit 4D bool causal mask
[B,1,S_q,max_seq] built from write_indices (keep j <= write_indices[b]+t). Keeps
nonpad_kv_seqlen to select the external-cache kernel path. New helper
create_static_cache_causal_mask in _common.py.

Tests: 5 CPU value-level mask tests, updated/added 3 static-cache graph tests,
and a new e2e CUDA prefill+decode regression test (tests/static_cache_decode_test.py).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/components/_attention.py   |  56 ++++++---
 src/mobius/components/_common.py      |  80 ++++++++++++
 src/mobius/components/_common_test.py |  94 ++++++++++++++
 tests/build_graph_test.py             |  50 ++++++--
 tests/static_cache_decode_test.py     | 172 ++++++++++++++++++++++++++
 5 files changed, 425 insertions(+), 27 deletions(-)
 create mode 100644 tests/static_cache_decode_test.py

diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py
index 9c076d2e..a8205942 100644
--- a/src/mobius/components/_attention.py
+++ b/src/mobius/components/_attention.py
@@ -10,7 +10,7 @@
 from onnxscript import OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig
-from mobius.components._common import Linear
+from mobius.components._common import Linear, create_static_cache_causal_mask
 from mobius.components._rms_norm import OffsetRMSNorm, RMSNorm
 from mobius.components._rotary_embedding import apply_rotary_pos_emb
 
@@ -99,14 +99,19 @@ def _apply_attention(
 
     Static cache mode (``static_cache is not None``):
         Scatters new key/value into the static cache via TensorScatter,
-        then attends over the full cache using ``nonpad_kv_seqlen``.
-        Also uses ``is_causal=1``.
+        then attends over the full cache using ``nonpad_kv_seqlen`` with
+        ``is_causal=0`` plus an explicit causal mask derived from
+        ``write_indices`` (see :func:`create_static_cache_causal_mask`).
+        ``is_causal=1`` cannot be used here: the opset-24 Attention kernel
+        rejects it together with ``nonpad_kv_seqlen`` when ``S_q`` differs
+        from the (pre-allocated) cache length, which is always the case.
         Returns ``(attn_output, updated_key_cache, updated_value_cache)``.
 
     Note:
-        Both paths set ``is_causal=1`` on the Attention op, which enables
-        built-in causal masking. This means ``attn_mask`` should encode
-        only padding information (as a bool mask), not causality.
+        The dynamic path sets ``is_causal=1`` so callers only provide a
+        bool padding mask.  The static path instead uses ``is_causal=0``
+        with a full causal+padding mask because the external-cache kernel
+        does not accept ``is_causal=1`` alongside ``nonpad_kv_seqlen``.
 
     Note:
         ``nonpad_kv_seqlen`` (input #6) is only valid in static cache mode
@@ -138,28 +143,43 @@ def _apply_attention(
             axis=1,
         )  # [B, max_seq, kv_hidden]
 
-        # Attend over the full cache.  We pass None for attn_mask and use
-        # is_causal=1 instead — the Attention op handles causal + padding
-        # masking internally via is_causal + nonpad_kv_seqlen.  Using
-        # create_attention_bias() here would produce incorrect causality
-        # during prefill because it cannot represent the relationship
-        # between query positions and the full cache length.
+        # Attend over the full cache with is_causal=0 plus an explicit
+        # causal mask.  The opset-24 Attention CUDA kernel rejects
+        # is_causal=1 together with nonpad_kv_seqlen when S_q != total_kv
+        # and there is no past_key (the causal_cross_no_past guard in
+        # attention.cc).  With a pre-allocated [B, max_seq, ...] cache the
+        # query length never equals the total cache length, so that guard
+        # fires in BOTH prefill and decode.  Per ORT's own guidance we set
+        # is_causal=0 and pass an explicit causal mask built from
+        # write_indices: a query token at offset t attends to cache slots
+        # j <= write_indices[b] + t.  This single rule serves prefill
+        # (write_indices=0 -> triangular) and decode (write_indices=N,
+        # S_q=1 -> keep slots 0..N), and subsumes padding so it stays
+        # consistent with nonpad_kv_seqlen.
+        #
+        # nonpad_kv_seqlen is still passed: it selects the external-cache
+        # (TensorScatter) kernel path.  When both nonpad_kv_seqlen and an
+        # attn_mask are present, ORT skips Flash and routes to the
+        # memory-efficient / unfused path, which converts the bool mask to
+        # an additive bias (attention.cc ConvertAttnMaskToBias).
         #
-        # NOTE: The ONNX Attention spec supports attn_mask alongside
-        # nonpad_kv_seqlen for custom masking (e.g., user-defined masks
-        # beyond causal + padding).  Currently we rely on is_causal=1 +
-        # nonpad_kv_seqlen for standard LLM causal + padding masking.
         # TODO(titaiwang): Support user-provided attn_mask in external
         # cache mode for advanced use cases (e.g., prefix masking,
         # document boundaries in batched inference).
         # TODO(titaiwang): Support sliding window (circular cache mode)
         # with static cache for long-context models that use local
         # attention windows.
+        static_causal_mask = create_static_cache_causal_mask(
+            op,
+            query,
+            updated_k,
+            static_cache.write_indices,
+        )
         attn_output, _, _ = op.Attention(
             query,
             updated_k,
             updated_v,
-            None,  # no attn_mask — is_causal handles masking
+            static_causal_mask,  # explicit causal mask (is_causal=0)
             None,  # no past_key (full cache is already provided)
             None,  # no past_value
             static_cache.nonpad_kv_seqlen,
@@ -167,7 +187,7 @@ def _apply_attention(
             kv_num_heads=num_key_value_heads,
             scale=scale,
             softcap=softcap,
-            is_causal=1,
+            is_causal=0,
             _outputs=3,
         )
         return attn_output, updated_k, updated_v
diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py
index 2d4836f5..179c13b0 100644
--- a/src/mobius/components/_common.py
+++ b/src/mobius/components/_common.py
@@ -321,3 +321,83 @@ def create_sliding_window_mask(
     # Combine with padding mask
     padding_mask = op.Cast(op.Unsqueeze(attention_mask, [1]), to=ir.DataType.BOOL)
     return op.And(within_window, padding_mask)
+
+
+def create_static_cache_causal_mask(
+    op: OpBuilder,
+    query: ir.Value,
+    key_cache: ir.Value,
+    write_indices: ir.Value,
+):
+    """Build a causal attention mask for the static (TensorScatter) KV cache.
+
+    The opset-24 ONNX ``Attention`` CUDA kernel rejects ``is_causal=1``
+    together with ``nonpad_kv_seqlen`` whenever the query length differs
+    from the total KV length and there is no ``past_key`` (see
+    ``onnxruntime/core/providers/cuda/llm/attention.cc`` — the
+    ``causal_cross_no_past`` guard).  In the static cache the KV buffers
+    are pre-allocated to ``max_seq_len``, so ``S_q != total_kv`` in **both**
+    prefill and decode.  The fix, per ORT's own guidance, is to drive the
+    Attention op with ``is_causal=0`` and supply an explicit causal mask.
+
+    A query token at position ``t`` within this step writes into cache slot
+    ``write_indices[b] + t`` (``write_indices[b]`` is the number of valid
+    cache tokens *before* this step).  Causality means it may attend to
+    every cache slot ``j`` with ``j <= write_indices[b] + t``.  This single
+    rule serves both phases:
+
+    * **Prefill** (``write_indices=0``, ``S_q=N``): triangular causal mask
+      ``j <= t``.
+    * **Decode** (``write_indices=N``, ``S_q=1``): keep slots ``j <= N``,
+      i.e. all previously written tokens plus the just-written one.
+
+    Because padding slots ``j >= nonpad_kv_seqlen[b]`` are always greater
+    than ``write_indices[b] + t``, they are masked out too, so the mask is
+    consistent with the ``nonpad_kv_seqlen`` bounds (which are still passed
+    to the Attention op to select the external-cache kernel path).
+
+    The mask is 4D ``[batch, 1, S_q, max_seq]`` rather than 3D on purpose:
+    ``ConvertAttnMaskToBias`` (attention.cc) treats a 3D mask as
+    ``[heads, q, kv]`` and broadcasts over the batch dimension, which would
+    be incorrect because ``write_indices`` is per-batch.  A 4D mask with a
+    leading batch dim is honored per-batch (and ``dim1 == 1`` broadcasts
+    over heads).
+
+    Args:
+        op: The OpBuilder.
+        query: Query tensor ``[batch, S_q, hidden]``; only dims 0/1 are
+            read to derive the query sequence length ``S_q``.
+        key_cache: Pre-allocated key cache ``[batch, max_seq, kv_hidden]``;
+            dim 1 supplies the total KV length ``max_seq``.
+        write_indices: Per-batch write start position ``[batch]`` INT64 —
+            the number of valid cache tokens before this step.
+
+    Returns:
+        Bool mask ``[batch, 1, S_q, max_seq]``. ``True`` = attend,
+        ``False`` = mask out.
+    """
+    zero = op.Constant(value_int=0)
+    one = op.Constant(value_int=1)
+
+    # Scalar S_q (query length) and total KV length (max_seq) for Range.
+    q_len = op.Squeeze(op.Shape(query, start=1, end=2), op.Constant(value_ints=[0]))
+    total_kv = op.Squeeze(
+        op.Shape(key_cache, start=1, end=2), op.Constant(value_ints=[0])
+    )
+
+    # Per-step query offsets 0..S_q-1 and key slot indices 0..max_seq-1.
+    q_offsets = op.Range(zero, q_len, one)  # [S_q] int64
+    key_positions = op.Range(zero, total_kv, one)  # [max_seq] int64
+
+    # Absolute query positions: write_indices[b] + t  →  [batch, S_q].
+    query_positions = op.Add(
+        op.Unsqueeze(write_indices, [1]),  # [batch, 1]
+        op.Unsqueeze(q_offsets, [0]),  # [1, S_q]
+    )
+
+    # Reshape for broadcasting to [batch, 1, S_q, max_seq].
+    query_positions = op.Unsqueeze(query_positions, [1, 3])  # [batch, 1, S_q, 1]
+    key_positions = op.Unsqueeze(key_positions, [0, 1, 2])  # [1, 1, 1, max_seq]
+
+    # Keep key slot j for query at position p iff j <= p (causal + padding).
+    return op.GreaterOrEqual(query_positions, key_positions)
diff --git a/src/mobius/components/_common_test.py b/src/mobius/components/_common_test.py
index 3ebc1bf9..93bc2c4b 100644
--- a/src/mobius/components/_common_test.py
+++ b/src/mobius/components/_common_test.py
@@ -5,7 +5,9 @@
 
 from __future__ import annotations
 
+import numpy as np
 import onnx_ir as ir
+import onnxruntime as ort
 
 from mobius._testing import count_op_type, create_test_builder, create_test_input
 from mobius.components._common import (
@@ -13,9 +15,39 @@
     Linear,
     create_attention_bias,
     create_padding_mask,
+    create_static_cache_causal_mask,
 )
 
 
+def _run_static_cache_mask(max_seq: int, query_len: int, write_index: int) -> np.ndarray:
+    """Build the static-cache causal mask subgraph and evaluate it on CPU.
+
+    Returns the bool mask array of shape ``[1, 1, query_len, max_seq]`` where
+    ``True`` means the query token attends to that cache slot.
+    """
+    builder, op, graph = create_test_builder()
+    query = create_test_input(builder, "query", [1, "S_q", 32], dtype=ir.DataType.FLOAT)
+    key_cache = create_test_input(
+        builder, "key_cache", [1, max_seq, 16], dtype=ir.DataType.FLOAT
+    )
+    write_indices = create_test_input(
+        builder, "write_indices", [1], dtype=ir.DataType.INT64
+    )
+    mask = create_static_cache_causal_mask(op, query, key_cache, write_indices)
+    mask.name = "mask"
+    graph.outputs.append(mask)
+    proto = ir.to_proto(ir.Model(graph, ir_version=10))
+    session = ort.InferenceSession(
+        proto.SerializeToString(), providers=["CPUExecutionProvider"]
+    )
+    feeds = {
+        "query": np.zeros((1, query_len, 32), np.float32),
+        "key_cache": np.zeros((1, max_seq, 16), np.float32),
+        "write_indices": np.array([write_index], np.int64),
+    }
+    return session.run(None, feeds)[0]
+
+
 class TestLinear:
     def test_linear_with_bias(self):
         linear = Linear(64, 128, bias=True)
@@ -178,3 +210,65 @@ def test_uses_simpler_ops_than_attention_bias(self):
         assert count_op_type(graph_bias, "CumSum") >= 1
         assert count_op_type(graph_bias, "GreaterOrEqual") >= 1
         assert count_op_type(graph_bias, "Where") >= 1
+
+
+class TestCreateStaticCacheCausalMask:
+    """Value-level checks for the static-cache causal mask.
+
+    Query token ``t`` writes to absolute cache slot ``write_index + t`` and must
+    attend to every slot ``j <= write_index + t`` (causal), while padded/future
+    slots stay masked. This is verified by running the subgraph on CPU.
+    """
+
+    def test_structure_uses_range_and_greater_or_equal(self):
+        builder, op, graph = create_test_builder()
+        query = create_test_input(builder, "query", [1, "S_q", 32], dtype=ir.DataType.FLOAT)
+        key_cache = create_test_input(
+            builder, "key_cache", [1, 8, 16], dtype=ir.DataType.FLOAT
+        )
+        write_indices = create_test_input(
+            builder, "write_indices", [1], dtype=ir.DataType.INT64
+        )
+        mask = create_static_cache_causal_mask(op, query, key_cache, write_indices)
+        assert mask is not None
+        assert count_op_type(graph, "Range") >= 1
+        assert count_op_type(graph, "GreaterOrEqual") >= 1
+
+    def test_mask_shape_is_4d_batch_one_head(self):
+        mask = _run_static_cache_mask(max_seq=8, query_len=4, write_index=0)
+        # [B, 1, S_q, max_seq] — dim1=1 broadcasts over heads, dim0=batch.
+        assert mask.shape == (1, 1, 4, 8)
+        assert mask.dtype == bool
+
+    def test_prefill_is_triangular(self):
+        """write_index=0: each query token t attends to slots 0..t only."""
+        mask = _run_static_cache_mask(max_seq=8, query_len=4, write_index=0)
+        expected = np.array(
+            [
+                [1, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 0, 0, 0, 0],
+            ],
+            dtype=bool,
+        )
+        np.testing.assert_array_equal(mask[0, 0], expected)
+
+    def test_decode_keeps_prefix_through_write_index(self):
+        """Single decode token at slot N attends to slots 0..N, masks the rest."""
+        mask = _run_static_cache_mask(max_seq=8, query_len=1, write_index=3)
+        expected = np.array([[1, 1, 1, 1, 0, 0, 0, 0]], dtype=bool)
+        np.testing.assert_array_equal(mask[0, 0], expected)
+
+    def test_chunked_prefill_offset_write_index(self):
+        """Second prefill chunk (write_index=2, S_q=3) stays causal at absolute pos."""
+        mask = _run_static_cache_mask(max_seq=8, query_len=3, write_index=2)
+        expected = np.array(
+            [
+                [1, 1, 1, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 0, 0, 0],
+            ],
+            dtype=bool,
+        )
+        np.testing.assert_array_equal(mask[0, 0], expected)
diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index 7c15be06..dee8a9fb 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -4437,7 +4437,13 @@ def test_static_cache_graph_validates(self):
         assert len(proto.SerializeToString()) > 0
 
     def test_static_cache_attention_is_causal(self):
-        """Verify Attention ops use is_causal=1 in static cache mode."""
+        """Verify Attention ops use is_causal=0 in static cache mode.
+
+        The opset-24 Attention CUDA kernel rejects is_causal=1 together
+        with nonpad_kv_seqlen when S_q != total_kv with no past_key (always
+        true for a pre-allocated cache).  The static path therefore uses
+        is_causal=0 plus an explicit causal mask.
+        """
         model, config = self._build_static_cache_model()
 
         attention_nodes = [n for n in model.graph if n.op_type == "Attention"]
@@ -4448,25 +4454,51 @@ def test_static_cache_attention_is_causal(self):
             assert is_causal is not None, (
                 f"Attention node {node.name} missing is_causal attribute"
             )
-            assert is_causal.as_int() == 1, (
-                f"Attention node {node.name} should have is_causal=1"
+            assert is_causal.as_int() == 0, (
+                f"Attention node {node.name} should have is_causal=0 "
+                f"(causality is supplied via an explicit attn_mask)"
             )
 
-    def test_static_cache_attention_no_attn_mask_input(self):
-        """Verify Attention ops do NOT receive attn_mask in static cache mode."""
+    def test_static_cache_attention_has_causal_mask_input(self):
+        """Verify Attention ops receive an explicit causal attn_mask.
+
+        With is_causal=0, causality must come from input 3 (attn_mask).
+        The mask is produced by create_static_cache_causal_mask, whose
+        final op is a GreaterOrEqual yielding a bool mask.
+        """
         model, config = self._build_static_cache_model()
 
         attention_nodes = [n for n in model.graph if n.op_type == "Attention"]
         assert len(attention_nodes) == config.num_hidden_layers
 
         for node in attention_nodes:
-            # Input 3 (0-indexed) is attn_mask — should be empty/None
+            # Input 3 (0-indexed) is attn_mask — should be a real value now.
             attn_mask_input = node.inputs[3]
-            assert attn_mask_input is None or attn_mask_input.name == "", (
-                f"Attention node {node.name} should not have attn_mask "
-                f"connected, but got input: {attn_mask_input}"
+            assert attn_mask_input is not None and attn_mask_input.name != "", (
+                f"Attention node {node.name} should have an explicit causal "
+                f"attn_mask connected, but got: {attn_mask_input}"
+            )
+            producer = attn_mask_input.producer()
+            assert producer is not None and producer.op_type == "GreaterOrEqual", (
+                f"Attention node {node.name} attn_mask should be produced by "
+                f"the causal-mask GreaterOrEqual, got "
+                f"{None if producer is None else producer.op_type}"
             )
 
+    def test_static_cache_has_no_tensorscatter_left_unmasked(self):
+        """Static cache graph must contain the causal-mask construction ops.
+
+        Guards against regressing back to the is_causal=1 / no-mask form
+        that ORT rejects: the graph must build per-step query positions
+        (Range + Add) feeding a GreaterOrEqual mask.
+        """
+        model, _ = self._build_static_cache_model()
+        op_types = [n.op_type for n in model.graph]
+        assert "Range" in op_types, "Causal mask should use Range for positions"
+        assert "GreaterOrEqual" in op_types, (
+            "Causal mask should compare query/key positions with GreaterOrEqual"
+        )
+
     def test_static_cache_moe_graph_builds(self):
         """Build a MoE model (qwen2_moe) with static cache."""
         model, _config = self._build_static_cache_model(
diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
new file mode 100644
index 00000000..f9f3b1fb
--- /dev/null
+++ b/tests/static_cache_decode_test.py
@@ -0,0 +1,172 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""End-to-end regression test for the static (TensorScatter) KV cache.
+
+mobius used to emit the static-cache ``Attention`` op with ``is_causal=1``
+alongside ``nonpad_kv_seqlen``.  The opset-24 ONNX ``Attention`` CUDA kernel
+rejects that combination whenever the query length differs from the total KV
+length and there is no ``past_key`` (the ``causal_cross_no_past`` guard in
+``onnxruntime/core/providers/cuda/llm/attention.cc``).  Because the static
+cache is pre-allocated to ``max_seq_len``, that guard fires in **both**
+prefill (``S_q = N``) and decode (``S_q = 1``), raising ``NOT_IMPLEMENTED``.
+
+The fix sets ``is_causal=0`` and supplies an explicit causal mask
+(:func:`mobius.components._common.create_static_cache_causal_mask`).  This test
+exercises the actual ONNX Runtime kernel for both phases so the regression
+cannot silently come back.  It requires the CUDA Execution Provider because
+``TensorScatter`` and the external-cache ``Attention`` path are CUDA-only.
+
+The model is built fp32: the ``is_causal`` guard fires in ORT *before* kernel
+dtype dispatch, so fp32 exercises the same external-cache code path as a
+production fp16 export while avoiding the cos/sin-cache dtype casting that
+only the full CLI build pipeline applies.  A raw ``InferenceSession`` is used
+(instead of the ``OnnxModelSession`` test helper) to keep this regression
+guard free of the optional ``onnxruntime-easy`` dependency.
+"""
+
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import onnx_ir as ir
+import onnxruntime as ort
+import pytest
+from _test_configs import _base_config
+
+from mobius._registry import registry
+from mobius.tasks import CausalLMTask
+
+pytestmark = pytest.mark.skipif(
+    "CUDAExecutionProvider" not in ort.get_available_providers(),
+    reason="static-cache TensorScatter / external-cache Attention are CUDA-only",
+)
+
+_MAX_SEQ_LEN = 16
+_MODEL_TYPE = "qwen2"
+_CACHE_DTYPE = np.float32
+
+
+def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None:
+    """Fill empty initializers with small random values of their dtype.
+
+    The graph is built without real weights; ORT still needs concrete
+    initializers to run.  Small values keep logits finite and well-scaled.
+    """
+    for initializer in model.graph.initializers.values():
+        if initializer.const_value is not None:
+            continue
+        shape = initializer.shape
+        dims = [d if isinstance(d, int) else 1 for d in shape] if shape else [1]
+        dtype = initializer.dtype or ir.DataType.FLOAT
+        np_dtype = dtype.numpy()
+        if np.issubdtype(np_dtype, np.floating):
+            data = (rng.standard_normal(dims) * 0.02).astype(np_dtype)
+        else:
+            data = np.zeros(dims, dtype=np_dtype)
+        initializer.const_value = ir.Tensor(data)
+
+
+def _build_static_cache_session(
+    tmp_dir: str,
+) -> tuple[ort.InferenceSession, object]:
+    """Build a tiny static-cache qwen2 graph and load it on CUDA."""
+    config = _base_config()
+    module = registry.get(_MODEL_TYPE)(config)
+    task = CausalLMTask(static_cache=True, max_seq_len=_MAX_SEQ_LEN)
+    model = task.build(module, config)["model"]
+    _fill_random_weights(model, np.random.default_rng(0))
+
+    model_path = str(Path(tmp_dir) / "model.onnx")
+    ir.save(model, model_path, external_data="model.onnx.data")
+    session = ort.InferenceSession(
+        model_path,
+        providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+    )
+    assert "CUDAExecutionProvider" in session.get_providers(), (
+        "static-cache regression test must run on CUDA"
+    )
+    return session, config
+
+
+def _empty_caches(num_layers: int, kv_hidden: int) -> dict[str, np.ndarray]:
+    """Zeroed ``[1, max_seq, kv_hidden]`` cache buffers for every layer."""
+    feeds: dict[str, np.ndarray] = {}
+    for layer in range(num_layers):
+        zeros = np.zeros((1, _MAX_SEQ_LEN, kv_hidden), dtype=_CACHE_DTYPE)
+        feeds[f"key_cache.{layer}"] = zeros.copy()
+        feeds[f"value_cache.{layer}"] = zeros.copy()
+    return feeds
+
+
+def _carry_caches(
+    outputs: dict[str, np.ndarray], num_layers: int
+) -> dict[str, np.ndarray]:
+    """Feed the prefill ``updated_*`` caches back in as decode inputs."""
+    feeds: dict[str, np.ndarray] = {}
+    for layer in range(num_layers):
+        feeds[f"key_cache.{layer}"] = outputs[f"updated_key_cache.{layer}"]
+        feeds[f"value_cache.{layer}"] = outputs[f"updated_value_cache.{layer}"]
+    return feeds
+
+
+def test_static_cache_prefill_and_decode_run_on_cuda():
+    """Prefill (S_q>1) and decode (S_q=1) both run without NOT_IMPLEMENTED.
+
+    This is the regression guard the codebase previously lacked: it loads
+    a real static-cache graph and runs it on ``CUDAExecutionProvider`` for
+    both phases.  Reverting to ``is_causal=1`` (no mask) makes ORT raise
+    ``NOT_IMPLEMENTED`` here, failing the test.
+    """
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        session, config = _build_static_cache_session(tmp_dir)
+        num_layers = config.num_hidden_layers
+        kv_hidden = config.num_key_value_heads * config.head_dim
+        vocab = config.vocab_size
+        output_names = [out.name for out in session.get_outputs()]
+        rng = np.random.default_rng(1)
+
+        # --- Prefill: write N tokens from slot 0 (S_q = N != max_seq). ---
+        prefill_len = 4
+        prefill_feeds: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(
+                0, vocab, size=(1, prefill_len), dtype=np.int64
+            ),
+            "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
+            "write_indices": np.array([0], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
+        }
+        prefill_feeds.update(_empty_caches(num_layers, kv_hidden))
+
+        prefill_out = dict(
+            zip(output_names, session.run(output_names, prefill_feeds))
+        )
+        prefill_logits = prefill_out["logits"]
+        assert prefill_logits.shape == (1, prefill_len, vocab)
+        assert np.isfinite(prefill_logits).all(), "prefill logits must be finite"
+
+        # --- Decode: one token at slot N (S_q = 1 != max_seq). ---
+        decode_feeds: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64),
+            "position_ids": np.array([[prefill_len]], dtype=np.int64),
+            "write_indices": np.array([prefill_len], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([prefill_len + 1], dtype=np.int64),
+        }
+        decode_feeds.update(_carry_caches(prefill_out, num_layers))
+
+        decode_out = dict(
+            zip(output_names, session.run(output_names, decode_feeds))
+        )
+        decode_logits = decode_out["logits"]
+        assert decode_logits.shape == (1, 1, vocab)
+        assert np.isfinite(decode_logits).all(), "decode logits must be finite"
+
+        # The decode step must have scattered its key into slot N (the
+        # previously-empty tail), confirming the in-place cache advanced.
+        advanced_key = decode_out["updated_key_cache.0"]
+        assert advanced_key.shape == (1, _MAX_SEQ_LEN, kv_hidden)
+        assert np.any(advanced_key[0, prefill_len] != 0), (
+            "decode should scatter the new key into cache slot N"
+        )

From ce410677840d8af260e4fa6e636f87266ef53a58 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:37:12 +0000
Subject: [PATCH 03/38] fix(static-cache): phase-split static Attention to keep
 decode on Flash/XQA

Follow-up to the is_causal=0 fix per architect review. An always-present
attn_mask disables Flash Attention in ORT (the kernel selection checks
attn_mask != nullptr by pointer, not content), so a single always-masked
static-cache Attention forced single-token DECODE onto the memory-efficient
path -- not apples-to-apples with the GQA variant's Flash/XQA decode, which
contaminates the headline decode profiling metric.

Phase-split the static-cache attention behind an If keyed on Shape(query)[1]>1:
- multi-token step (S_q>1, prefill / chunked / speculative): explicit causal
  mask -> memory-efficient path (unavoidable; static buffer makes K_seq=total
  so Flash prefill is guard-blocked regardless, and it is the amortized path).
- single-token decode (S_q==1): omit attn_mask; nonpad_kv_seqlen alone bounds
  attention to the valid prefix -> stays on Flash/XQA.

New helper _attend_over_static_cache builds the two-branch If (reusing
rename_subgraph_values for SSA-safe subgraphs). Validated through the full
build_from_module optimize pipeline + a CUDA prefill+decode run.

Tests: recurse into If subgraphs for op assertions; add phase-split mask-presence
and If-structure tests; upgrade the e2e regression to build_from_module so the
real export pipeline is exercised.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/components/_attention.py | 165 +++++++++++++++++++++-------
 tests/build_graph_test.py           | 134 +++++++++++++++-------
 tests/static_cache_decode_test.py   |  28 +++--
 3 files changed, 240 insertions(+), 87 deletions(-)

diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py
index a8205942..7ee1bd7d 100644
--- a/src/mobius/components/_attention.py
+++ b/src/mobius/components/_attention.py
@@ -7,12 +7,114 @@
 from typing import NamedTuple
 
 import onnx_ir as ir
-from onnxscript import OpBuilder, nn
+from onnxscript import GraphBuilder, OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig
+from mobius._constants import OPSET_VERSION
 from mobius.components._common import Linear, create_static_cache_causal_mask
 from mobius.components._rms_norm import OffsetRMSNorm, RMSNorm
 from mobius.components._rotary_embedding import apply_rotary_pos_emb
+from mobius.components._scan_utils import rename_subgraph_values
+
+
+def _attend_over_static_cache(
+    op: OpBuilder,
+    query: ir.Value,
+    key_cache: ir.Value,
+    value_cache: ir.Value,
+    write_indices: ir.Value,
+    nonpad_kv_seqlen: ir.Value,
+    num_attention_heads: int,
+    num_key_value_heads: int,
+    scale: float | None,
+    softcap: float | None,
+) -> ir.Value:
+    """Attend over the static KV cache, phase-split for decode kernel fidelity.
+
+    The opset-24 ONNX ``Attention`` CUDA kernel rejects ``is_causal=1``
+    together with ``nonpad_kv_seqlen`` when ``S_q != total_kv`` with no
+    ``past_key`` (the ``causal_cross_no_past`` guard in ``attention.cc``).
+    With a pre-allocated ``[B, max_seq, ...]`` cache that condition holds in
+    both prefill and decode, so ``is_causal=0`` must be used.
+
+    Both phases run with ``is_causal=0`` and keep ``nonpad_kv_seqlen`` (which
+    selects the external-cache / TensorScatter kernel path).  The phases
+    differ in whether an explicit ``attn_mask`` is supplied, because the
+    *presence* of ``attn_mask`` — regardless of its contents — disables Flash
+    Attention in ORT (``attn_mask != nullptr`` routes to the memory-efficient
+    or unfused path; see the kernel-selection cascade in ``attention.cc``):
+
+    * **Multi-token step** (``S_q > 1``: prefill or speculative/chunked decode):
+      needs intra-query causality, so it passes an explicit causal mask built
+      from ``write_indices`` (:func:`create_static_cache_causal_mask`) and
+      therefore runs on the memory-efficient path.  This is unavoidable — the
+      static buffer makes ``K_seq == total``, so Flash prefill is blocked by
+      the same guard regardless — and is the cheap, amortized path anyway.
+    * **Single-token decode** (``S_q == 1``): a lone query needs no
+      intra-query causal mask; ``nonpad_kv_seqlen`` alone bounds attention to
+      the valid prefix ``0..write_indices[b]``.  Omitting ``attn_mask`` keeps
+      this hot path on Flash / XQA — the kernel the GQA variant also uses,
+      so the comparison stays apples-to-apples.
+
+    The two phases are emitted as the branches of an ``If`` keyed on
+    ``Shape(query)[1] > 1`` so a single exported graph serves both, while
+    decode structurally omits the mask input.
+
+    Returns:
+        The attention output for the active phase, shape ``[B, S_q, hidden]``.
+    """
+    seq_len = op.Squeeze(
+        op.Shape(query, start=1, end=2), op.Constant(value_ints=[0])
+    )
+    is_multi_token_step = op.Greater(seq_len, op.Constant(value_int=1))
+
+    def _build_attention_branch(name: str, use_causal_mask: bool) -> ir.Graph:
+        branch = ir.Graph(
+            [], [], nodes=[], name=name, opset_imports={"": OPSET_VERSION}
+        )
+        branch_op = GraphBuilder(branch).op
+        attn_mask = (
+            create_static_cache_causal_mask(
+                branch_op, query, key_cache, write_indices
+            )
+            if use_causal_mask
+            else None
+        )
+        attn_output, _, _ = branch_op.Attention(
+            query,
+            key_cache,
+            value_cache,
+            attn_mask,
+            None,  # no past_key (full cache is already provided)
+            None,  # no past_value
+            nonpad_kv_seqlen,
+            q_num_heads=num_attention_heads,
+            kv_num_heads=num_key_value_heads,
+            scale=scale,
+            softcap=softcap,
+            is_causal=0,
+            _outputs=3,
+        )
+        # Prefix internal node/value names so the two branches stay in SSA
+        # form when merged under the parent graph, then pin the branch
+        # output name (the If wires branches by output position).
+        rename_subgraph_values(branch, f"{name}_")
+        attn_output.name = f"{name}_attn_output"
+        branch.outputs.append(attn_output)
+        return branch
+
+    prefill_branch = _build_attention_branch(
+        "static_cache_prefill", use_causal_mask=True
+    )
+    decode_branch = _build_attention_branch(
+        "static_cache_decode", use_causal_mask=False
+    )
+    return op.If(
+        is_multi_token_step,
+        then_branch=prefill_branch,
+        else_branch=decode_branch,
+        _outputs=1,
+    )
 
 
 class GQAContext(NamedTuple):
@@ -100,18 +202,22 @@ def _apply_attention(
     Static cache mode (``static_cache is not None``):
         Scatters new key/value into the static cache via TensorScatter,
         then attends over the full cache using ``nonpad_kv_seqlen`` with
-        ``is_causal=0`` plus an explicit causal mask derived from
-        ``write_indices`` (see :func:`create_static_cache_causal_mask`).
-        ``is_causal=1`` cannot be used here: the opset-24 Attention kernel
-        rejects it together with ``nonpad_kv_seqlen`` when ``S_q`` differs
-        from the (pre-allocated) cache length, which is always the case.
+        ``is_causal=0``.  Causality is phase-split (see
+        :func:`_attend_over_static_cache`): multi-token steps (``S_q > 1``)
+        use an explicit causal mask derived from ``write_indices`` (memory-
+        efficient path), while single-token decode (``S_q == 1``) omits the
+        mask to stay on Flash/XQA.  ``is_causal=1`` cannot be used here: the
+        opset-24 Attention kernel rejects it together with
+        ``nonpad_kv_seqlen`` when ``S_q`` differs from the (pre-allocated)
+        cache length, which is always the case.
         Returns ``(attn_output, updated_key_cache, updated_value_cache)``.
 
     Note:
         The dynamic path sets ``is_causal=1`` so callers only provide a
         bool padding mask.  The static path instead uses ``is_causal=0``
-        with a full causal+padding mask because the external-cache kernel
-        does not accept ``is_causal=1`` alongside ``nonpad_kv_seqlen``.
+        because the external-cache kernel does not accept ``is_causal=1``
+        alongside ``nonpad_kv_seqlen``; causality is supplied per phase via
+        an explicit mask (prefill) or ``nonpad_kv_seqlen`` alone (decode).
 
     Note:
         ``nonpad_kv_seqlen`` (input #6) is only valid in static cache mode
@@ -143,25 +249,13 @@ def _apply_attention(
             axis=1,
         )  # [B, max_seq, kv_hidden]
 
-        # Attend over the full cache with is_causal=0 plus an explicit
-        # causal mask.  The opset-24 Attention CUDA kernel rejects
-        # is_causal=1 together with nonpad_kv_seqlen when S_q != total_kv
-        # and there is no past_key (the causal_cross_no_past guard in
-        # attention.cc).  With a pre-allocated [B, max_seq, ...] cache the
-        # query length never equals the total cache length, so that guard
-        # fires in BOTH prefill and decode.  Per ORT's own guidance we set
-        # is_causal=0 and pass an explicit causal mask built from
-        # write_indices: a query token at offset t attends to cache slots
-        # j <= write_indices[b] + t.  This single rule serves prefill
-        # (write_indices=0 -> triangular) and decode (write_indices=N,
-        # S_q=1 -> keep slots 0..N), and subsumes padding so it stays
-        # consistent with nonpad_kv_seqlen.
-        #
-        # nonpad_kv_seqlen is still passed: it selects the external-cache
-        # (TensorScatter) kernel path.  When both nonpad_kv_seqlen and an
-        # attn_mask are present, ORT skips Flash and routes to the
-        # memory-efficient / unfused path, which converts the bool mask to
-        # an additive bias (attention.cc ConvertAttnMaskToBias).
+        # Attend over the full cache.  Both phases use is_causal=0 (the
+        # opset-24 Attention kernel rejects is_causal=1 + nonpad_kv_seqlen
+        # for a pre-allocated cache) and keep nonpad_kv_seqlen to select the
+        # external-cache kernel path.  Causality is enforced per-phase: the
+        # multi-token branch supplies an explicit causal mask (MEA path),
+        # while single-token decode omits the mask to stay on Flash/XQA.
+        # See _attend_over_static_cache for the full rationale.
         #
         # TODO(titaiwang): Support user-provided attn_mask in external
         # cache mode for advanced use cases (e.g., prefix masking,
@@ -169,26 +263,17 @@ def _apply_attention(
         # TODO(titaiwang): Support sliding window (circular cache mode)
         # with static cache for long-context models that use local
         # attention windows.
-        static_causal_mask = create_static_cache_causal_mask(
+        attn_output = _attend_over_static_cache(
             op,
-            query,
-            updated_k,
-            static_cache.write_indices,
-        )
-        attn_output, _, _ = op.Attention(
             query,
             updated_k,
             updated_v,
-            static_causal_mask,  # explicit causal mask (is_causal=0)
-            None,  # no past_key (full cache is already provided)
-            None,  # no past_value
+            static_cache.write_indices,
             static_cache.nonpad_kv_seqlen,
-            q_num_heads=num_attention_heads,
-            kv_num_heads=num_key_value_heads,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
             scale=scale,
             softcap=softcap,
-            is_causal=0,
-            _outputs=3,
         )
         return attn_output, updated_k, updated_v
 
diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index dee8a9fb..81dfe45b 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -4338,6 +4338,23 @@ class TestBuildStaticCacheGraph:
 
     MAX_SEQ_LEN = 128
 
+    @staticmethod
+    def _walk_nodes(graph):
+        """Yield every node in ``graph``, recursing into subgraphs (e.g. If).
+
+        The static-cache attention is emitted as an ``If`` whose branches
+        each contain an ``Attention`` op, so structural assertions must look
+        inside subgraph attributes rather than only the top-level graph.
+        """
+        for node in graph:
+            yield node
+            for attr in node.attributes.values():
+                if attr.type == ir.AttributeType.GRAPH and attr.value is not None:
+                    yield from TestBuildStaticCacheGraph._walk_nodes(attr.value)
+                elif attr.type == ir.AttributeType.GRAPHS:
+                    for subgraph in attr.value:
+                        yield from TestBuildStaticCacheGraph._walk_nodes(subgraph)
+
     def _build_static_cache_model(self, model_type: str = "qwen2", **config_overrides):
         """Build a model with CausalLMTask(static_cache=True) and return (model, config)."""
         from mobius.tasks import CausalLMTask
@@ -4413,12 +4430,20 @@ def test_static_cache_graph_outputs(self):
         )
 
     def test_static_cache_has_tensorscatter_and_attention(self):
-        """Verify graph contains TensorScatter and Attention ops."""
+        """Verify graph contains TensorScatter and Attention ops.
+
+        TensorScatter lives at the top level (shared by both phases);
+        Attention lives inside the phase-split If branches, so the search
+        recurses into subgraphs.
+        """
         model, _ = self._build_static_cache_model()
 
-        op_types = {n.op_type for n in model.graph}
+        op_types = {n.op_type for n in self._walk_nodes(model.graph)}
         assert "TensorScatter" in op_types, "Static cache graph should use TensorScatter"
         assert "Attention" in op_types, "Static cache graph should use Attention"
+        assert "If" in op_types, (
+            "Static cache attention should be phase-split behind an If"
+        )
 
     def test_static_cache_has_initializers(self):
         """Verify the graph has model parameters."""
@@ -4437,17 +4462,21 @@ def test_static_cache_graph_validates(self):
         assert len(proto.SerializeToString()) > 0
 
     def test_static_cache_attention_is_causal(self):
-        """Verify Attention ops use is_causal=0 in static cache mode.
+        """Verify every Attention op uses is_causal=0 in static cache mode.
 
         The opset-24 Attention CUDA kernel rejects is_causal=1 together
         with nonpad_kv_seqlen when S_q != total_kv with no past_key (always
-        true for a pre-allocated cache).  The static path therefore uses
-        is_causal=0 plus an explicit causal mask.
+        true for a pre-allocated cache).  Both phase-split branches therefore
+        use is_causal=0; causality comes from an explicit mask (prefill) or
+        nonpad_kv_seqlen alone (decode).
         """
         model, config = self._build_static_cache_model()
 
-        attention_nodes = [n for n in model.graph if n.op_type == "Attention"]
-        assert len(attention_nodes) == config.num_hidden_layers
+        attention_nodes = [
+            n for n in self._walk_nodes(model.graph) if n.op_type == "Attention"
+        ]
+        # Two branches (prefill + decode) per layer.
+        assert len(attention_nodes) == 2 * config.num_hidden_layers
 
         for node in attention_nodes:
             is_causal = node.attributes.get("is_causal")
@@ -4456,47 +4485,71 @@ def test_static_cache_attention_is_causal(self):
             )
             assert is_causal.as_int() == 0, (
                 f"Attention node {node.name} should have is_causal=0 "
-                f"(causality is supplied via an explicit attn_mask)"
+                f"(causality is supplied via an explicit attn_mask or "
+                f"nonpad_kv_seqlen)"
             )
 
-    def test_static_cache_attention_has_causal_mask_input(self):
-        """Verify Attention ops receive an explicit causal attn_mask.
+    def test_static_cache_phase_split_mask_presence(self):
+        """Prefill branch carries an explicit causal mask; decode omits it.
 
-        With is_causal=0, causality must come from input 3 (attn_mask).
-        The mask is produced by create_static_cache_causal_mask, whose
-        final op is a GreaterOrEqual yielding a bool mask.
+        This is the Flash-eligibility invariant: ORT disables Flash whenever
+        attn_mask is present (by pointer, not content), so single-token
+        decode MUST structurally omit the mask to stay on Flash/XQA, while
+        multi-token prefill supplies a causal mask (memory-efficient path).
         """
         model, config = self._build_static_cache_model()
 
-        attention_nodes = [n for n in model.graph if n.op_type == "Attention"]
-        assert len(attention_nodes) == config.num_hidden_layers
-
+        attention_nodes = [
+            n for n in self._walk_nodes(model.graph) if n.op_type == "Attention"
+        ]
+        with_mask = []
+        without_mask = []
         for node in attention_nodes:
-            # Input 3 (0-indexed) is attn_mask — should be a real value now.
-            attn_mask_input = node.inputs[3]
-            assert attn_mask_input is not None and attn_mask_input.name != "", (
-                f"Attention node {node.name} should have an explicit causal "
-                f"attn_mask connected, but got: {attn_mask_input}"
-            )
-            producer = attn_mask_input.producer()
-            assert producer is not None and producer.op_type == "GreaterOrEqual", (
-                f"Attention node {node.name} attn_mask should be produced by "
-                f"the causal-mask GreaterOrEqual, got "
-                f"{None if producer is None else producer.op_type}"
-            )
+            attn_mask_input = node.inputs[3] if len(node.inputs) > 3 else None
+            if attn_mask_input is not None and attn_mask_input.name != "":
+                producer = attn_mask_input.producer()
+                assert producer is not None and producer.op_type == "GreaterOrEqual", (
+                    f"Masked Attention {node.name} attn_mask should come from "
+                    f"the causal-mask GreaterOrEqual, got "
+                    f"{None if producer is None else producer.op_type}"
+                )
+                with_mask.append(node)
+            else:
+                without_mask.append(node)
+
+        # Exactly one masked (prefill) and one maskless (decode) per layer.
+        assert len(with_mask) == config.num_hidden_layers, (
+            f"Expected {config.num_hidden_layers} masked (prefill) Attention "
+            f"ops, got {len(with_mask)}"
+        )
+        assert len(without_mask) == config.num_hidden_layers, (
+            f"Expected {config.num_hidden_layers} maskless (decode) Attention "
+            f"ops for Flash eligibility, got {len(without_mask)}"
+        )
 
-    def test_static_cache_has_no_tensorscatter_left_unmasked(self):
-        """Static cache graph must contain the causal-mask construction ops.
+    def test_static_cache_is_phase_split_behind_if(self):
+        """Static cache attention must be a per-layer If over masked/maskless.
 
-        Guards against regressing back to the is_causal=1 / no-mask form
-        that ORT rejects: the graph must build per-step query positions
-        (Range + Add) feeding a GreaterOrEqual mask.
+        Guards against regressing to (a) is_causal=1 (ORT-rejected) or
+        (b) a single always-masked Attention that would force decode off
+        Flash.  Expect one If per layer plus the causal-mask ops (Range +
+        GreaterOrEqual) used only in the prefill branch.
         """
-        model, _ = self._build_static_cache_model()
-        op_types = [n.op_type for n in model.graph]
-        assert "Range" in op_types, "Causal mask should use Range for positions"
-        assert "GreaterOrEqual" in op_types, (
-            "Causal mask should compare query/key positions with GreaterOrEqual"
+        model, config = self._build_static_cache_model()
+        op_counts: dict[str, int] = {}
+        for node in self._walk_nodes(model.graph):
+            op_counts[node.op_type] = op_counts.get(node.op_type, 0) + 1
+
+        assert op_counts.get("If", 0) == config.num_hidden_layers, (
+            "Each layer's static-cache attention should be phase-split via If"
+        )
+        # Causal mask (prefill branch only): one Range + GreaterOrEqual/layer.
+        assert op_counts.get("Range", 0) >= config.num_hidden_layers, (
+            "Prefill branch should build query positions with Range"
+        )
+        assert op_counts.get("GreaterOrEqual", 0) == config.num_hidden_layers, (
+            "Causal mask (GreaterOrEqual) should appear once per layer "
+            "(prefill branch only — decode is maskless)"
         )
 
     def test_static_cache_moe_graph_builds(self):
@@ -4518,8 +4571,9 @@ def test_static_cache_moe_graph_builds(self):
         assert "position_ids" in input_names
         assert "attention_mask" not in input_names
 
-        # Verify TensorScatter and Attention ops are present
-        op_types = {n.op_type for n in model.graph}
+        # Verify TensorScatter (top level) and Attention (inside the
+        # phase-split If branches) ops are present.
+        op_types = {n.op_type for n in self._walk_nodes(model.graph)}
         assert "TensorScatter" in op_types
         assert "Attention" in op_types
 
diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index f9f3b1fb..08effb20 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -11,11 +11,15 @@
 cache is pre-allocated to ``max_seq_len``, that guard fires in **both**
 prefill (``S_q = N``) and decode (``S_q = 1``), raising ``NOT_IMPLEMENTED``.
 
-The fix sets ``is_causal=0`` and supplies an explicit causal mask
-(:func:`mobius.components._common.create_static_cache_causal_mask`).  This test
-exercises the actual ONNX Runtime kernel for both phases so the regression
-cannot silently come back.  It requires the CUDA Execution Provider because
-``TensorScatter`` and the external-cache ``Attention`` path are CUDA-only.
+The fix sets ``is_causal=0`` and phase-splits the attention behind an ``If``
+keyed on ``Shape(query)[1] > 1``: the multi-token (prefill) branch supplies an
+explicit causal mask (:func:`mobius.components._common.create_static_cache_causal_mask`,
+memory-efficient path), while the single-token decode branch omits the mask so
+ORT keeps it on Flash/XQA — the same kernel the GQA variant uses, so the
+profiling comparison stays apples-to-apples.  This test exercises the actual
+ONNX Runtime kernel for both phases so the regression cannot silently come
+back.  It requires the CUDA Execution Provider because ``TensorScatter`` and
+the external-cache ``Attention`` path are CUDA-only.
 
 The model is built fp32: the ``is_causal`` guard fires in ORT *before* kernel
 dtype dispatch, so fp32 exercises the same external-cache code path as a
@@ -36,6 +40,7 @@
 import pytest
 from _test_configs import _base_config
 
+from mobius._builder import build_from_module
 from mobius._registry import registry
 from mobius.tasks import CausalLMTask
 
@@ -72,11 +77,20 @@ def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None:
 def _build_static_cache_session(
     tmp_dir: str,
 ) -> tuple[ort.InferenceSession, object]:
-    """Build a tiny static-cache qwen2 graph and load it on CUDA."""
+    """Build a tiny static-cache qwen2 graph and load it on CUDA.
+
+    Uses the full ``build_from_module`` export path (not bare ``task.build``)
+    so the phase-split ``If`` subgraphs are exercised through the real
+    ``optimize_model`` pipeline — that is where a structural regression in
+    the static-cache attention would surface.
+    """
     config = _base_config()
     module = registry.get(_MODEL_TYPE)(config)
     task = CausalLMTask(static_cache=True, max_seq_len=_MAX_SEQ_LEN)
-    model = task.build(module, config)["model"]
+    package = build_from_module(
+        module, config, task=task, execution_provider="default"
+    )
+    model = package["model"]
     _fill_random_weights(model, np.random.default_rng(0))
 
     model_path = str(Path(tmp_dir) / "model.onnx")

From ed57be8bb0e3901bbe1250c0231039597e6f9773 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:41:20 +0000
Subject: [PATCH 04/38] docs(static-cache): clarify Flash->MEA decode-latency
 rationale

Per readability review: make explicit that the attn_mask presence routes to
the *slower* memory-efficient path, and that the phase split exists precisely
to pay that Flash->MEA cost only on prefill, never on the per-token decode
hot path. Doc-only; no behavior change.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/components/_attention.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py
index 7ee1bd7d..93b57dd9 100644
--- a/src/mobius/components/_attention.py
+++ b/src/mobius/components/_attention.py
@@ -41,8 +41,11 @@ def _attend_over_static_cache(
     selects the external-cache / TensorScatter kernel path).  The phases
     differ in whether an explicit ``attn_mask`` is supplied, because the
     *presence* of ``attn_mask`` — regardless of its contents — disables Flash
-    Attention in ORT (``attn_mask != nullptr`` routes to the memory-efficient
-    or unfused path; see the kernel-selection cascade in ``attention.cc``):
+    Attention in ORT (``attn_mask != nullptr`` routes to the slower memory-
+    efficient or unfused path; see the kernel-selection cascade in
+    ``attention.cc``).  The phase split exists precisely to pay that
+    Flash→MEA latency cost only where it is unavoidable (multi-token prefill)
+    and never on the per-token decode hot path:
 
     * **Multi-token step** (``S_q > 1``: prefill or speculative/chunked decode):
       needs intra-query causality, so it passes an explicit causal mask built

From df203cce92b88acde5293f7bcc2447a8d9604667 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:42:48 +0000
Subject: [PATCH 05/38] Fix fp16 GQA export emitting fp32 packed weights in
 fold passes

When building fp16 models, _cast_module_dtype casts params to fp16 but
the resulting initializer Values lose their declared .dtype (None) while
const_value stays fp16. FoldConcatInitializersPass and
FoldTransposedInitializerPass then defaulted the folded initializer's
dtype to FLOAT, serializing the packed QKV / transposed weights as fp32.
ORT rejected the model with a fp16/fp32 MatMul type-parameter error on
both CPU and CUDA EPs, breaking GQA export.

- Add shared helper _dtype_utils.initializer_dtype() that resolves the
  effective dtype from the declared type, falling back to const_value
  when the type annotation was dropped; prefers the data dtype and warns
  on stale-metadata disagreement.
- Use it in both fold passes to stamp the correct dtype on the new
  initializer's TensorType and LazyTensor.
- Guard FoldConcatInitializersPass against folding before weights load
  (mirrors FoldTransposedInitializerPass).
- Add regression tests, including an end-to-end ORT CPU-EP load test that
  reproduces the original MatMul fp16/fp32 failure without the fix.

Verified end-to-end: native fp16 Phi-3.5 GQA export now loads in ORT
CUDA EP with no manual post-cast (32 GroupQueryAttention nodes, all
fp16 initializers).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_dtype_utils.py         |  55 +++++++++
 src/mobius/_passes/_dtype_utils_test.py    |  41 +++++++
 src/mobius/_passes/_fold_concat.py         |  37 +++++-
 src/mobius/_passes/_fold_concat_test.py    | 130 +++++++++++++++++++++
 src/mobius/_passes/_fold_transpose.py      |  15 ++-
 src/mobius/_passes/_fold_transpose_test.py |  50 ++++++++
 6 files changed, 321 insertions(+), 7 deletions(-)
 create mode 100644 src/mobius/_passes/_dtype_utils.py
 create mode 100644 src/mobius/_passes/_dtype_utils_test.py

diff --git a/src/mobius/_passes/_dtype_utils.py b/src/mobius/_passes/_dtype_utils.py
new file mode 100644
index 00000000..d5368dd2
--- /dev/null
+++ b/src/mobius/_passes/_dtype_utils.py
@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Shared dtype helpers for graph passes that materialize new initializers.
+
+Passes such as :class:`~mobius._passes.FoldConcatInitializersPass` and
+:class:`~mobius._passes.FoldTransposedInitializerPass` pre-compute new
+initializers from existing ones. They must stamp the *correct* dtype on the
+result, otherwise an fp16 model can silently end up with fp32 weights that
+onnxruntime rejects at load time (a MatMul binding fp16 and fp32 to the same
+type parameter ``T``).
+"""
+
+from __future__ import annotations
+
+import logging
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+def initializer_dtype(value: ir.Value) -> ir.DataType | None:
+    """Return the effective dtype of an initializer ``value``.
+
+    Uses the value's declared ``type`` dtype, but falls back to the dtype of its
+    ``const_value`` when the type annotation is missing. When both are present
+    but disagree, the ``const_value`` dtype wins (it is the data actually
+    serialized) and a warning is logged, since a healthy initializer should
+    never have a declared type that contradicts its data.
+
+    Graph building can drop the declared ``type`` on an initializer while its
+    actual tensor data (``const_value``) still carries the correct dtype. In
+    that situation, defaulting to ``ir.DataType.FLOAT`` would emit fp32 weights
+    into an otherwise fp16 model. Reading the dtype from ``const_value`` keeps
+    folded initializers consistent with the weights they are derived from.
+
+    Returns ``None`` only when neither the declared type nor ``const_value`` is
+    available; callers decide on a final fallback.
+    """
+    declared = value.dtype
+    const_dtype = value.const_value.dtype if value.const_value is not None else None
+
+    if declared is not None and const_dtype is not None and declared != const_dtype:
+        logger.warning(
+            "Initializer %r declares dtype %s but its data is %s; using the data "
+            "dtype. This indicates stale type metadata.",
+            value.name,
+            declared,
+            const_dtype,
+        )
+        return const_dtype
+    if declared is not None:
+        return declared
+    return const_dtype
diff --git a/src/mobius/_passes/_dtype_utils_test.py b/src/mobius/_passes/_dtype_utils_test.py
new file mode 100644
index 00000000..881ab0d7
--- /dev/null
+++ b/src/mobius/_passes/_dtype_utils_test.py
@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for the shared initializer dtype helper."""
+
+from __future__ import annotations
+
+import numpy as np
+import onnx_ir as ir
+
+from mobius._passes._dtype_utils import initializer_dtype
+
+
+def _value(dtype: ir.DataType | None, const: np.ndarray | None) -> ir.Value:
+    v = ir.Value(name="w")
+    if dtype is not None:
+        v.dtype = dtype
+    if const is not None:
+        v.const_value = ir.tensor(const)
+    return v
+
+
+class TestInitializerDtype:
+    def test_uses_declared_dtype_when_present(self):
+        v = _value(ir.DataType.FLOAT16, np.ones((2,), np.float16))
+        assert initializer_dtype(v) == ir.DataType.FLOAT16
+
+    def test_falls_back_to_const_value_when_declared_missing(self):
+        """The core fix: a dropped declared type must not hide the real dtype."""
+        v = _value(None, np.ones((2,), np.float16))
+        assert v.dtype is None
+        assert initializer_dtype(v) == ir.DataType.FLOAT16
+
+    def test_const_value_wins_on_disagreement(self):
+        """Stale declared metadata must not override the serialized data dtype."""
+        v = _value(ir.DataType.FLOAT, np.ones((2,), np.float16))
+        assert initializer_dtype(v) == ir.DataType.FLOAT16
+
+    def test_returns_none_when_nothing_available(self):
+        v = _value(None, None)
+        assert initializer_dtype(v) is None
diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py
index b8f2eed9..2e1e926d 100644
--- a/src/mobius/_passes/_fold_concat.py
+++ b/src/mobius/_passes/_fold_concat.py
@@ -27,6 +27,8 @@
 import numpy as np
 import onnx_ir as ir
 
+from mobius._passes._dtype_utils import initializer_dtype
+
 logger = logging.getLogger(__name__)
 
 
@@ -63,6 +65,19 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
             if not all(v is not None and v.is_initializer() for v in inputs):
                 continue
 
+            # All inputs must have loaded weights. Folding before weights are
+            # loaded would force the dtype to default to FLOAT and bake a wrong
+            # type into the packed initializer. Mirror FoldTransposedInitializer's
+            # guard and skip with a warning so the Concat survives for a later run.
+            if any(v.const_value is None for v in inputs):  # type: ignore[union-attr]
+                logger.warning(
+                    "FoldConcatInitializersPass: skipping Concat %r — an input "
+                    "initializer has no const_value (pass ran before weights "
+                    "were loaded, or a weight is missing).",
+                    node.name,
+                )
+                continue
+
             axis_attr = node.attributes.get("axis")
             axis: int = axis_attr.value if axis_attr is not None else 0
 
@@ -79,8 +94,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
 
             # Require uniform dtype — mixed-dtype concat is unusual and likely
             # a modelling error; skip and warn rather than silently produce
-            # a wrong result.
-            dtypes = [v.dtype for v in inputs]  # type: ignore[union-attr]
+            # a wrong result. Resolve each dtype from the declared type or, when
+            # that is missing, from the loaded ``const_value`` so an fp16 weight
+            # whose type annotation was dropped is not mistaken for fp32.
+            dtypes = [initializer_dtype(v) for v in inputs]  # type: ignore[union-attr]
             if len(set(dtypes)) > 1:
                 logger.warning(
                     "FoldConcatInitializersPass: skipping Concat with mixed dtypes %s"
@@ -89,6 +106,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                 )
                 continue
 
+            # Authoritative dtype for the packed initializer. Falls back to FLOAT
+            # only when neither the declared type nor const_value is available.
+            packed_dtype = dtypes[0] or ir.DataType.FLOAT
+
             # Build a name for the packed initializer from the input names and axis.
             # The name encodes both the ordered input names and the axis, so two Concat
             # nodes with the same inputs in the same order along the same axis will
@@ -117,7 +138,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                 )
                 continue
 
-            new_val = ir.Value(name=packed_name, shape=out_shape, type=out_val.type)
+            # Stamp the resolved dtype on the new initializer's type so the
+            # declared type, the LazyTensor, and the materialized data all agree
+            # (out_val.type may be missing after stage-2 rewrites).
+            new_val = ir.Value(
+                name=packed_name, shape=out_shape, type=ir.TensorType(packed_dtype)
+            )
 
             captured_inputs = list(inputs)  # capture for closure
 
@@ -133,8 +159,9 @@ def _make_packed(
                     arrays.append(v.const_value.numpy())
                 return ir.tensor(np.concatenate(arrays, axis=ax))
 
-            dtype = inputs[0].dtype or ir.DataType.FLOAT  # type: ignore[union-attr]
-            new_val.const_value = ir.LazyTensor(_make_packed, dtype=dtype, shape=out_shape)
+            new_val.const_value = ir.LazyTensor(
+                _make_packed, dtype=packed_dtype, shape=out_shape
+            )
 
             model.graph.initializers[new_val.name] = new_val
 
diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py
index 492538d6..af53b362 100644
--- a/src/mobius/_passes/_fold_concat_test.py
+++ b/src/mobius/_passes/_fold_concat_test.py
@@ -467,6 +467,57 @@ def test_concat_with_no_inputs_skipped(self):
         result = FoldConcatInitializersPass()(model)
         assert not result.modified
 
+    def test_packed_dtype_follows_const_value_when_declared_dtype_missing(self):
+        """Regression: fp16 weights with an unset declared dtype must fold to fp16.
+
+        After ``_cast_module_dtype`` casts parameters to fp16, the resulting
+        initializer Values keep an fp16 ``const_value`` but lose their declared
+        ``.dtype`` (it becomes ``None``).  The pass must resolve the packed dtype
+        from ``const_value`` rather than defaulting to FLOAT, otherwise GQA's
+        packed-QKV weight is serialized as fp32 and ORT rejects the model with a
+        fp16/fp32 MatMul type-mismatch error.
+        """
+        q = ir.Value(name="q_weight")
+        q.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16))
+        k = ir.Value(name="k_weight")
+        k.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16))
+        v = ir.Value(name="v_weight")
+        v.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16))
+        # Critically: do NOT set .dtype — declared dtype stays None.
+        assert q.dtype is None and k.dtype is None and v.dtype is None
+
+        concat_node = ir.Node(
+            "",
+            "Concat",
+            inputs=[q, k, v],
+            attributes=[ir.Attr("axis", ir.AttributeType.INT, 0)],
+            num_outputs=1,
+        )
+        graph = ir.Graph(
+            inputs=[ir.Value(name="x")],
+            outputs=[concat_node.outputs[0]],
+            nodes=[concat_node],
+            name="test_graph",
+            opset_imports={"": 20},
+        )
+        graph.register_initializer(q)
+        graph.register_initializer(k)
+        graph.register_initializer(v)
+        model = ir.Model(graph, ir_version=10)
+
+        result = FoldConcatInitializersPass()(model)
+        assert result.modified
+
+        packed = model.graph.initializers["q_weight__k_weight__v_weight__axis_0__concat"]
+        assert packed.dtype == ir.DataType.FLOAT16, (
+            f"Packed initializer declared dtype should be FLOAT16, got {packed.dtype}"
+        )
+        assert packed.const_value.dtype == ir.DataType.FLOAT16, (
+            "Packed LazyTensor dtype should be FLOAT16, got "
+            f"{packed.const_value.dtype}"
+        )
+        assert packed.const_value.numpy().dtype == np.float16
+
     def test_shape_computed_from_inputs_when_output_shape_missing(self):
         """When the Concat output has no shape, shape is inferred from input shapes."""
         a = np.ones((4, 3), dtype=np.float32)
@@ -512,3 +563,82 @@ def test_shape_computed_from_inputs_when_output_shape_missing(self):
         # Shape should be computed from inputs: (4,3) + (4,3) along axis 0 → (8,3).
         packed = model.graph.initializers[packed_name]
         assert list(packed.shape) == [8, 3]
+
+
+class TestFoldConcatOrtLoad:
+    """End-to-end regression: a folded fp16 GQA-style graph must load in ORT.
+
+    Reproduces the original failure where the packed-QKV initializer was
+    serialized as fp32 while the rest of the graph was fp16, causing ORT to
+    reject the model with::
+
+        Type parameter (T) of Optype (MatMul) bound to different types
+        (tensor(float16) and tensor(float))
+
+    The failure occurs on the CPU EP too, so no GPU is needed.
+    """
+
+    def test_folded_fp16_concat_matmul_loads_in_ort(self, tmp_path):
+        import onnx
+        import onnxruntime as ort
+
+        # Three fp16 weights with UNSET declared dtype (as after _cast_module_dtype),
+        # but with shape retained so they serialize before they are pruned.
+        q = ir.Value(name="q_weight", shape=ir.Shape([8, 4]))
+        q.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16))
+        k = ir.Value(name="k_weight", shape=ir.Shape([8, 4]))
+        k.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16))
+        v = ir.Value(name="v_weight", shape=ir.Shape([8, 4]))
+        v.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16))
+
+        concat_node = ir.Node(
+            "",
+            "Concat",
+            inputs=[q, k, v],
+            attributes=[ir.Attr("axis", ir.AttributeType.INT, 1)],
+            num_outputs=1,
+        )
+        packed_out = concat_node.outputs[0]  # shape (8, 12), packed QKV weight
+        packed_out.shape = ir.Shape([8, 12])
+
+        # hidden @ packed_qkv : (N, 8) @ (8, 12) -> (N, 12), all fp16.
+        hidden = ir.Value(
+            name="hidden", shape=ir.Shape(["N", 8]), type=ir.TensorType(ir.DataType.FLOAT16)
+        )
+        matmul_node = ir.Node(
+            "", "MatMul", inputs=[hidden, packed_out], num_outputs=1
+        )
+        out = matmul_node.outputs[0]
+        out.shape = ir.Shape(["N", 12])
+        out.dtype = ir.DataType.FLOAT16
+
+        graph = ir.Graph(
+            inputs=[hidden],
+            outputs=[out],
+            nodes=[concat_node, matmul_node],
+            name="qkv_matmul",
+            opset_imports={"": 20},
+        )
+        graph.register_initializer(q)
+        graph.register_initializer(k)
+        graph.register_initializer(v)
+        model = ir.Model(graph, ir_version=10)
+
+        FoldConcatInitializersPass()(model)
+
+        # Drop the now-unused source initializers (a DCE step runs after folding
+        # in the real export pipeline). Without this they linger as dead weights.
+        for dead in ("q_weight", "k_weight", "v_weight"):
+            del model.graph.initializers[dead]
+
+        model_path = tmp_path / "fp16_qkv.onnx"
+        onnx.save(ir.to_proto(model), str(model_path))
+
+        # Before the fix this raises a fp16/fp32 MatMul type-mismatch at load.
+        sess = ort.InferenceSession(
+            str(model_path), providers=["CPUExecutionProvider"]
+        )
+        feed = {"hidden": np.random.randn(2, 8).astype(np.float16)}
+        (result,) = sess.run(None, feed)
+        assert result.shape == (2, 12)
+        assert result.dtype == np.float16
diff --git a/src/mobius/_passes/_fold_transpose.py b/src/mobius/_passes/_fold_transpose.py
index 402cad7c..e7f6cae3 100644
--- a/src/mobius/_passes/_fold_transpose.py
+++ b/src/mobius/_passes/_fold_transpose.py
@@ -26,6 +26,8 @@
 
 import onnx_ir as ir
 
+from mobius._passes._dtype_utils import initializer_dtype
+
 logger = logging.getLogger(__name__)
 
 
@@ -107,6 +109,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
 
                 new_name = f"{inp.name}_t"
 
+                # Resolve the dtype from the declared type or, when it was
+                # dropped during graph building, from the loaded const_value, so
+                # a transposed fp16 weight is not silently emitted as fp32.
+                transposed_dtype = initializer_dtype(inp) or ir.DataType.FLOAT
+
                 # Derive shape of the transposed tensor from the Transpose output.
                 # Fall back to computing from the input shape if shape inference
                 # did not propagate to this new node (e.g. after stage-2 rewrites).
@@ -115,7 +122,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                     perm = list(perm_attr.value)
                     t_shape = ir.Shape([inp.shape[p] for p in perm])
 
-                new_val = ir.Value(name=new_name, shape=t_shape, type=inp.type)
+                new_val = ir.Value(
+                    name=new_name,
+                    shape=t_shape,
+                    type=ir.TensorType(transposed_dtype),
+                )
 
                 # Create a LazyTensor that transposes the original data on demand.
                 # The actual numpy transposition is deferred until serialization,
@@ -123,7 +134,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                 src = inp  # captured for the closure below
                 new_val.const_value = ir.LazyTensor(
                     lambda s=src: ir.tensor(s.const_value.numpy().T),  # type: ignore[union-attr]
-                    dtype=inp.dtype or ir.DataType.FLOAT,
+                    dtype=transposed_dtype,
                     shape=t_shape,
                     name=new_val.name,
                 )
diff --git a/src/mobius/_passes/_fold_transpose_test.py b/src/mobius/_passes/_fold_transpose_test.py
index 93c74b37..0e3def2a 100644
--- a/src/mobius/_passes/_fold_transpose_test.py
+++ b/src/mobius/_passes/_fold_transpose_test.py
@@ -508,3 +508,53 @@ def test_original_weight_pruned_after_full_pipeline():
     assert "weight_t" not in model.graph.initializers
     transpose_nodes = [n for n in model.graph.all_nodes() if n.op_type == "Transpose"]
     assert len(transpose_nodes) == 1
+
+
+class TestFoldTransposeDtype:
+    """Regression: a transposed fp16 weight with unset declared dtype stays fp16."""
+
+    def test_transposed_dtype_follows_const_value_when_declared_dtype_missing(self):
+        """fp16 weights lose their declared dtype after _cast_module_dtype.
+
+        The pass must resolve the transposed initializer's dtype from
+        ``const_value`` rather than defaulting to FLOAT, otherwise a fp16 weight
+        is emitted as fp32 and ORT rejects the downstream MatMul.
+        """
+        weight = ir.Value(name="weight")
+        weight.const_value = ir.tensor(np.arange(12, dtype=np.float16).reshape(4, 3))
+        # Do NOT set .dtype — declared dtype stays None.
+        assert weight.dtype is None
+
+        x = ir.Value(name="x")
+        x.shape = ir.Shape([2, 4])
+        x.dtype = ir.DataType.FLOAT16
+
+        transpose_node = ir.Node(
+            "",
+            "Transpose",
+            inputs=[weight],
+            attributes=[ir.Attr("perm", ir.AttributeType.INTS, [1, 0])],
+            num_outputs=1,
+        )
+        w_t = transpose_node.outputs[0]
+        matmul_node = ir.Node("", "MatMul", inputs=[x, w_t], num_outputs=1)
+
+        graph = ir.Graph(
+            inputs=[x],
+            outputs=[matmul_node.outputs[0]],
+            nodes=[transpose_node, matmul_node],
+            name="test_graph",
+            opset_imports={"": 20},
+        )
+        graph.register_initializer(weight)
+        model = ir.Model(graph, ir_version=10)
+
+        result = FoldTransposedInitializerPass()(model)
+        assert result.modified
+
+        packed = model.graph.initializers["weight_t"]
+        assert packed.dtype == ir.DataType.FLOAT16, (
+            f"Transposed initializer declared dtype should be FLOAT16, got {packed.dtype}"
+        )
+        assert packed.const_value.dtype == ir.DataType.FLOAT16
+        assert packed.const_value.numpy().dtype == np.float16

From dd6bc15825386431207166feff199522f2bca1db Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:49:08 +0000
Subject: [PATCH 06/38] test(static-cache): add fp16 runtime maskless-decode
 (Flash-eligibility) guard

Code review asked for an e2e guard that the decode phase stays Flash-eligible.
ORT 1.27's Python profiler emits only op-level Node events (no CUDA Kernel
events), so the internal Flash-vs-MEA kernel choice is not observable from
end_profiling(). Instead assert the deterministic structural precondition that
governs Flash eligibility: profile an fp16 decode + prefill on CUDA and assert
the executed decode-branch Attention carries NO attn_mask input (rank-4) while
prefill does, per layer. This catches a regression that rewires the mask onto
the hot decode path and silently forces it onto the memory-efficient kernel.

Builder now parametrized by dtype + profiling; existing fp32 runnability test
unchanged.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/static_cache_decode_test.py | 148 ++++++++++++++++++++++++++++--
 1 file changed, 140 insertions(+), 8 deletions(-)

diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index 08effb20..2787e72b 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -21,16 +21,23 @@
 back.  It requires the CUDA Execution Provider because ``TensorScatter`` and
 the external-cache ``Attention`` path are CUDA-only.
 
-The model is built fp32: the ``is_causal`` guard fires in ORT *before* kernel
-dtype dispatch, so fp32 exercises the same external-cache code path as a
-production fp16 export while avoiding the cos/sin-cache dtype casting that
-only the full CLI build pipeline applies.  A raw ``InferenceSession`` is used
-(instead of the ``OnnxModelSession`` test helper) to keep this regression
-guard free of the optional ``onnxruntime-easy`` dependency.
+The runnability test (:func:`test_static_cache_prefill_and_decode_run_on_cuda`)
+is built fp32: the ``is_causal`` guard fires in ORT *before* kernel dtype
+dispatch, so fp32 exercises the same external-cache code path as a production
+fp16 export while avoiding the cos/sin-cache dtype casting that only the full
+CLI build pipeline applies.  The Flash-eligibility guard
+(:func:`test_static_cache_decode_runs_maskless_on_cuda`) builds fp16 — the
+production precision — and asserts the decode ``If`` branch executes without an
+``attn_mask`` input (the structural precondition for ORT to keep decode on
+Flash).  A raw ``InferenceSession`` is used (instead of the
+``OnnxModelSession`` test helper) to keep this regression guard free of the
+optional ``onnxruntime-easy`` dependency.
 """
 
 from __future__ import annotations
 
+import dataclasses
+import json
 import tempfile
 from pathlib import Path
 
@@ -76,6 +83,9 @@ def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None:
 
 def _build_static_cache_session(
     tmp_dir: str,
+    *,
+    ir_dtype: ir.DataType = ir.DataType.FLOAT,
+    enable_profiling: bool = False,
 ) -> tuple[ort.InferenceSession, object]:
     """Build a tiny static-cache qwen2 graph and load it on CUDA.
 
@@ -83,8 +93,19 @@ def _build_static_cache_session(
     so the phase-split ``If`` subgraphs are exercised through the real
     ``optimize_model`` pipeline — that is where a structural regression in
     the static-cache attention would surface.
+
+    Args:
+        tmp_dir: Directory for the saved ONNX model + external data.
+        ir_dtype: Weight/activation precision.  Defaults to fp32 (the
+            ``is_causal`` guard fires before kernel dtype dispatch, so fp32
+            exercises the same external-cache path while avoiding cos/sin
+            cache casting).  Pass ``ir.DataType.FLOAT16`` to build the
+            production-precision graph used for the Flash-eligibility guard.
+        enable_profiling: Turn on ORT op-level profiling so callers can
+            inspect which ``If`` branch executed and with which inputs.
     """
     config = _base_config()
+    config = dataclasses.replace(config, dtype=ir_dtype)
     module = registry.get(_MODEL_TYPE)(config)
     task = CausalLMTask(static_cache=True, max_seq_len=_MAX_SEQ_LEN)
     package = build_from_module(
@@ -95,8 +116,13 @@ def _build_static_cache_session(
 
     model_path = str(Path(tmp_dir) / "model.onnx")
     ir.save(model, model_path, external_data="model.onnx.data")
+    session_options = ort.SessionOptions()
+    if enable_profiling:
+        session_options.enable_profiling = True
+        session_options.profile_file_prefix = str(Path(tmp_dir) / "prof")
     session = ort.InferenceSession(
         model_path,
+        session_options,
         providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
     )
     assert "CUDAExecutionProvider" in session.get_providers(), (
@@ -105,11 +131,13 @@ def _build_static_cache_session(
     return session, config
 
 
-def _empty_caches(num_layers: int, kv_hidden: int) -> dict[str, np.ndarray]:
+def _empty_caches(
+    num_layers: int, kv_hidden: int, np_dtype: np.dtype = _CACHE_DTYPE
+) -> dict[str, np.ndarray]:
     """Zeroed ``[1, max_seq, kv_hidden]`` cache buffers for every layer."""
     feeds: dict[str, np.ndarray] = {}
     for layer in range(num_layers):
-        zeros = np.zeros((1, _MAX_SEQ_LEN, kv_hidden), dtype=_CACHE_DTYPE)
+        zeros = np.zeros((1, _MAX_SEQ_LEN, kv_hidden), dtype=np_dtype)
         feeds[f"key_cache.{layer}"] = zeros.copy()
         feeds[f"value_cache.{layer}"] = zeros.copy()
     return feeds
@@ -184,3 +212,107 @@ def test_static_cache_prefill_and_decode_run_on_cuda():
         assert np.any(advanced_key[0, prefill_len] != 0), (
             "decode should scatter the new key into cache slot N"
         )
+
+
+def _executed_attention_events(profile_path: str) -> list[dict]:
+    """Op-level ``Attention`` events from an ORT profiling JSON.
+
+    Each returned dict has ``name`` (carries the ``static_cache_prefill`` /
+    ``static_cache_decode`` branch tag) and ``has_mask`` — True when the
+    executed node received a rank-4 ``attn_mask`` input.  ORT 1.27's Python
+    profiler emits only op-level ``Node`` events (no CUDA ``Kernel`` events),
+    so the *internal* attention kernel (Flash vs memory-efficient) is not
+    observable here; the mask-input signature is, and the presence of an
+    ``attn_mask`` is exactly what makes ORT ineligible for Flash.
+    """
+    with open(profile_path) as handle:
+        events = json.load(handle)
+
+    attention_events: list[dict] = []
+    for event in events:
+        args = event.get("args", {})
+        if args.get("op_name") != "Attention":
+            continue
+        input_shapes = args.get("input_type_shape", [])
+        has_mask = any(
+            len(next(iter(shape.values()))) == 4 for shape in input_shapes
+        )
+        attention_events.append({"name": event["name"], "has_mask": has_mask})
+    return attention_events
+
+
+def test_static_cache_decode_runs_maskless_on_cuda():
+    """The decode branch executes maskless at runtime (Flash-eligible).
+
+    Build-time tests assert the *graph* phase-splits the static-cache
+    attention behind an ``If`` (decode branch omits ``attn_mask``).  This
+    test closes the loop at runtime on the production fp16 path: it profiles
+    a single-token decode and a multi-token prefill on CUDA and asserts the
+    ``If`` routed correctly and that the executed decode ``Attention`` carries
+    **no** ``attn_mask`` input.
+
+    Why this matters: ORT disables Flash whenever ``attn_mask`` is present
+    (by pointer, not content), so a regression that wired the mask onto the
+    decode branch would silently push the hot decode path onto the slower
+    memory-efficient kernel and invalidate the Attention-vs-GQA decode
+    comparison.  ORT 1.27's Python profiler does not surface the internal
+    kernel name, so we assert the structural precondition (mask absence)
+    that deterministically governs Flash eligibility.
+    """
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        session, config = _build_static_cache_session(
+            tmp_dir, ir_dtype=ir.DataType.FLOAT16, enable_profiling=True
+        )
+        num_layers = config.num_hidden_layers
+        kv_hidden = config.num_key_value_heads * config.head_dim
+        vocab = config.vocab_size
+        output_names = [out.name for out in session.get_outputs()]
+        rng = np.random.default_rng(2)
+
+        # Single-token decode (S_q = 1): the If must take the maskless branch.
+        decode_feeds: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64),
+            "position_ids": np.array([[4]], dtype=np.int64),
+            "write_indices": np.array([4], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([5], dtype=np.int64),
+        }
+        decode_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16))
+        session.run(output_names, decode_feeds)
+
+        # Multi-token prefill (S_q = 4): the If must take the masked branch.
+        prefill_len = 4
+        prefill_feeds: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(
+                0, vocab, size=(1, prefill_len), dtype=np.int64
+            ),
+            "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
+            "write_indices": np.array([0], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
+        }
+        prefill_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16))
+        session.run(output_names, prefill_feeds)
+
+        events = _executed_attention_events(session.end_profiling())
+
+    decode_events = [e for e in events if "static_cache_decode" in e["name"]]
+    prefill_events = [e for e in events if "static_cache_prefill" in e["name"]]
+
+    # Decode took the maskless branch on every layer (Flash-eligible).
+    assert len(decode_events) == num_layers, (
+        f"expected {num_layers} decode-branch Attention executions, "
+        f"got {len(decode_events)}"
+    )
+    assert all(not e["has_mask"] for e in decode_events), (
+        "decode-branch Attention must run WITHOUT an attn_mask input so ORT "
+        "keeps it on Flash; a mask here forces the slower memory-efficient "
+        "path and breaks the decode-latency comparison"
+    )
+
+    # Prefill took the masked branch on every layer (memory-efficient path).
+    assert len(prefill_events) == num_layers, (
+        f"expected {num_layers} prefill-branch Attention executions, "
+        f"got {len(prefill_events)}"
+    )
+    assert all(e["has_mask"] for e in prefill_events), (
+        "prefill-branch Attention must carry the explicit causal mask"
+    )

From 0859a357478a148e7afdc1cdc07db4a774ecd2e7 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:51:20 +0000
Subject: [PATCH 07/38] test(static-cache): make phase-split mask guard
 fail-closed per-If

Per hardened Code+Critical review spec: instead of counting masked/maskless
Attention nodes globally (which could pass vacuously), locate the If nodes
directly and inspect BOTH branch subgraphs. Assert (a) one If per layer exists
-- a regression to a single unconditional-mask Attention (no If) now fails the
If-count assertion -- and (b) then-branch (prefill) Attention carries the
GreaterOrEqual causal mask while else-branch (decode) omits attn_mask
(Flash-eligibility precondition). Catches missing-If, single-sided, and
inverted-mask regressions.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/build_graph_test.py | 87 +++++++++++++++++++++++++--------------
 1 file changed, 56 insertions(+), 31 deletions(-)

diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index 81dfe45b..3ce03ce4 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -4490,42 +4490,67 @@ def test_static_cache_attention_is_causal(self):
             )
 
     def test_static_cache_phase_split_mask_presence(self):
-        """Prefill branch carries an explicit causal mask; decode omits it.
-
-        This is the Flash-eligibility invariant: ORT disables Flash whenever
-        attn_mask is present (by pointer, not content), so single-token
-        decode MUST structurally omit the mask to stay on Flash/XQA, while
-        multi-token prefill supplies a causal mask (memory-efficient path).
+        """Each layer is an If whose two branches differ in mask presence.
+
+        Fail-closed structural guard for the phase split.  It locates the
+        ``If`` nodes directly and inspects *both* subgraphs, so it cannot pass
+        vacuously: a regression to a single unconditional Attention (no If)
+        fails the If-count assertion, and an inverted or single-sided mask
+        fails the per-branch checks.
+
+        Invariant (If condition is ``Greater(Shape(query)[1], 1)``):
+        * ``then_branch`` (multi-token / prefill) Attention HAS an explicit
+          causal ``attn_mask`` produced by ``GreaterOrEqual`` → memory-
+          efficient path.
+        * ``else_branch`` (single-token / decode) Attention OMITS ``attn_mask``
+          (input[3] None/empty) → Flash-eligible (ORT disables Flash whenever
+          attn_mask is present, by pointer not content).
         """
         model, config = self._build_static_cache_model()
 
-        attention_nodes = [
-            n for n in self._walk_nodes(model.graph) if n.op_type == "Attention"
-        ]
-        with_mask = []
-        without_mask = []
-        for node in attention_nodes:
-            attn_mask_input = node.inputs[3] if len(node.inputs) > 3 else None
-            if attn_mask_input is not None and attn_mask_input.name != "":
-                producer = attn_mask_input.producer()
-                assert producer is not None and producer.op_type == "GreaterOrEqual", (
-                    f"Masked Attention {node.name} attn_mask should come from "
-                    f"the causal-mask GreaterOrEqual, got "
-                    f"{None if producer is None else producer.op_type}"
-                )
-                with_mask.append(node)
-            else:
-                without_mask.append(node)
-
-        # Exactly one masked (prefill) and one maskless (decode) per layer.
-        assert len(with_mask) == config.num_hidden_layers, (
-            f"Expected {config.num_hidden_layers} masked (prefill) Attention "
-            f"ops, got {len(with_mask)}"
+        if_nodes = [n for n in self._walk_nodes(model.graph) if n.op_type == "If"]
+        # (a) Fail-closed: the phase-split If must exist, one per layer.
+        assert len(if_nodes) == config.num_hidden_layers, (
+            f"static-cache attention must phase-split via If "
+            f"(expected {config.num_hidden_layers}, got {len(if_nodes)}); a "
+            f"single unconditional-mask Attention would force decode off Flash"
         )
-        assert len(without_mask) == config.num_hidden_layers, (
-            f"Expected {config.num_hidden_layers} maskless (decode) Attention "
-            f"ops for Flash eligibility, got {len(without_mask)}"
+
+        for if_node in if_nodes:
+            then_branch = if_node.attributes["then_branch"].value
+            else_branch = if_node.attributes["else_branch"].value
+            then_attn = self._single_attention(then_branch)
+            else_attn = self._single_attention(else_branch)
+
+            # (b) then-branch (prefill) MUST carry the GreaterOrEqual mask.
+            then_mask = then_attn.inputs[3] if len(then_attn.inputs) > 3 else None
+            assert then_mask is not None and then_mask.name != "", (
+                "prefill (then) branch Attention must carry an attn_mask"
+            )
+            producer = then_mask.producer()
+            assert producer is not None and producer.op_type == "GreaterOrEqual", (
+                f"prefill attn_mask should come from the causal-mask "
+                f"GreaterOrEqual, got "
+                f"{None if producer is None else producer.op_type}"
+            )
+
+            # (b) else-branch (decode) MUST omit the mask to stay Flash-eligible.
+            else_mask = else_attn.inputs[3] if len(else_attn.inputs) > 3 else None
+            assert else_mask is None or else_mask.name == "", (
+                "decode (else) branch Attention must OMIT attn_mask so ORT "
+                "keeps it on Flash; a mask here forces the slower memory-"
+                "efficient path"
+            )
+
+    @staticmethod
+    def _single_attention(graph):
+        """Return the sole ``Attention`` node in an If branch subgraph."""
+        attention_nodes = [n for n in graph if n.op_type == "Attention"]
+        assert len(attention_nodes) == 1, (
+            f"expected exactly one Attention per If branch, "
+            f"got {len(attention_nodes)}"
         )
+        return attention_nodes[0]
 
     def test_static_cache_is_phase_split_behind_if(self):
         """Static cache attention must be a per-layer If over masked/maskless.

From a28b4e1f452e3a45a7c6886571f0954754355500 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:54:07 +0000
Subject: [PATCH 08/38] docs: update obsolete fp16 GQA export bug docs after
 df203cc fix

Section 3 of the mobius-onnx-export-gotchas skill is no longer a live bug.
df203cc fixed fp16 GQA exports emitting fp32 packed-QKV/transposed weights
(fold passes now stamp the correct dtype via _dtype_utils.initializer_dtype).
Rewrote section 3 from a BUG+manual-post-cast workaround to a FIXED note with
corrected root cause (fold passes, not the PackQKV rewrite), kept the FLOAT32
verification snippet, and demoted the post-cast script to a fallback for stale
pre-fix artifacts. Updated the skill frontmatter and added a CHANGELOG entry.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../mobius-onnx-export-gotchas/SKILL.md       | 52 ++++++++++++-------
 CHANGELOG.md                                  | 18 +++++++
 2 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
index 65f3b640..d3442f15 100644
--- a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -1,6 +1,6 @@
 ---
 name: mobius-onnx-export-gotchas
-description: Use when building/exporting ONNX models with the `mobius build` CLI (especially Phi-3 / Phi-3.5 or any model with `--execution-provider cuda` GQA fusion and/or `--static-cache`). Covers the current CLI syntax, the dtype flag values, the GQA-vs-static-cache interaction, and a known fp16 packed-QKV FLOAT32 bug that makes GQA fp16 exports fail to load in onnxruntime.
+description: Use when building/exporting ONNX models with the `mobius build` CLI (especially Phi-3 / Phi-3.5 or any model with `--execution-provider cuda` GQA fusion and/or `--static-cache`). Covers the current CLI syntax, the dtype flag values, the GQA-vs-static-cache interaction, and how to verify fp16 GQA exports load in onnxruntime (the historical packed-QKV FLOAT32 load bug is fixed as of df203cc).
 ---
 
 # mobius ONNX export gotchas
@@ -31,24 +31,40 @@ op). That breaks the pattern the GQA rewrite matches, so combining
   OrtValue — NOT via `--static-cache`.
 - **ONNX-Attention + in-place cache:** `--execution-provider default --static-cache --max-seq-len N`.
 
-## 3. BUG: GQA fp16 export leaves packed-QKV weights as FLOAT32 → model won't load
-For an fp16 GQA export, the per-layer packed QKV weight
-(`..q_proj.weight__k_proj.weight__v_proj.weight__axis_0__concat`) is emitted as **FLOAT32**, while its
-MatMul's other input is fp16. onnxruntime then rejects the model at load:
+## 3. FIXED: fp16 GQA export previously left packed-QKV weights as FLOAT32 → model wouldn't load
+**Status: fixed as of commit `df203cc`.** Native fp16 Phi-3.5 GQA export now loads directly in the ORT
+CUDA EP with **no manual post-cast** (32 GroupQueryAttention nodes, all-fp16 initializers). If you are on
+that commit or later, you should not hit this — skip to the verification snippet below. The history is
+kept here because old artifacts exported before the fix still carry fp32 packed weights.
+
+### Symptom (pre-fix)
+For an fp16 GQA export, a folded per-layer packed QKV weight
+(`..q_proj.weight__k_proj.weight__v_proj.weight__axis_0__concat`) was emitted as **FLOAT32**, while its
+MatMul's other input was fp16. onnxruntime then rejected the model at load on both CPU and CUDA EPs:
 
 ```
 Type Error: Type parameter (T) of Optype (MatMul) bound to different types
 (tensor(float16) and tensor(float)) in node (node_MatMul_*)
 ```
 
-You'll also see at save time: `The value type for shape [H, 3H] is not known. Skipping serialization`.
+You'd also see at save time: `The value type for shape [H, 3H] is not known. Skipping serialization`.
+
+### Root cause
+`_cast_module_dtype` casts module params to fp16, but the resulting initializer `Value`s lose their
+declared `.dtype` (it becomes `None`) while their `const_value` stays fp16. The fold passes
+`FoldConcatInitializersPass` (`src/mobius/_passes/_fold_concat.py`) and `FoldTransposedInitializerPass`
+(`src/mobius/_passes/_fold_transpose.py`) then defaulted the folded initializer's dtype to `FLOAT`,
+serializing the packed QKV / transposed weights as fp32.
 
-**Root cause:** `_cast_module_dtype` (`src/mobius/_builder.py:84`) casts module params to fp16 *before*
-graph build. The GQA `PackQKVWithBias` rewrite (`src/mobius/rewrite_rules/_group_query_attention.py`)
-then emits the packed weight as a graph-level `op.Concat(q_w,k_w,v_w)` that a constant-fold collapses
-into a NEW initializer whose dtype is FLOAT32/untyped — the fp16 cast never reaches it.
+### The fix
+A shared helper `initializer_dtype()` (`src/mobius/_passes/_dtype_utils.py`) resolves the effective dtype
+from the declared type, **falling back to `const_value` when the type annotation was dropped** (preferring
+the data dtype and warning on stale-metadata disagreement). Both fold passes use it to stamp the correct
+dtype on the new initializer's `TensorType` and `LazyTensor`, and `FoldConcatInitializersPass` now also
+skips folding before weights are loaded (mirroring `FoldTransposedInitializerPass`). A regression test
+loads the fp16 GQA export in the ORT CPU EP to lock this in.
 
-### Detect
+### Verify (still worth running on any fp16 build)
 ```python
 import onnx
 m = onnx.load("model.onnx", load_external_data=False)
@@ -56,11 +72,11 @@ fp32 = [i.name for i in m.graph.initializer if i.data_type == onnx.TensorProto.F
 print(len(fp32), "FLOAT32 initializers (should be 0 for fp16)")
 ```
 
-### Fix (post-export, numerically == intended fp16)
-Cast the FLOAT32 initializers to fp16, optionally strip dead pre-pack q/k/v initializers, re-save.
-**Gotcha when re-saving with external data:** if you save with `location="X.data"` and then rename the
-file, the references inside `model.onnx` still point to `X.data`. Either save directly with
-`location="model.onnx.data"`, or rewrite each initializer's `external_data` `location` entry.
+### Salvaging a stale pre-fix artifact (only if re-exporting is not an option)
+Prefer re-exporting on the fixed code. If you must repair an old model, cast its FLOAT32 initializers to
+fp16 and re-save. **Gotcha when re-saving with external data:** if you save with `location="X.data"` and
+then rename the file, the references inside `model.onnx` still point to `X.data`. Either save directly
+with `location="model.onnx.data"`, or rewrite each initializer's `external_data` `location` entry.
 
 ```python
 import onnx, numpy as np
@@ -74,10 +90,6 @@ onnx.save(m, "model.onnx", save_as_external_data=True, all_tensors_to_one_file=T
           location="model.onnx.data", size_threshold=1024, convert_attribute=False)
 ```
 
-**Proper upstream fix:** set the packed-Concat output type to the model dtype in the GQA pack rewrite,
-or cast ALL float initializers at save time regardless of registered value_info; add an e2e test that
-loads the fp16 GQA export in onnxruntime.
-
 ## 4. Always validate the export in ORT before profiling
 Load the model on `CUDAExecutionProvider` and run one prefill + one decode `session.run`. Confirm:
 (a) the expected attention op (`com.microsoft::GroupQueryAttention` vs `ai.onnx::Attention`),
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 83edfb24..dee0e635 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### fp16 GQA Export Fix
+
+#### Fixed
+
+- Native fp16 GroupQueryAttention exports (e.g. `microsoft/Phi-3.5-mini-instruct`
+  with `--dtype f16 --execution-provider cuda`) no longer emit fp32 packed-QKV /
+  transposed weights. Previously the fold passes (`FoldConcatInitializersPass`,
+  `FoldTransposedInitializerPass`) defaulted a folded initializer's dtype to
+  `FLOAT` when the source `Value`'s declared type had been dropped during fp16
+  casting, producing a model onnxruntime rejected at load with a
+  `MatMul` type-parameter error (`tensor(float16)` vs `tensor(float)`) on both
+  CPU and CUDA EPs. A new `mobius._passes._dtype_utils.initializer_dtype()`
+  helper now resolves the effective dtype from `const_value` when the type
+  annotation is missing, so fp16 GQA models load directly with no manual
+  post-cast.
+
+---
+
 ### WebGPU Shape Op Support
 
 #### Changed

From 71e84b371bc7cae567de9fec1d92acae3b39b760 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 22:55:45 +0000
Subject: [PATCH 09/38] Strip dead pre-pack weights in
 FoldConcatInitializersPass
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

FoldConcatInitializersPass removed the QKV-pack Concat node with
`graph.remove(node)`, which detaches the node from the graph's node list
but NOT from its input Values. The folded q/k/v source initializers kept
a stale use pointing at the removed Concat, so the downstream
RemoveUnusedNodesPass (run by fold_initializers_after_weights) treated
them as live and left them in the graph. For fp16 Phi-3.5 GQA that
serialized 96 orphaned pre-pack q/k/v_proj weights (~1.8 GB) into the
exported model.

Use `graph.remove(node, safe=True)` at both removal sites so the node
detaches from its inputs, clearing the source initializers' use lists.
The existing RemoveUnusedNodesPass then strips the dead pre-pack weights
as part of the proper export — no post-hoc patch needed.
FoldTransposedInitializerPass already does this; this aligns FoldConcat.

Add a regression test asserting the source initializers are detached
(zero uses) after folding and removed by RemoveUnusedNodesPass.

Verified end-to-end: native fp16 Phi-3.5 GQA export drops from 8.9 GB to
7.2 GB (199 initializers, 0 unused), all fp16, loads + runs on ORT CUDA EP.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_concat.py      | 12 +++++--
 src/mobius/_passes/_fold_concat_test.py | 45 +++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py
index 2e1e926d..7035f1e9 100644
--- a/src/mobius/_passes/_fold_concat.py
+++ b/src/mobius/_passes/_fold_concat.py
@@ -125,7 +125,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
             if packed_name in model.graph.initializers:
                 existing_val = model.graph.initializers[packed_name]
                 out_val.replace_all_uses_with(existing_val, replace_graph_outputs=True)
-                model.graph.remove(node)
+                # safe=True detaches the node from its inputs before removal so
+                # the folded q/k/v initializers' use lists are cleared. Otherwise
+                # `inp.uses()` still points at the removed Concat and the dead
+                # pre-pack weights survive RemoveUnusedNodesPass (~1.8 GB of
+                # orphaned weights serialized into the fp16 GQA model).
+                model.graph.remove(node, safe=True)
                 folded_count += 1
                 modified = True
 
@@ -166,7 +171,10 @@ def _make_packed(
             model.graph.initializers[new_val.name] = new_val
 
             out_val.replace_all_uses_with(new_val, replace_graph_outputs=True)
-            model.graph.remove(node)
+            # safe=True detaches the node from its inputs before removal so the
+            # folded q/k/v initializers' use lists are cleared; otherwise the
+            # dead pre-pack weights survive RemoveUnusedNodesPass.
+            model.graph.remove(node, safe=True)
             folded_count += 1
             modified = True
 
diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py
index af53b362..bff9f2ee 100644
--- a/src/mobius/_passes/_fold_concat_test.py
+++ b/src/mobius/_passes/_fold_concat_test.py
@@ -124,6 +124,51 @@ def test_dynamic_input_not_folded(self):
         result = FoldConcatInitializersPass()(model)
         assert not result.modified
 
+    def test_folded_inputs_detached_so_dce_strips_dead_weights(self):
+        """After folding, the source initializers must be detached and DCE-able.
+
+        FoldConcat removes the Concat node; if it does not detach the node from
+        its inputs (``safe=True``), the source q/k/v initializers keep a stale
+        use pointing at the removed node, so RemoveUnusedNodesPass cannot strip
+        them. For fp16 GQA that leaves ~1.8 GB of dead pre-pack weights in the
+        exported model.
+        """
+        from onnx_ir.passes import common as common_passes
+
+        a = np.ones((4, 8), dtype=np.float16)
+        b = np.ones((4, 8), dtype=np.float16) * 2
+        model, init_vals = _make_concat_model([a, b], axis=0)
+        # The Concat output is the only graph output; route it through a MatMul
+        # so the packed initializer stays live while the sources become dead.
+        packed_out = model.graph.outputs[0]
+        hidden = ir.Value(
+            name="hidden",
+            shape=ir.Shape(["N", 8]),
+            type=ir.TensorType(ir.DataType.FLOAT16),
+        )
+        matmul = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1)
+        mm_out = matmul.outputs[0]
+        mm_out.shape = ir.Shape(["N", 8])
+        mm_out.dtype = ir.DataType.FLOAT16
+        model.graph.append(matmul)
+        model.graph.inputs.append(hidden)
+        model.graph.outputs[0] = mm_out
+
+        FoldConcatInitializersPass()(model)
+
+        # Sources must be detached (zero uses) so DCE can remove them.
+        for v in init_vals:
+            assert len(list(v.uses())) == 0, (
+                f"{v.name} still has uses after fold — node not detached"
+            )
+
+        common_passes.RemoveUnusedNodesPass()(model)
+        remaining = set(model.graph.initializers)
+        assert "init_0" not in remaining and "init_1" not in remaining, (
+            f"Dead pre-pack weights survived DCE: {remaining}"
+        )
+        assert "init_0__init_1__axis_0__concat" in remaining
+
     def test_uses_lazy_tensor(self):
         """The packed initializer wraps sources in a LazyTensor.
 

From 80f09022cf803051340ae3b57dca8dcfa9477f1d Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:01:33 +0000
Subject: [PATCH 10/38] test(static-cache): assert decode selects Flash,
 prefill MEA (direct kernel proof)

The reviewers + architect required proof that the phase split actually keeps
decode on Flash at runtime (structural maskless-ness is necessary but not the
proof). ORT 1.27's Python profiler exposes no kernel name, but the opset-24 LLM
Attention kernel logs its choice at VERBOSE via the default logger
(attention.cc: "ONNX Attention: using Flash Attention" / "... Memory Efficient
Attention"). New fp16 CUDA test raises the default logger to VERBOSE, redirects
fd 2 around decode + prefill runs, and asserts decode (q_seq=1) selects Flash
on every layer while prefill (q_seq>1) selects Memory-Efficient. A mask wired
onto decode flips it to MEA and fails this test. Kept the structural maskless
test too (robust to log-format changes).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/static_cache_decode_test.py | 144 +++++++++++++++++++++++++++++-
 1 file changed, 142 insertions(+), 2 deletions(-)

diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index 2787e72b..575a0ebe 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -11,14 +11,27 @@
 cache is pre-allocated to ``max_seq_len``, that guard fires in **both**
 prefill (``S_q = N``) and decode (``S_q = 1``), raising ``NOT_IMPLEMENTED``.
 
+Three CUDA regression tests here guard against (a) the original
+``is_causal=1`` ``NOT_IMPLEMENTED`` regression and (b) a mask wired onto the
+decode branch that would silently push decode off Flash:
+
+* :func:`test_static_cache_prefill_and_decode_run_on_cuda` (fp32) — both
+  phases run without ``NOT_IMPLEMENTED``.
+* :func:`test_static_cache_decode_runs_maskless_on_cuda` (fp16) — the
+  executed decode ``Attention`` carries no ``attn_mask`` input (the
+  structural Flash-eligibility precondition), robust to log-format changes.
+* :func:`test_static_cache_decode_selects_flash_kernel_on_cuda` (fp16) — the
+  direct proof: ORT's VERBOSE kernel-selection log shows decode on **Flash**
+  and prefill on **Memory-Efficient**.
+
 The fix sets ``is_causal=0`` and phase-splits the attention behind an ``If``
 keyed on ``Shape(query)[1] > 1``: the multi-token (prefill) branch supplies an
 explicit causal mask (:func:`mobius.components._common.create_static_cache_causal_mask`,
 memory-efficient path), while the single-token decode branch omits the mask so
 ORT keeps it on Flash/XQA — the same kernel the GQA variant uses, so the
-profiling comparison stays apples-to-apples.  This test exercises the actual
+profiling comparison stays apples-to-apples.  These tests exercise the actual
 ONNX Runtime kernel for both phases so the regression cannot silently come
-back.  It requires the CUDA Execution Provider because ``TensorScatter`` and
+back.  They require the CUDA Execution Provider because ``TensorScatter`` and
 the external-cache ``Attention`` path are CUDA-only.
 
 The runnability test (:func:`test_static_cache_prefill_and_decode_run_on_cuda`)
@@ -36,8 +49,11 @@
 
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 import json
+import os
+import re
 import tempfile
 from pathlib import Path
 
@@ -316,3 +332,127 @@ def test_static_cache_decode_runs_maskless_on_cuda():
     assert all(e["has_mask"] for e in prefill_events), (
         "prefill-branch Attention must carry the explicit causal mask"
     )
+
+
+# Matches the opset-24 LLM Attention kernel-selection log line emitted at
+# VERBOSE by onnxruntime/core/providers/cuda/llm/attention.cc, e.g.
+#   "ONNX Attention: using Flash Attention (batch=1, q_seq=1, total_seq=16, ...)"
+#   "ONNX Attention: using Memory Efficient Attention (batch=1, q_seq=4, ...)"
+_ATTENTION_KERNEL_LINE = re.compile(
+    r"ONNX Attention: using (?P<kernel>.+?) \(batch=\d+, q_seq=(?P<q_seq>\d+)"
+)
+
+
+@contextlib.contextmanager
+def _capture_attention_kernel_log():
+    """Capture ORT's per-op attention kernel-selection log lines.
+
+    The opset-24 LLM ``Attention`` CUDA kernel logs which kernel it selected
+    (Flash / Memory-Efficient / unfused) at VERBOSE through the *default*
+    (process-global) logger, written to the C++ ``stderr`` (fd 2).  ORT's
+    Python profiler does not surface this, so to read it we raise the default
+    logger severity to VERBOSE and redirect fd 2 around the run.  Callers must
+    read the yielded file *inside* the ``with`` block (it is closed on exit);
+    fd 2 and the logger severity are restored before exit so assertions made
+    after the block still report normally.
+    """
+    ort.set_default_logger_severity(0)
+    saved_stderr_fd = os.dup(2)
+    try:
+        with tempfile.TemporaryFile(mode="w+b") as capture_file:
+            os.dup2(capture_file.fileno(), 2)
+            try:
+                yield capture_file
+            finally:
+                capture_file.flush()
+                os.dup2(saved_stderr_fd, 2)
+    finally:
+        os.close(saved_stderr_fd)
+        ort.set_default_logger_severity(2)  # back to ORT's default (WARNING)
+
+
+def _selected_attention_kernels(capture_file) -> list[tuple[str, int]]:
+    """Parse ``(kernel_name, q_seq)`` pairs from a captured verbose log."""
+    capture_file.seek(0)
+    text = capture_file.read().decode("utf-8", "replace")
+    kernels: list[tuple[str, int]] = []
+    for line in text.splitlines():
+        match = _ATTENTION_KERNEL_LINE.search(line)
+        if match is not None:
+            kernels.append((match.group("kernel").strip(), int(match.group("q_seq"))))
+    return kernels
+
+
+def test_static_cache_decode_selects_flash_kernel_on_cuda():
+    """Decode actually selects Flash; prefill selects Memory-Efficient.
+
+    Structural maskless-ness (the test above) is necessary, but the *proof*
+    that the phase split achieves its purpose is the kernel ORT actually runs.
+    The opset-24 LLM Attention kernel logs its choice at VERBOSE; this test
+    captures that log and asserts the single-token decode runs **Flash**
+    (the same external-cache kernel the GQA variant uses, so the decode-
+    latency comparison is apples-to-apples) while the multi-token prefill
+    runs **Memory-Efficient** (mask present, the cheap amortized path).
+
+    This is the regression guard the reviewers required: a change that wired
+    the causal mask onto the decode branch would flip its kernel from Flash
+    to Memory-Efficient and fail here, even though finiteness / scatter / If-
+    count assertions would all still pass.
+    """
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        session, config = _build_static_cache_session(
+            tmp_dir, ir_dtype=ir.DataType.FLOAT16
+        )
+        num_layers = config.num_hidden_layers
+        kv_hidden = config.num_key_value_heads * config.head_dim
+        vocab = config.vocab_size
+        output_names = [out.name for out in session.get_outputs()]
+        rng = np.random.default_rng(3)
+
+        decode_feeds: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64),
+            "position_ids": np.array([[4]], dtype=np.int64),
+            "write_indices": np.array([4], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([5], dtype=np.int64),
+        }
+        decode_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16))
+
+        prefill_len = 4
+        prefill_feeds: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(
+                0, vocab, size=(1, prefill_len), dtype=np.int64
+            ),
+            "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
+            "write_indices": np.array([0], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
+        }
+        prefill_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16))
+
+        with _capture_attention_kernel_log() as capture_file:
+            session.run(output_names, decode_feeds)
+            session.run(output_names, prefill_feeds)
+            kernels = _selected_attention_kernels(capture_file)
+
+    # The kernel log distinguishes phases by query length (q_seq).
+    decode_kernels = [name for name, q_seq in kernels if q_seq == 1]
+    prefill_kernels = [name for name, q_seq in kernels if q_seq == prefill_len]
+
+    assert len(decode_kernels) == num_layers, (
+        f"expected {num_layers} decode (q_seq=1) Attention kernel-selection "
+        f"log lines, got {len(decode_kernels)} (all parsed: {kernels}). If "
+        f"empty, ORT's verbose attention-kernel log format may have changed."
+    )
+    assert all("Flash" in name for name in decode_kernels), (
+        f"decode MUST select Flash Attention to stay apples-to-apples with "
+        f"GQA's decode kernel; got {decode_kernels}. A mask on the decode "
+        f"branch flips this to Memory-Efficient."
+    )
+
+    assert len(prefill_kernels) == num_layers, (
+        f"expected {num_layers} prefill (q_seq={prefill_len}) kernel-selection "
+        f"log lines, got {len(prefill_kernels)} (all parsed: {kernels})"
+    )
+    assert all("Memory Efficient" in name for name in prefill_kernels), (
+        f"prefill should select Memory-Efficient Attention (causal mask "
+        f"present); got {prefill_kernels}"
+    )

From b39d99a67039f47564137d7a9ad9492632efc5da Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:04:41 +0000
Subject: [PATCH 11/38] docs(static-cache): document load-bearing
 rename/pin/append ordering

The phase-split branch builder relies on a specific ordering of
rename_subgraph_values -> pin attn_output.name -> append to branch.outputs.
Per architect review, reordering these would either rename the pinned
output or expose it to renaming before it is protected. Make the
invariant explicit so future edits don't silently break SSA / If wiring.
Comment-only; no behavior change (11/11 static-cache graph tests pass).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/components/_attention.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py
index 93b57dd9..801b8155 100644
--- a/src/mobius/components/_attention.py
+++ b/src/mobius/components/_attention.py
@@ -101,6 +101,22 @@ def _build_attention_branch(name: str, use_causal_mask: bool) -> ir.Graph:
         # Prefix internal node/value names so the two branches stay in SSA
         # form when merged under the parent graph, then pin the branch
         # output name (the If wires branches by output position).
+        #
+        # Ordering is load-bearing — do NOT reorder these three lines:
+        #   1. rename first: rename_subgraph_values renames node OUTPUT value
+        #      names but skips graph.inputs/outputs. Running it before the
+        #      output is appended (branch.outputs is still empty) means
+        #      attn_output is renamed here as an internal node output, not
+        #      protected as a graph output.
+        #   2. pin the name next: this deterministic name must be set AFTER
+        #      the rename so the rename does not clobber it, and BEFORE the
+        #      append so the If sees a stable output name.
+        #   3. append last: registers the now-stable value as the branch
+        #      output the parent If wires by position.
+        # Outer-scope captures (query/key_cache/value_cache/nonpad_kv_seqlen/
+        # write_indices) are only ever node INPUTS inside the branch, never
+        # outputs, so the rename never touches them and implicit-input capture
+        # stays intact.
         rename_subgraph_values(branch, f"{name}_")
         attn_output.name = f"{name}_attn_output"
         branch.outputs.append(attn_output)

From b094af37ed4d9c696c9ed40d22ef36149a7b29ac Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:08:04 +0000
Subject: [PATCH 12/38] test(static-cache): assert decode honors nonpad bound
 (cache-poison guard)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds a semantic regression guard closing code-review MINOR-3: the existing
e2e tests prove decode runs, stays maskless, and selects Flash, but not that
the nonpad_kv_seqlen bound is actually applied. This test poisons every cache
slot at/beyond nonpad with large garbage and asserts decode logits are
bit-identical to the clean-cache decode — proving out-of-range keys are never
attended (fp16, CUDA).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/static_cache_decode_test.py | 76 +++++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)

diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index 575a0ebe..8374cce0 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -456,3 +456,79 @@ def test_static_cache_decode_selects_flash_kernel_on_cuda():
         f"prefill should select Memory-Efficient Attention (causal mask "
         f"present); got {prefill_kernels}"
     )
+
+
+def test_static_cache_decode_ignores_keys_beyond_nonpad_on_cuda():
+    """Out-of-bound cache slots cannot change a decode's output.
+
+    The decode branch runs maskless and relies solely on ``nonpad_kv_seqlen``
+    to bound attention to the valid keys.  The other e2e tests prove decode
+    *runs* and stays on Flash, but not that the bound is semantically applied.
+    This test closes that gap: it compares a decode over a clean carried cache
+    against a decode over the *same* cache with every slot at or beyond
+    ``nonpad`` overwritten with large garbage.  If the bound is honored the two
+    decodes produce bit-identical logits; if a regression let the kernel read
+    the whole pre-allocated cache, the garbage would perturb the softmax and
+    the logits would diverge.
+    """
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        session, config = _build_static_cache_session(
+            tmp_dir, ir_dtype=ir.DataType.FLOAT16
+        )
+        num_layers = config.num_hidden_layers
+        kv_hidden = config.num_key_value_heads * config.head_dim
+        vocab = config.vocab_size
+        output_names = [out.name for out in session.get_outputs()]
+        rng = np.random.default_rng(4)
+
+        # Prefill four real tokens into slots 0..3 to populate the cache.
+        prefill_len = 4
+        prefill_feeds: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(
+                0, vocab, size=(1, prefill_len), dtype=np.int64
+            ),
+            "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
+            "write_indices": np.array([0], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
+        }
+        prefill_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16))
+        prefill_out = dict(
+            zip(output_names, session.run(output_names, prefill_feeds))
+        )
+
+        # Decode one token into slot 4; valid keys are slots 0..4 (nonpad=5).
+        nonpad = prefill_len + 1
+        decode_inputs: dict[str, np.ndarray] = {
+            "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64),
+            "position_ids": np.array([[prefill_len]], dtype=np.int64),
+            "write_indices": np.array([prefill_len], dtype=np.int64),
+            "nonpad_kv_seqlen": np.array([nonpad], dtype=np.int64),
+        }
+
+        clean_feeds = {
+            **decode_inputs,
+            **_carry_caches(prefill_out, num_layers),
+        }
+        baseline = dict(
+            zip(output_names, session.run(output_names, clean_feeds))
+        )
+
+        # Poison every cache slot at or beyond ``nonpad`` with large garbage;
+        # those positions must never be attended during decode.  Slot 4 (the
+        # decode write target, within nonpad) is left untouched.
+        poisoned_caches = _carry_caches(prefill_out, num_layers)
+        for layer in range(num_layers):
+            for name in (f"key_cache.{layer}", f"value_cache.{layer}"):
+                buf = poisoned_caches[name].copy()
+                buf[:, nonpad:, :] = np.float16(50.0)
+                poisoned_caches[name] = buf
+        poisoned_feeds = {**decode_inputs, **poisoned_caches}
+        perturbed = dict(
+            zip(output_names, session.run(output_names, poisoned_feeds))
+        )
+
+    assert np.array_equal(baseline["logits"], perturbed["logits"]), (
+        "decode logits changed when cache slots beyond nonpad_kv_seqlen were "
+        "poisoned — the nonpad bound is not being honored, so decode is "
+        "attending to invalid (out-of-range) keys"
+    )

From f74d812b82d864500f5b3ec9b04e1a2fec5b21b9 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:11:04 +0000
Subject: [PATCH 13/38] test(dtype): assert warning fires on
 declared/const_value dtype disagreement

Completes the initializer_dtype contract test (readability nit from
9abb0595): verify the documented 'stale type metadata' warning is
actually emitted when an initializer's declared dtype disagrees with its
const_value, not just that const_value wins the return value.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_dtype_utils_test.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/mobius/_passes/_dtype_utils_test.py b/src/mobius/_passes/_dtype_utils_test.py
index 881ab0d7..9bbdf0c6 100644
--- a/src/mobius/_passes/_dtype_utils_test.py
+++ b/src/mobius/_passes/_dtype_utils_test.py
@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+import logging
+
 import numpy as np
 import onnx_ir as ir
 
@@ -31,10 +33,15 @@ def test_falls_back_to_const_value_when_declared_missing(self):
         assert v.dtype is None
         assert initializer_dtype(v) == ir.DataType.FLOAT16
 
-    def test_const_value_wins_on_disagreement(self):
+    def test_const_value_wins_on_disagreement(self, caplog):
         """Stale declared metadata must not override the serialized data dtype."""
         v = _value(ir.DataType.FLOAT, np.ones((2,), np.float16))
-        assert initializer_dtype(v) == ir.DataType.FLOAT16
+        with caplog.at_level(logging.WARNING, logger="mobius._passes._dtype_utils"):
+            assert initializer_dtype(v) == ir.DataType.FLOAT16
+        assert any(
+            record.levelno == logging.WARNING and "stale type metadata" in record.getMessage().lower()
+            for record in caplog.records
+        ), "expected a warning when declared dtype disagrees with const_value"
 
     def test_returns_none_when_nothing_available(self):
         v = _value(None, None)

From 215df6354da92fb9ce3995a0b974fef5c955bca9 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:14:41 +0000
Subject: [PATCH 14/38] test(static-cache): make decode-Flash proof dispositive
 at head_dim=96
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The decode-on-Flash guard previously ran only at the tiny default head_dim=16,
proving the phase-split wiring but not that ORT's Flash kernel accepts
Phi-3.5's production head_dim=96 (fp16) on the target GPU — the open question
gating variant#2's decode-on-Flash premise. Parametrize the kernel-selection
test over head_dim in {16, 96} (heads/hidden scaled to stay tiny) so it
empirically confirms decode selects Flash and prefill Memory-Efficient at the
real model's head dimension. Verified on A100 (SM80): both head_dims pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/static_cache_decode_test.py | 37 ++++++++++++++++++++++++++++---
 1 file changed, 34 insertions(+), 3 deletions(-)

diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index 8374cce0..2cb6f9c5 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -102,6 +102,7 @@ def _build_static_cache_session(
     *,
     ir_dtype: ir.DataType = ir.DataType.FLOAT,
     enable_profiling: bool = False,
+    config_overrides: dict | None = None,
 ) -> tuple[ort.InferenceSession, object]:
     """Build a tiny static-cache qwen2 graph and load it on CUDA.
 
@@ -119,8 +120,12 @@ def _build_static_cache_session(
             production-precision graph used for the Flash-eligibility guard.
         enable_profiling: Turn on ORT op-level profiling so callers can
             inspect which ``If`` branch executed and with which inputs.
+        config_overrides: Optional ``_base_config`` field overrides.  Used to
+            build at the production ``head_dim`` (Phi-3.5 = 96) so the
+            Flash-eligibility guard is dispositive for the real model's head
+            dimension, not just the tiny default (``head_dim=16``).
     """
-    config = _base_config()
+    config = _base_config(**(config_overrides or {}))
     config = dataclasses.replace(config, dtype=ir_dtype)
     module = registry.get(_MODEL_TYPE)(config)
     task = CausalLMTask(static_cache=True, max_seq_len=_MAX_SEQ_LEN)
@@ -383,7 +388,25 @@ def _selected_attention_kernels(capture_file) -> list[tuple[str, int]]:
     return kernels
 
 
-def test_static_cache_decode_selects_flash_kernel_on_cuda():
+# Phi-3.5's real attention head dimension (32 heads x 96 = 3072 hidden).  The
+# tiny default config uses head_dim=16, which exercises the phase-split wiring
+# but NOT whether ORT's Flash kernel accepts the production head_dim on this
+# GPU.  Building at head_dim=96 (heads/hidden scaled down to stay tiny) makes
+# the Flash-eligibility assertion dispositive for the model we actually ship.
+_PHI35_HEAD_DIM_OVERRIDES = {
+    "hidden_size": 384,
+    "num_attention_heads": 4,
+    "num_key_value_heads": 2,
+    "head_dim": 96,
+}
+
+
+@pytest.mark.parametrize(
+    "config_overrides",
+    [None, _PHI35_HEAD_DIM_OVERRIDES],
+    ids=["tiny-head-dim-16", "phi35-head-dim-96"],
+)
+def test_static_cache_decode_selects_flash_kernel_on_cuda(config_overrides):
     """Decode actually selects Flash; prefill selects Memory-Efficient.
 
     Structural maskless-ness (the test above) is necessary, but the *proof*
@@ -398,10 +421,18 @@ def test_static_cache_decode_selects_flash_kernel_on_cuda():
     the causal mask onto the decode branch would flip its kernel from Flash
     to Memory-Efficient and fail here, even though finiteness / scatter / If-
     count assertions would all still pass.
+
+    Parametrized over head dimension: the tiny default (``head_dim=16``) and
+    Phi-3.5's production ``head_dim=96``.  The latter makes the assertion
+    dispositive — it empirically confirms ORT's Flash kernel accepts the real
+    model's head dimension (fp16) on this GPU, rather than silently routing to
+    Memory-Efficient, which would invalidate the decode-on-Flash premise.
     """
     with tempfile.TemporaryDirectory() as tmp_dir:
         session, config = _build_static_cache_session(
-            tmp_dir, ir_dtype=ir.DataType.FLOAT16
+            tmp_dir,
+            ir_dtype=ir.DataType.FLOAT16,
+            config_overrides=config_overrides,
         )
         num_layers = config.num_hidden_layers
         kv_hidden = config.num_key_value_heads * config.head_dim

From e0fd0c616f23c7409a18e77ea08b7d86476b74a5 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:16:22 +0000
Subject: [PATCH 15/38] docs(static-cache): note _capture_attention_kernel_log
 parallel-safety

Per readability review: the helper mutates process-global state (default
logger severity + fd 2). It restores in finally, so it is safe under the
sequential/xdist-multiprocess way these tests run, but document that it is
not safe to call from threads sharing the process. Comment-only.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/static_cache_decode_test.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index 2cb6f9c5..a8befa20 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -360,6 +360,11 @@ def _capture_attention_kernel_log():
     read the yielded file *inside* the ``with`` block (it is closed on exit);
     fd 2 and the logger severity are restored before exit so assertions made
     after the block still report normally.
+
+    Not safe under *in-process* parallelism: it mutates process-global state
+    (the default logger severity and fd 2).  That is fine here — pytest runs
+    tests sequentially in-process, and ``pytest-xdist`` isolates workers in
+    separate processes — but do not call it from threads sharing this process.
     """
     ort.set_default_logger_severity(0)
     saved_stderr_fd = os.dup(2)

From 03e827bd3ec4e669135530aa7af607ffc149a368 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:21:39 +0000
Subject: [PATCH 16/38] test(static-cache): version-gate the verbose-log Flash
 kernel proof
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Per code review: the decode-on-Flash kernel-name proof reads ORT's internal
VERBOSE attention.cc selection log, an unstable contract. Skip it on any ORT
version != 1.27.x (the validated version) with a message telling maintainers
to re-validate the log strings and bump _VALIDATED_ORT_VERSION, or demote the
proof to the profiling harness. This makes an ORT bump self-announce as a skip
at the moment the assumption is invalidated, instead of a confusing regex-miss
assertion failure. The deterministic structural maskless guard is intentionally
NOT gated — it stays the version-robust backstop. Still runs/passes on 1.27.0.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/static_cache_decode_test.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index a8befa20..2654b662 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -406,6 +406,15 @@ def _selected_attention_kernels(capture_file) -> list[tuple[str, int]]:
 }
 
 
+# ORT version whose VERBOSE attention.cc kernel-selection log strings
+# (``_ATTENTION_KERNEL_LINE``) this guard was validated against.  The log text
+# is an internal, unstable contract; on an ORT bump the kernel-name proof must
+# be re-validated (or demoted to the profiling harness) rather than fail with a
+# confusing regex-miss assertion.  Gating on this version makes that churn
+# self-announce as a skip at the exact moment the assumption is invalidated.
+_VALIDATED_ORT_VERSION = "1.27"
+
+
 @pytest.mark.parametrize(
     "config_overrides",
     [None, _PHI35_HEAD_DIM_OVERRIDES],
@@ -432,7 +441,20 @@ def test_static_cache_decode_selects_flash_kernel_on_cuda(config_overrides):
     dispositive — it empirically confirms ORT's Flash kernel accepts the real
     model's head dimension (fp16) on this GPU, rather than silently routing to
     Memory-Efficient, which would invalidate the decode-on-Flash premise.
+
+    Skips on an unvalidated ORT version: this proof reads ORT's internal
+    VERBOSE kernel-selection log strings, so a version bump must re-validate
+    those strings (the deterministic ``test_static_cache_decode_runs_maskless``
+    structural guard remains the version-robust backstop and is not gated).
     """
+    if not ort.__version__.startswith(_VALIDATED_ORT_VERSION):
+        pytest.skip(
+            f"decode-on-Flash kernel proof reads ORT's VERBOSE attention.cc "
+            f"selection log, validated only on ORT {_VALIDATED_ORT_VERSION}.x "
+            f"(running {ort.__version__}). Re-validate the log strings in "
+            f"_ATTENTION_KERNEL_LINE and bump _VALIDATED_ORT_VERSION, or demote "
+            f"this proof to the profiling harness."
+        )
     with tempfile.TemporaryDirectory() as tmp_dir:
         session, config = _build_static_cache_session(
             tmp_dir,

From 26d74d228feb01808c12da31aee13296db8fb63f Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:23:50 +0000
Subject: [PATCH 17/38] docs(skill): add VALUE-based fp16 GQA weight-integrity
 gotcha (section 5)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A fp16 GQA export can be all-fp16, right-count, and still all-zeros: if the
fold passes leave the packed-QKV Concat output dtype unknown/fp32 while the
data is fp16, the serializer skips it and it loads as near-zero. Count/dtype
checks (section 4c) do NOT catch this — broken and fixed builds can share the
same 197-fp16 initializer profile. Document the canonical VALUE-based gate:
per-slice packed-QKV correlation ~= 1.0 (broken ~= 0) and L2 norm ~= 126.6 at
layer 0, plus end-to-end greedy-argmax parity vs attn_dynamic (~19-20/20).
Cross-link from section 4 and update frontmatter. Per QA @b5d02a20.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../mobius-onnx-export-gotchas/SKILL.md       | 40 ++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
index d3442f15..1ec6e062 100644
--- a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -1,6 +1,6 @@
 ---
 name: mobius-onnx-export-gotchas
-description: Use when building/exporting ONNX models with the `mobius build` CLI (especially Phi-3 / Phi-3.5 or any model with `--execution-provider cuda` GQA fusion and/or `--static-cache`). Covers the current CLI syntax, the dtype flag values, the GQA-vs-static-cache interaction, and how to verify fp16 GQA exports load in onnxruntime (the historical packed-QKV FLOAT32 load bug is fixed as of df203cc).
+description: Use when building/exporting ONNX models with the `mobius build` CLI (especially Phi-3 / Phi-3.5 or any model with `--execution-provider cuda` GQA fusion and/or `--static-cache`). Covers the current CLI syntax, the dtype flag values, the GQA-vs-static-cache interaction, how to verify fp16 GQA exports load in onnxruntime (the historical packed-QKV FLOAT32 load bug is fixed as of df203cc), and why fp16 GQA exports need VALUE-based weight checks (corr≈1.0 / norm), not just initializer count/dtype, to catch silently-zeroed packed-QKV weights.
 ---
 
 # mobius ONNX export gotchas
@@ -94,3 +94,41 @@ onnx.save(m, "model.onnx", save_as_external_data=True, all_tensors_to_one_file=T
 Load the model on `CUDAExecutionProvider` and run one prefill + one decode `session.run`. Confirm:
 (a) the expected attention op (`com.microsoft::GroupQueryAttention` vs `ai.onnx::Attention`),
 (b) finite fp16 logits, (c) no FLOAT32 initializers for an fp16 build.
+
+These checks are **necessary but NOT sufficient** for a fp16 GQA export — see §5. A model can pass all
+three and still have silently-zeroed packed-QKV weights.
+
+## 5. Verifying a fp16 GQA export: use VALUE-based weight checks, NOT initializer count/dtype
+**A fp16 GQA export can be all-fp16, right-count, and still all-zeros — only a corr≈1.0 / norm≈126 VALUE
+check on the packed QKV proves the weights are real.**
+
+### Symptom
+The GQA model loads cleanly (32 `GroupQueryAttention` nodes, all-fp16, finite logits) but generates
+garbage (e.g. `holdou_(...artersarters`). Prefill logits come out ~3× the reference scale, with
+`max|Δlogit|` ~50+ versus the reference.
+
+### Root cause
+The packed-QKV initializer is `Transpose(Concat(q, k, v, axis=0))`. If the fold passes
+(`FoldConcatInitializersPass` / `FoldTransposedInitializerPass`) leave the packed-Concat output dtype
+UNKNOWN / defaulted-to-fp32 while the data is fp16, the serializer **skips** it and it loads as
+**near-zero** — the weights are silently dead. (This is the §3 failure mode; the upstream fix in
+`df203cc` stamps the fp16 dtype at the fold-pass source. A post-hoc cast is NOT a fix — it re-corrupts.)
+
+### Why count/dtype checks fail (the trap)
+The BROKEN export and the FIXED export can have the **same initializer count and the same fp16/fp32 dtype
+ratio** (e.g. both 197 fp16 after dead-weight stripping). Counting initializers or checking
+"0 fp32 / all fp16" does **not** distinguish a healthy model from a zeroed-weight one. §4(c) alone will
+pass a dead model.
+
+### Canonical verification (load-bearing, not optional)
+VALUE-based per-slice check on each packed-QKV initializer against its source q/k/v weights:
+- per-slice correlation **≈ 1.000** (broken ≈ 0.000), AND
+- packed-QKV L2 norm **≈ 126.6** at layer 0 / mean(|abs|) **≈ 0.013** (broken ≈ 0.80 / ≈ 5e-6).
+
+Plus an end-to-end next-token greedy-argmax parity check vs the `attn_dynamic` reference (expect
+**~19–20 / 20**). Isolated single-token divergences are fp16 dead-ties (reference top1−top2 gap = 0.0000),
+not bugs. Optional hardening: assert **0 unused initializers** and that all N packed-QKV initializers are
+present, to catch dead-weight OVER-stripping.
+
+QA's `gqa_weight_integrity_gate.py` (`--self-check --strip-audit --scan-all`, per-layer corr/norm)
+implements exactly this gate.

From cf6c5c45b562a85321864894d003718313c1062d Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:29:23 +0000
Subject: [PATCH 18/38] fix(tasks): stamp explicit present KV-cache output
 shapes for GQA

GroupQueryAttention's contrib-op shape inference mis-derives the present
KV head_dim (32 instead of 96), so present.{i}.key/value graph outputs
declared the wrong head_dim while past_key_values inputs were correct.
ORT logged 'Error merging shape info ... lenient merge' (64 warnings on
Phi-3.5) and any consumer trusting declared shapes (e.g. onnxruntime-genai)
would see inconsistent past-vs-present KV cache types.

_register_kv_cache_outputs now accepts optional batch/num_kv_heads/
key_head_dim/value_head_dim/total_seq_len/dtype; when all provided it
stamps present.{i}.{key,value} symmetric to the past inputs before
add_output. Opt-in: omitting them preserves inference-only behavior for
the other callers. _causal_lm wires concrete values through.

Verified on a real Phi-3.5 GQA export: present.0.key now
[batch,32,past_sequence_len + sequence_len,96]; the 64 present-KV merge
warnings are eliminated; weights byte-identical (corr 1.0 x32); next-token
parity vs attn_dynamic 20/20 identical.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/tasks/_cache_utils.py      |  45 +++++++++-
 src/mobius/tasks/_cache_utils_test.py | 115 ++++++++++++++++++++++++++
 src/mobius/tasks/_causal_lm.py        |  19 ++++-
 src/mobius/tasks/_task_test.py        |  30 +++++++
 4 files changed, 204 insertions(+), 5 deletions(-)
 create mode 100644 src/mobius/tasks/_cache_utils_test.py

diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py
index 41cf5071..9d4671fe 100644
--- a/src/mobius/tasks/_cache_utils.py
+++ b/src/mobius/tasks/_cache_utils.py
@@ -114,13 +114,54 @@ def _register_kv_cache_outputs(
     present_key_values: list[tuple[ir.Value, ir.Value]],
     *,
     prefix: str = "present",
+    batch: ir.SymbolicDim | str | None = None,
+    num_kv_heads: int | None = None,
+    key_head_dim: int | None = None,
+    value_head_dim: int | None = None,
+    total_seq_len: ir.SymbolicDim | str | int | None = None,
+    dtype: ir.DataType | None = None,
 ) -> None:
     """Name and register KV cache outputs on the graph.
 
-    Output shapes and dtypes are inferred by the shape inference pass
-    that runs during model optimization.
+    When every present-shape parameter (``batch``, ``num_kv_heads``,
+    ``key_head_dim``, ``value_head_dim``, ``total_seq_len``, ``dtype``) is
+    supplied, each ``present.{i}.{key,value}`` output is stamped with an
+    explicit ``[batch, num_kv_heads, total_seq_len, head_dim]`` type, symmetric
+    to the ``past_key_values`` inputs created by :func:`_make_kv_cache_inputs`.
+
+    This explicit stamp is required for ``com.microsoft::GroupQueryAttention``
+    exports: that contrib op's shape inference mis-derives the present
+    ``head_dim`` (it divides the *packed* QKV query hidden by
+    ``num_heads + 2 * kv_num_heads`` and lands on the wrong value), so the
+    ``present.*`` outputs would otherwise declare the wrong ``head_dim`` even
+    though the kernel produces correct data at runtime. The mismatch makes ORT
+    log a shape-merge warning and breaks present->past chaining in consumers
+    such as onnxruntime-genai that trust the declared shapes. The plain ONNX
+    ``Attention`` op infers the present shape correctly, so the stamp is a
+    no-op there.
+
+    When the parameters are omitted, output shapes/dtypes are left to the
+    shape inference pass that runs during model optimization.
     """
+    stamp = all(
+        param is not None
+        for param in (
+            batch,
+            num_kv_heads,
+            key_head_dim,
+            value_head_dim,
+            total_seq_len,
+            dtype,
+        )
+    )
     for i, (present_key, present_value) in enumerate(present_key_values):
+        if stamp:
+            present_key.shape = ir.Shape([batch, num_kv_heads, total_seq_len, key_head_dim])
+            present_key.type = ir.TensorType(dtype)
+            present_value.shape = ir.Shape(
+                [batch, num_kv_heads, total_seq_len, value_head_dim]
+            )
+            present_value.type = ir.TensorType(dtype)
         builder.add_output(present_key, f"{prefix}.{i}.key")
         builder.add_output(present_value, f"{prefix}.{i}.value")
 
diff --git a/src/mobius/tasks/_cache_utils_test.py b/src/mobius/tasks/_cache_utils_test.py
new file mode 100644
index 00000000..065bda85
--- /dev/null
+++ b/src/mobius/tasks/_cache_utils_test.py
@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for KV-cache I/O helpers in :mod:`mobius.tasks._cache_utils`."""
+
+from __future__ import annotations
+
+import onnx_ir as ir
+
+from mobius.tasks._base import _make_graph
+from mobius.tasks._cache_utils import _register_kv_cache_outputs
+
+
+def _present_pair(name: str, wrong_head_dim: int) -> tuple[ir.Value, ir.Value]:
+    """A present key/value pair carrying a deliberately wrong inferred shape.
+
+    Mimics ``com.microsoft::GroupQueryAttention`` shape inference, which
+    mis-derives ``head_dim`` (e.g. 32 instead of 96) for the present outputs.
+    """
+    key = ir.Value(name=f"{name}_key")
+    value = ir.Value(name=f"{name}_value")
+    for v in (key, value):
+        v.shape = ir.Shape(["batch", 32, "seq", wrong_head_dim])
+        v.type = ir.TensorType(ir.DataType.FLOAT16)
+    return key, value
+
+
+def _dims(value: ir.Value) -> list[object]:
+    return [d if isinstance(d, int) else str(d) for d in value.shape]
+
+
+class TestRegisterKVCacheOutputs:
+    def test_stamps_explicit_shapes_over_wrong_inference(self):
+        """Explicit params must override mis-inferred present shapes (GQA bug)."""
+        _, builder = _make_graph()
+        pairs = [_present_pair("present.0", wrong_head_dim=32)]
+
+        _register_kv_cache_outputs(
+            builder,
+            pairs,
+            batch=ir.SymbolicDim("batch"),
+            num_kv_heads=32,
+            key_head_dim=96,
+            value_head_dim=96,
+            total_seq_len="past_sequence_len + sequence_len",
+            dtype=ir.DataType.FLOAT16,
+        )
+
+        key, value = pairs[0]
+        assert _dims(key) == ["batch", 32, "past_sequence_len + sequence_len", 96]
+        assert _dims(value) == ["batch", 32, "past_sequence_len + sequence_len", 96]
+        assert key.dtype == ir.DataType.FLOAT16
+        assert value.dtype == ir.DataType.FLOAT16
+
+    def test_distinct_key_value_head_dims(self):
+        """MLA-style caches may use different key/value head dims."""
+        _, builder = _make_graph()
+        pairs = [_present_pair("present.0", wrong_head_dim=32)]
+
+        _register_kv_cache_outputs(
+            builder,
+            pairs,
+            batch=ir.SymbolicDim("batch"),
+            num_kv_heads=16,
+            key_head_dim=192,
+            value_head_dim=128,
+            total_seq_len="past_sequence_len + sequence_len",
+            dtype=ir.DataType.FLOAT16,
+        )
+
+        key, value = pairs[0]
+        assert _dims(key)[1] == 16 and _dims(key)[3] == 192
+        assert _dims(value)[1] == 16 and _dims(value)[3] == 128
+
+    def test_no_params_leaves_shapes_untouched(self):
+        """Without shape params the helper must not stamp (inference path)."""
+        _, builder = _make_graph()
+        pairs = [_present_pair("present.0", wrong_head_dim=32)]
+
+        _register_kv_cache_outputs(builder, pairs)
+
+        key, _ = pairs[0]
+        # Unchanged: still the (wrong) pre-existing inferred shape.
+        assert _dims(key) == ["batch", 32, "seq", 32]
+
+    def test_partial_params_do_not_stamp(self):
+        """All shape params are required; a partial set falls back to inference."""
+        _, builder = _make_graph()
+        pairs = [_present_pair("present.0", wrong_head_dim=32)]
+
+        _register_kv_cache_outputs(
+            builder,
+            pairs,
+            batch=ir.SymbolicDim("batch"),
+            num_kv_heads=32,
+            # key_head_dim / value_head_dim / total_seq_len / dtype omitted
+        )
+
+        key, _ = pairs[0]
+        assert _dims(key) == ["batch", 32, "seq", 32]
+
+    def test_registers_named_outputs(self):
+        """Outputs are registered with the conventional present.{i}.* names."""
+        _, builder = _make_graph()
+        pairs = [_present_pair("a", 32), _present_pair("b", 32)]
+
+        _register_kv_cache_outputs(builder, pairs)
+
+        names = [v.name for v in builder.graph.outputs]
+        assert names == [
+            "present.0.key",
+            "present.0.value",
+            "present.1.key",
+            "present.1.value",
+        ]
diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py
index 471d350b..972437c7 100644
--- a/src/mobius/tasks/_causal_lm.py
+++ b/src/mobius/tasks/_causal_lm.py
@@ -153,6 +153,10 @@ def build(
             num_kv_cache_heads = (
                 config.num_attention_heads if use_mla else config.num_key_value_heads
             )
+            kv_key_head_dim = (
+                (config.qk_nope_head_dim or 0) + (config.qk_rope_head_dim or 0)
+            ) or config.head_dim
+            kv_value_head_dim = config.v_head_dim or config.head_dim
 
             past_key_values = _make_kv_cache_inputs(
                 builder,
@@ -162,9 +166,8 @@ def build(
                 config.dtype,
                 batch,
                 past_seq_len,
-                key_head_dim=((config.qk_nope_head_dim or 0) + (config.qk_rope_head_dim or 0))
-                or None,
-                value_head_dim=config.v_head_dim or None,
+                key_head_dim=kv_key_head_dim,
+                value_head_dim=kv_value_head_dim,
             )
 
         logits, present_key_values = module(
@@ -184,9 +187,19 @@ def build(
                 present_key_values,
             )
         else:
+            # Stamp explicit present shapes symmetric to the past inputs so the
+            # com.microsoft::GroupQueryAttention export declares the correct
+            # head_dim (its contrib-op shape inference otherwise mis-derives it
+            # from the packed QKV hidden). total_seq = past + current sequence.
             _register_kv_cache_outputs(
                 builder,
                 present_key_values,
+                batch=batch,
+                num_kv_heads=num_kv_cache_heads,
+                key_head_dim=kv_key_head_dim,
+                value_head_dim=kv_value_head_dim,
+                total_seq_len="past_sequence_len + sequence_len",
+                dtype=config.dtype,
             )
 
         return ModelPackage({"model": _make_model(graph)}, config=config)
diff --git a/src/mobius/tasks/_task_test.py b/src/mobius/tasks/_task_test.py
index d50f4337..875618a9 100644
--- a/src/mobius/tasks/_task_test.py
+++ b/src/mobius/tasks/_task_test.py
@@ -91,6 +91,36 @@ def test_build_outputs(self):
         assert "present.0.key" in output_names
         assert "present.0.value" in output_names
 
+    def test_present_outputs_match_past_inputs(self):
+        """present.{i}.* must declare the same kv_heads/head_dim/dtype as the
+        corresponding past_key_values.{i}.* inputs, with total (past+current)
+        sequence length. Guards the GQA present-head_dim export bug, where the
+        contrib-op shape inference would otherwise mis-declare head_dim.
+        """
+        config = make_config()
+        module = CausalLMModel(config)
+        pkg = CausalLMTask().build(module, config)
+        model = pkg["model"]
+        inputs = {v.name: v for v in model.graph.inputs}
+        outputs = {v.name: v for v in model.graph.outputs}
+
+        def dims(value):
+            return [d if isinstance(d, int) else str(d) for d in value.shape]
+
+        for i in range(config.num_hidden_layers):
+            for kind in ("key", "value"):
+                past = inputs[f"past_key_values.{i}.{kind}"]
+                present = outputs[f"present.{i}.{kind}"]
+                past_dims = dims(past)
+                present_dims = dims(present)
+                # batch, kv_heads, head_dim must match the past input exactly.
+                assert present_dims[0] == past_dims[0]
+                assert present_dims[1] == past_dims[1]
+                assert present_dims[3] == past_dims[3]
+                # present covers past + current tokens.
+                assert present_dims[2] == "past_sequence_len + sequence_len"
+                assert present.dtype == past.dtype
+
     def test_build_producer_info(self):
         config = make_config()
         module = CausalLMModel(config)

From 091174981e1fbf648b0fc2837176544d3c2dd943 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:32:55 +0000
Subject: [PATCH 19/38] docs(skill): add GQA present-output head_dim gotcha
 (section 6)

Document the GQA present.{i}.key/value head_dim metadata bug and its fix
(cf6c5c4): symptom (64 'lenient merge' warnings, declared head_dim 32 vs
past's 96), root cause (GroupQueryAttention contrib-op shape inference),
the opt-in shape-stamping fix, a verify snippet, and the separate
pre-existing internal GQA-hidden-output value_info warning (1024 vs 3072)
that remains as a follow-up.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../mobius-onnx-export-gotchas/SKILL.md       | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
index 1ec6e062..7cbf1578 100644
--- a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -132,3 +132,49 @@ present, to catch dead-weight OVER-stripping.
 
 QA's `gqa_weight_integrity_gate.py` (`--self-check --strip-audit --scan-all`, per-layer corr/norm)
 implements exactly this gate.
+
+## 6. FIXED: GQA `present.*` KV-cache outputs declared the wrong `head_dim`
+**Status: fixed as of commit `cf6c5c4`.** A native fp16 GQA export now declares
+`present.{i}.{key,value}` with the correct `head_dim`, symmetric to its `past_key_values.{i}.*` inputs.
+
+### Symptom (pre-fix)
+The graph **output** `present.{i}.key/value` declared the wrong `head_dim` (e.g. `32` instead of the real
+`96` on Phi-3.5) while the matching `past_key_values.{i}.*` **input** was correct (`96`). At load ORT logged
+(once per key+value per layer — 64 on Phi-3.5):
+
+```
+[W ...MergeShapeInfo] Error merging shape info for output. 'present.0.key'
+source:{-1,32,-1,96} target:{-1,32,-1,32}. Falling back to lenient merge.
+```
+
+Runtime still produced correct (96-wide) arrays via lenient merge, but any consumer that **trusts declared
+shapes** (e.g. `onnxruntime-genai`) would see inconsistent past-vs-present KV cache types.
+
+### Root cause
+`GroupQueryAttention`'s contrib-op shape inference mis-derives the present `head_dim` (it does **not**
+reproduce on the plain `Attention` op, which infers correctly). `_register_kv_cache_outputs`
+(`src/mobius/tasks/_cache_utils.py`) added the present outputs with **no explicit shape**, so the buggy
+inference won.
+
+### The fix
+`_register_kv_cache_outputs` now opt-in **stamps** `present.{i}.{key,value}` shape+dtype symmetric to the
+past inputs when the caller passes `batch`/`num_kv_heads`/`key_head_dim`/`value_head_dim`/`total_seq_len`/
+`dtype` (wired from `_causal_lm.py`). Omitting them preserves inference-only behavior, so the other ~10
+callers are unaffected. The stamp survives `SymbolicShapeInferencePass` (policy `refine` only tightens
+unknown dims; it won't replace a concrete `96` with a conflicting `32`).
+
+### Verify
+```python
+import onnx
+m = onnx.load("model.onnx", load_external_data=False)
+d = lambda vi: [(x.dim_param or x.dim_value) for x in vi.type.tensor_type.shape.dim]
+o = {v.name: v for v in m.graph.output}
+print("present.0.key:", d(o["present.0.key"]))   # head_dim must equal the past input's (e.g. 96, NOT 32)
+```
+
+### Known remaining (separate, pre-existing, harmless)
+ORT still logs ~32 `Error merging shape info ... source:{-1,-1,3072} target:{-1,-1,1024}` warnings on the
+GQA op's **internal hidden-state output** value_info (`v_*.GroupQueryAttention_*_0`, `1024`=32×32 vs the
+correct `3072`=32×96). That value is **not** a declared graph I/O — runtime is correct and `onnxruntime-genai`
+does not trust it — so it does not bite shape-trusting consumers the way the present-output bug did. Tracked
+as a follow-up in the GQA rewrite emission path (not the KV-cache output path).

From 5286ddad3d23a3613edc956601057961854f6789 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:33:27 +0000
Subject: [PATCH 20/38] docs(skill): add initializer_dtype() convention to
 prevent fp32-default footgun

Per critical-reviewer note: any future pass materializing an initializer
should resolve dtype via initializer_dtype() instead of 'value.dtype or
ir.DataType.FLOAT', since _cast_module_dtype drops the declared .dtype while
keeping the fp16 const_value. Documents the convention and the follow-up to
grep _passes/ for siblings / re-stamp at source in _cast_module_dtype.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .agents/skills/mobius-onnx-export-gotchas/SKILL.md | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
index 7cbf1578..1a04652b 100644
--- a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -72,6 +72,16 @@ fp32 = [i.name for i in m.graph.initializer if i.data_type == onnx.TensorProto.F
 print(len(fp32), "FLOAT32 initializers (should be 0 for fp16)")
 ```
 
+### Convention (prevents the whole class from reappearing)
+Any pass that **materializes a new initializer** must resolve its dtype via
+`initializer_dtype()` (`src/mobius/_passes/_dtype_utils.py`), **never** `value.dtype or ir.DataType.FLOAT`.
+The bug class originates in `_cast_module_dtype` dropping a `Value`'s declared `.dtype` (→ `None`) while its
+`const_value` stays fp16; a bare `.dtype or FLOAT` fallback then silently mis-types the result as fp32. Fold
+passes (`_fold_concat.py`, `_fold_transpose.py`) already follow this; mirror it in any future
+initializer-producing pass. Siblings still reading `.dtype` directly remain exposed — a follow-up should
+grep `_passes/` for `.dtype or ir.DataType` and consider re-stamping the type in `_cast_module_dtype` to kill
+the class at source.
+
 ### Salvaging a stale pre-fix artifact (only if re-exporting is not an option)
 Prefer re-exporting on the fixed code. If you must repair an old model, cast its FLOAT32 initializers to
 fp16 and re-save. **Gotcha when re-saving with external data:** if you save with `location="X.data"` and

From cd2ea145e317e27d8c3b1b5b7739ae51c0519e94 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:38:33 +0000
Subject: [PATCH 21/38] docs(skill): correct fp16 GQA mean|abs| to 0.015 and
 add signed-mean caveat

Per QA @b5d02a20 re-measured live on canonical phi35_gqa: L0 packed-QKV
mean(|abs|) is 0.01505 (was 0.013). Add a critical caveat: the good model's
SIGNED mean is ~2.6e-6 (weights are symmetric +/-), coincidentally resembling
the broken model's mean|abs| ~5e-6, so the discriminator must be mean-of-abs
or norm, never signed mean -- signed mean falsely flags the good model and
already caused a crew false alarm.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .agents/skills/mobius-onnx-export-gotchas/SKILL.md | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
index 1a04652b..9077584f 100644
--- a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -133,7 +133,13 @@ pass a dead model.
 ### Canonical verification (load-bearing, not optional)
 VALUE-based per-slice check on each packed-QKV initializer against its source q/k/v weights:
 - per-slice correlation **≈ 1.000** (broken ≈ 0.000), AND
-- packed-QKV L2 norm **≈ 126.6** at layer 0 / mean(|abs|) **≈ 0.013** (broken ≈ 0.80 / ≈ 5e-6).
+- packed-QKV L2 norm **≈ 126.6** at layer 0 / mean(|abs|) **≈ 0.015** (broken ≈ 0.80 / ≈ 5e-6).
+
+> ⚠️ **Use mean-of-ABS or norm — NEVER the signed mean.** The good model's *signed* mean is ~2.6e-6
+> (near zero, because the weights are symmetric ±), which coincidentally looks just like the broken
+> model's mean(|abs|) ~5e-6. Checking signed mean would **falsely flag the good model as broken** — this
+> exact confusion has already caused a false alarm in this crew. Valid discriminators: mean(|abs|)
+> (good ≈ 0.015 vs broken ≈ 5e-6) or L2 norm (good ≈ 126.6 vs broken ≈ 0.80).
 
 Plus an end-to-end next-token greedy-argmax parity check vs the `attn_dynamic` reference (expect
 **~19–20 / 20**). Isolated single-token divergences are fp16 dead-ties (reference top1−top2 gap = 0.0000),

From 86a166c45cf25c364b70066e79dd35a9df154cd3 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:42:40 +0000
Subject: [PATCH 22/38] =?UTF-8?q?docs(skill):=20generalize=20fp16=20GQA=20?=
 =?UTF-8?q?count-trap=20=E2=80=94=20init=20count=20unstable=20across=20fix?=
 =?UTF-8?q?es?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Per Secretary's definitive broadcast: the fp16-init count is not just equal
between broken/fixed builds, it is unstable across fixes (Phi-3.5 shifted
~197 -> ~293 with dead-weight stripping / fold changes) and carries no
correctness signal. Replace the fixed '197' example with the moving range and
state explicitly: never gate on count, use the VALUE gate (corr/norm + parity).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .agents/skills/mobius-onnx-export-gotchas/SKILL.md | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
index 9077584f..71616fff 100644
--- a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -126,9 +126,11 @@ UNKNOWN / defaulted-to-fp32 while the data is fp16, the serializer **skips** it
 
 ### Why count/dtype checks fail (the trap)
 The BROKEN export and the FIXED export can have the **same initializer count and the same fp16/fp32 dtype
-ratio** (e.g. both 197 fp16 after dead-weight stripping). Counting initializers or checking
-"0 fp32 / all fp16" does **not** distinguish a healthy model from a zeroed-weight one. §4(c) alone will
-pass a dead model.
+ratio**, so neither is a validity signal. Worse, the fp16-init count is **not even stable across fixes**
+— on Phi-3.5 it shifted from ~197 to ~293 as dead-weight stripping and fold behavior changed, with no
+bearing on correctness. Counting initializers or checking "0 fp32 / all fp16" does **not** distinguish a
+healthy model from a zeroed-weight one. §4(c) alone will pass a dead model. **Never gate on the count;
+use the VALUE gate below.**
 
 ### Canonical verification (load-bearing, not optional)
 VALUE-based per-slice check on each packed-QKV initializer against its source q/k/v weights:

From 35d08f632486c6a931bda7f0bc90bd24b19be005 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:44:56 +0000
Subject: [PATCH 23/38] test(fold-concat): name the live-packed-result
 invariant in DCE assert

Add a failure message to the 'packed concat survives DCE' assertion so a
future regression self-describes the invariant (live packed-QKV result must
not be stripped) instead of failing bare. Readability-review nit on 71e84b3.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_concat_test.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py
index bff9f2ee..f8b1095f 100644
--- a/src/mobius/_passes/_fold_concat_test.py
+++ b/src/mobius/_passes/_fold_concat_test.py
@@ -167,7 +167,10 @@ def test_folded_inputs_detached_so_dce_strips_dead_weights(self):
         assert "init_0" not in remaining and "init_1" not in remaining, (
             f"Dead pre-pack weights survived DCE: {remaining}"
         )
-        assert "init_0__init_1__axis_0__concat" in remaining
+        assert "init_0__init_1__axis_0__concat" in remaining, (
+            "the live packed result must NOT be stripped by DCE: "
+            f"{remaining}"
+        )
 
     def test_uses_lazy_tensor(self):
         """The packed initializer wraps sources in a LazyTensor.

From 4e9bb5a29973162cf9fd99ca123e02ce24a4c3bb Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Mon, 1 Jun 2026 23:50:04 +0000
Subject: [PATCH 24/38] test(fold-concat): assert survived packed-QKV values
 are exact after DCE

Strengthen the live-weight guard from name-only to value-equality: compare
the survived packed initializer's const_value against the expected
concatenation, so a future DCE that mutates (not just drops) retained
tensors is caught. Code-review nit MINOR-2 on 71e84b3.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_concat_test.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py
index f8b1095f..db319d3e 100644
--- a/src/mobius/_passes/_fold_concat_test.py
+++ b/src/mobius/_passes/_fold_concat_test.py
@@ -171,6 +171,14 @@ def test_folded_inputs_detached_so_dce_strips_dead_weights(self):
             "the live packed result must NOT be stripped by DCE: "
             f"{remaining}"
         )
+        # The survived packed weight must retain its exact values, not just its
+        # name — guards against a future DCE that mutates retained tensors.
+        packed = model.graph.initializers["init_0__init_1__axis_0__concat"]
+        np.testing.assert_array_equal(
+            packed.const_value.numpy(),
+            np.concatenate([a, b], axis=0),
+            err_msg="packed-QKV values corrupted by DCE",
+        )
 
     def test_uses_lazy_tensor(self):
         """The packed initializer wraps sources in a LazyTensor.

From b3b08ccf98904039d1ecf6b176711d1eb367783a Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Tue, 2 Jun 2026 00:05:22 +0000
Subject: [PATCH 25/38] =?UTF-8?q?test(fold-concat):=20gate=20packed-QKV=20?=
 =?UTF-8?q?values=20through=20serialize=E2=86=92reload?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Existing fold-pass tests assert the packed value in memory and that ORT
can load+run the folded model, but none compare the *serialized*
packed-QKV weight to its source q/k/v projections. The original
garbage-export bug (df203cc) corrupted bytes at serialization — fp16 data
written under a defaulted FLOAT32 dtype — which an in-memory const_value
check cannot see and a load+run check misses (the model still loads and
emits a right-shaped fp16 output).

Add a value gate that round-trips through the production save path
(ir.save with external data, like the real fp16 export's model.onnx +
model.onnx.data), reloads, and asserts the packed weight matches its
sources per-slice (Pearson corr >= 0.99, norm rel_err <= 2%, plus exact
fp16 equality) and that ORT inference matches a numpy reference. This
converts the manual QA weight-integrity discriminator (corr=1.0/norm~126)
into a CI guard against a numerically-corrupt pack that still has the
right count and dtype.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_concat_test.py | 138 ++++++++++++++++++++++--
 1 file changed, 128 insertions(+), 10 deletions(-)

diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py
index db319d3e..3da9116b 100644
--- a/src/mobius/_passes/_fold_concat_test.py
+++ b/src/mobius/_passes/_fold_concat_test.py
@@ -168,8 +168,7 @@ def test_folded_inputs_detached_so_dce_strips_dead_weights(self):
             f"Dead pre-pack weights survived DCE: {remaining}"
         )
         assert "init_0__init_1__axis_0__concat" in remaining, (
-            "the live packed result must NOT be stripped by DCE: "
-            f"{remaining}"
+            f"the live packed result must NOT be stripped by DCE: {remaining}"
         )
         # The survived packed weight must retain its exact values, not just its
         # name — guards against a future DCE that mutates retained tensors.
@@ -569,8 +568,7 @@ def test_packed_dtype_follows_const_value_when_declared_dtype_missing(self):
             f"Packed initializer declared dtype should be FLOAT16, got {packed.dtype}"
         )
         assert packed.const_value.dtype == ir.DataType.FLOAT16, (
-            "Packed LazyTensor dtype should be FLOAT16, got "
-            f"{packed.const_value.dtype}"
+            f"Packed LazyTensor dtype should be FLOAT16, got {packed.const_value.dtype}"
         )
         assert packed.const_value.numpy().dtype == np.float16
 
@@ -661,9 +659,7 @@ def test_folded_fp16_concat_matmul_loads_in_ort(self, tmp_path):
         hidden = ir.Value(
             name="hidden", shape=ir.Shape(["N", 8]), type=ir.TensorType(ir.DataType.FLOAT16)
         )
-        matmul_node = ir.Node(
-            "", "MatMul", inputs=[hidden, packed_out], num_outputs=1
-        )
+        matmul_node = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1)
         out = matmul_node.outputs[0]
         out.shape = ir.Shape(["N", 12])
         out.dtype = ir.DataType.FLOAT16
@@ -691,10 +687,132 @@ def test_folded_fp16_concat_matmul_loads_in_ort(self, tmp_path):
         onnx.save(ir.to_proto(model), str(model_path))
 
         # Before the fix this raises a fp16/fp32 MatMul type-mismatch at load.
-        sess = ort.InferenceSession(
-            str(model_path), providers=["CPUExecutionProvider"]
-        )
+        sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
         feed = {"hidden": np.random.randn(2, 8).astype(np.float16)}
         (result,) = sess.run(None, feed)
         assert result.shape == (2, 12)
         assert result.dtype == np.float16
+
+    def test_folded_fp16_packed_qkv_values_survive_serialization(self, tmp_path):
+        """Value gate: packed-QKV must keep its fp16 *values* through serialize→reload.
+
+        The other fold tests assert the packed value in memory and that ORT can
+        load+run the model, but none compare the *serialized* packed-QKV weight
+        to its source q/k/v projections.  The original garbage-export bug
+        corrupted bytes at serialization (fp16 data written under a defaulted
+        FLOAT32 dtype → near-zero/garbage) — a failure an in-memory
+        ``const_value`` check cannot see, and which a load+run check misses
+        because the model still loads and produces a right-shaped fp16 output.
+
+        This round-trips through the production save path (``ir.save`` with
+        external data, exactly like the real fp16 export's model.onnx +
+        model.onnx.data) and asserts the reloaded packed weight matches its
+        sources per-slice, catching a numerically-corrupt pack that still has
+        the right count and dtype (the corr≈0 / norm≈0.8 failure mode).
+        """
+        import onnxruntime as ort
+
+        rng = np.random.default_rng(0)
+        # fp16 q/k/v projection weights with UNSET declared dtype (the state left
+        # by _cast_module_dtype). Sized > 256 bytes so ir.save externalizes them,
+        # exercising the same external-data path as the real export. Column-concat
+        # (axis=1) mirrors the real packed-QKV layout [Q | K | V].
+        q_arr = rng.standard_normal((16, 16)).astype(np.float16)
+        k_arr = rng.standard_normal((16, 16)).astype(np.float16)
+        v_arr = rng.standard_normal((16, 16)).astype(np.float16)
+        q = ir.Value(name="q_weight", shape=ir.Shape([16, 16]))
+        q.const_value = ir.tensor(q_arr)
+        k = ir.Value(name="k_weight", shape=ir.Shape([16, 16]))
+        k.const_value = ir.tensor(k_arr)
+        v = ir.Value(name="v_weight", shape=ir.Shape([16, 16]))
+        v.const_value = ir.tensor(v_arr)
+
+        concat_node = ir.Node(
+            "",
+            "Concat",
+            inputs=[q, k, v],
+            attributes=[ir.Attr("axis", ir.AttributeType.INT, 1)],
+            num_outputs=1,
+        )
+        packed_out = concat_node.outputs[0]  # (16, 48) packed QKV weight
+        packed_out.shape = ir.Shape([16, 48])
+
+        hidden = ir.Value(
+            name="hidden", shape=ir.Shape(["N", 16]), type=ir.TensorType(ir.DataType.FLOAT16)
+        )
+        matmul_node = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1)
+        out = matmul_node.outputs[0]
+        out.shape = ir.Shape(["N", 48])
+        out.dtype = ir.DataType.FLOAT16
+
+        graph = ir.Graph(
+            inputs=[hidden],
+            outputs=[out],
+            nodes=[concat_node, matmul_node],
+            name="qkv_matmul",
+            opset_imports={"": 20},
+        )
+        graph.register_initializer(q)
+        graph.register_initializer(k)
+        graph.register_initializer(v)
+        model = ir.Model(graph, ir_version=10)
+
+        FoldConcatInitializersPass()(model)
+        # DCE step the real export runs after folding.
+        for dead in ("q_weight", "k_weight", "v_weight"):
+            del model.graph.initializers[dead]
+
+        # Serialize through the production path: external data, exactly as the
+        # real fp16 export writes model.onnx + model.onnx.data. This is where the
+        # original dtype bug corrupted the packed bytes.
+        model_path = tmp_path / "model.onnx"
+        ir.save(model, model_path, external_data="model.onnx.data")
+        assert (tmp_path / "model.onnx.data").exists(), (
+            "weights should be externalized, exercising the real export's save path"
+        )
+
+        reloaded = ir.load(model_path)
+        packed_name = "q_weight__k_weight__v_weight__axis_1__concat"
+        packed = reloaded.graph.initializers[packed_name].const_value.numpy()
+        assert packed.dtype == np.float16, (
+            f"reloaded packed-QKV must serialize as fp16, got {packed.dtype}"
+        )
+        expected = np.concatenate([q_arr, k_arr, v_arr], axis=1)
+        assert packed.shape == expected.shape
+
+        # Per-slice discriminator (mirrors QA's weight-integrity gate): each q/k/v
+        # column block of the reloaded pack must match its source projection. A
+        # serialize-time dtype corruption surfaces here as corr≈0 / norm collapse
+        # even when the pack's shape and dtype look correct.
+        blocks = {
+            "q": (packed[:, 0:16], q_arr),
+            "k": (packed[:, 16:32], k_arr),
+            "v": (packed[:, 32:48], v_arr),
+        }
+        for name, (got, src) in blocks.items():
+            got32 = got.astype(np.float32).ravel()
+            src32 = src.astype(np.float32).ravel()
+            corr = np.corrcoef(got32, src32)[0, 1]
+            assert corr >= 0.99, (
+                f"{name}-slice packed-QKV corr={corr:.4f} (<0.99) — corrupted pack"
+            )
+            src_norm = np.linalg.norm(src32)
+            rel_err = abs(np.linalg.norm(got32) - src_norm) / src_norm
+            assert rel_err <= 0.02, (
+                f"{name}-slice packed-QKV norm rel_err={rel_err:.4f} (>2%) — corrupted pack"
+            )
+
+        # fp16 concat is lossless, so the strongest assert also holds end-to-end.
+        np.testing.assert_array_equal(
+            packed,
+            expected,
+            err_msg="packed-QKV bytes corrupted by the serialization round-trip",
+        )
+
+        # Functional check: ORT inference on the reloaded model must match a numpy
+        # reference, proving the packed weight is correct in use, not just on disk.
+        sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
+        feed = {"hidden": rng.standard_normal((2, 16)).astype(np.float16)}
+        (result,) = sess.run(None, feed)
+        reference = feed["hidden"].astype(np.float32) @ expected.astype(np.float32)
+        np.testing.assert_allclose(result.astype(np.float32), reference, rtol=1e-2, atol=1e-2)

From ec5afb98e997d8af5460ee5e60a5025e2726fbf2 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Tue, 2 Jun 2026 00:13:45 +0000
Subject: [PATCH 26/38] test(fold-concat): add mean|abs| degeneracy assert +
 poison negative control
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Strengthen the packed-QKV serialize→reload value gate per QA/code-review
follow-up:

- Add a per-slice mean|abs| >= 1e-3 non-degeneracy assert. mean of ABSOLUTE
  values (not signed mean) is the robust discriminator for the near-zero
  'unserialized' failure mode: symmetric fp16 weights have a signed mean
  ~1e-6 that is indistinguishable from a broken tensor, and corr is undefined
  (nan) for a zero-variance slice. mean|abs| separates healthy (~0.0x) from
  broken (~1e-6) cleanly.
- Add test_value_gate_catches_corrupted_packed_slice: a negative control that
  zeroes the K slice, round-trips through serialize→reload, and asserts the
  discriminators flag it (and survive the round-trip) while the untouched Q/V
  slices still read healthy. Proves the value gate actually has teeth, so a
  future change cannot silently neuter the asserts and stay green.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_concat_test.py | 74 +++++++++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py
index 3da9116b..2c5b309e 100644
--- a/src/mobius/_passes/_fold_concat_test.py
+++ b/src/mobius/_passes/_fold_concat_test.py
@@ -801,6 +801,16 @@ def test_folded_fp16_packed_qkv_values_survive_serialization(self, tmp_path):
             assert rel_err <= 0.02, (
                 f"{name}-slice packed-QKV norm rel_err={rel_err:.4f} (>2%) — corrupted pack"
             )
+            # Non-degenerate magnitude. Use mean of ABSOLUTE values, not signed
+            # mean: symmetric fp16 weights have a signed mean ~1e-6 that is
+            # indistinguishable from a broken near-zero tensor, whereas mean|abs|
+            # separates cleanly (healthy ~0.0x vs unserialized ~1e-6). This also
+            # catches the near-zero failure mode robustly where corr is undefined
+            # (a constant/zero slice has zero variance → corrcoef is nan).
+            mean_abs = float(np.abs(got32).mean())
+            assert mean_abs >= 1e-3, (
+                f"{name}-slice packed-QKV mean|abs|={mean_abs:.2e} (~0) — near-zero/unserialized"
+            )
 
         # fp16 concat is lossless, so the strongest assert also holds end-to-end.
         np.testing.assert_array_equal(
@@ -816,3 +826,67 @@ def test_folded_fp16_packed_qkv_values_survive_serialization(self, tmp_path):
         (result,) = sess.run(None, feed)
         reference = feed["hidden"].astype(np.float32) @ expected.astype(np.float32)
         np.testing.assert_allclose(result.astype(np.float32), reference, rtol=1e-2, atol=1e-2)
+
+    def test_value_gate_catches_corrupted_packed_slice(self, tmp_path):
+        """Negative control: the per-slice value gate must FAIL on a bad pack.
+
+        Proves the discriminators in
+        ``test_folded_fp16_packed_qkv_values_survive_serialization`` actually
+        catch the original failure signature — a packed-QKV with a near-zero
+        slice (corr≈0 / norm collapse) — and that the corruption survives the
+        serialize→reload round-trip rather than being silently "healed" by
+        save/load.  Without this negative control a future change could neuter
+        the value asserts and still go green.
+        """
+        rng = np.random.default_rng(1)
+        q_arr = rng.standard_normal((16, 16)).astype(np.float16)
+        k_arr = rng.standard_normal((16, 16)).astype(np.float16)
+        v_arr = rng.standard_normal((16, 16)).astype(np.float16)
+
+        # Poison the K slice to zero — the exact "unserialized / near-zero packed
+        # weight" signature of the original garbage export.
+        poisoned = np.concatenate([q_arr, k_arr, v_arr], axis=1).copy()
+        poisoned[:, 16:32] = 0
+
+        packed = ir.Value(
+            name="packed_qkv",
+            shape=ir.Shape([16, 48]),
+            type=ir.TensorType(ir.DataType.FLOAT16),
+        )
+        packed.const_value = ir.tensor(poisoned)
+        graph = ir.Graph(
+            inputs=[],
+            outputs=[packed],
+            nodes=[],
+            name="poisoned_pack",
+            opset_imports={"": 20},
+        )
+        graph.register_initializer(packed)
+        model = ir.Model(graph, ir_version=10)
+
+        model_path = tmp_path / "model.onnx"
+        ir.save(model, model_path, external_data="model.onnx.data")
+        reloaded = ir.load(model_path)
+        got = reloaded.graph.initializers["packed_qkv"].const_value.numpy()
+
+        # The corruption must survive the round-trip (save/load must not "fix" it).
+        np.testing.assert_array_equal(got[:, 16:32], np.zeros((16, 16), dtype=np.float16))
+
+        # The gate's discriminators must flag the zeroed K slice. mean|abs| is the
+        # robust primary signal: corr is undefined (nan) for a zero-variance slice,
+        # which is exactly why a corr-only gate would be unsafe here.
+        k_block = got[:, 16:32].astype(np.float32).ravel()
+        k_ref = k_arr.astype(np.float32).ravel()
+        assert np.abs(k_block).mean() < 1e-3, (
+            "mean|abs| discriminator must flag a zeroed packed slice"
+        )
+        rel_err = abs(np.linalg.norm(k_block) - np.linalg.norm(k_ref)) / np.linalg.norm(k_ref)
+        assert rel_err > 0.02, "norm rel_err discriminator must flag a zeroed packed slice"
+
+        # The untouched Q and V slices must still read as healthy — the gate is
+        # specific to the corrupted slice, not a blanket failure.
+        for sl, ref in ((got[:, 0:16], q_arr), (got[:, 32:48], v_arr)):
+            sl32 = sl.astype(np.float32).ravel()
+            assert np.abs(sl32).mean() >= 1e-3
+            corr = np.corrcoef(sl32, ref.astype(np.float32).ravel())[0, 1]
+            assert corr >= 0.99

From ef58730e25a27d8df5fda9a593fe2b57dc92ec0f Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Tue, 2 Jun 2026 00:17:00 +0000
Subject: [PATCH 27/38] docs(skill): fix fp16 GQA count-trap direction to
 293->197

Per Secretary + @69ff092d read-only verification: the count goes 293 -> 197,
not 197 -> 293. 293 = unstripped intermediate (packed-QKV plus ~96 dead
unpacked q/k/v source inits); the safe dead-weight strip removes the dead
pre-pack inits -> 197 (canonical published model). Also note the old broken
export was likewise 197 fp16, so even a correct final count proves nothing.
Conclusion (never gate on count, use the VALUE gate) is unchanged.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .agents/skills/mobius-onnx-export-gotchas/SKILL.md | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
index 71616fff..0ff4c8b5 100644
--- a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
+++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md
@@ -127,8 +127,10 @@ UNKNOWN / defaulted-to-fp32 while the data is fp16, the serializer **skips** it
 ### Why count/dtype checks fail (the trap)
 The BROKEN export and the FIXED export can have the **same initializer count and the same fp16/fp32 dtype
 ratio**, so neither is a validity signal. Worse, the fp16-init count is **not even stable across fixes**
-— on Phi-3.5 it shifted from ~197 to ~293 as dead-weight stripping and fold behavior changed, with no
-bearing on correctness. Counting initializers or checking "0 fp32 / all fp16" does **not** distinguish a
+— on Phi-3.5 it moved from ~293 down to ~197 (an unstripped intermediate carries the packed-QKV plus the
+now-dead unpacked q/k/v source initializers; a safe dead-weight strip then removes the ~96 dead pre-pack
+inits), with no bearing on correctness. Note the OLD broken export was *also* 197 fp16, so even a "right"
+final count proves nothing. Counting initializers or checking "0 fp32 / all fp16" does **not** distinguish a
 healthy model from a zeroed-weight one. §4(c) alone will pass a dead model. **Never gate on the count;
 use the VALUE gate below.**
 

From 24fec659100902a5c8a21d9170cc02ccfe05dd7a Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Tue, 2 Jun 2026 00:21:21 +0000
Subject: [PATCH 28/38] Warn on partial present-shape param set in
 _register_kv_cache_outputs

The present-shape stamp in _register_kv_cache_outputs is all-or-nothing:
all six params stamp the explicit GQA present.* type, none opts out to
inference. A partial set silently fell back to the known-wrong inference
path (the exact head_dim mis-derivation cf6c5c4 fixes), which is almost
always a caller wiring slip rather than an intentional opt-out.

Emit a logger.warning naming the missing parameters when a strict subset
is provided, so the slip is loud rather than silent. Behavior is otherwise
unchanged (still falls back to inference); document the all-or-nothing
contract in the docstring. Tests assert the partial path warns + names the
omitted params, and that the zero-param opt-out stays silent.

Addresses readability-review nit (9abb0595) on cf6c5c4.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/tasks/_cache_utils.py      | 40 ++++++++++++++++++++-------
 src/mobius/tasks/_cache_utils_test.py | 32 +++++++++++++--------
 2 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py
index 9d4671fe..3e001e7d 100644
--- a/src/mobius/tasks/_cache_utils.py
+++ b/src/mobius/tasks/_cache_utils.py
@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import logging
 from typing import NamedTuple
 
 import onnx_ir as ir
@@ -18,6 +19,8 @@
 
 from mobius._configs import BaseModelConfig
 
+logger = logging.getLogger(__name__)
+
 _FUNCTIONS_DOMAIN = "com.microsoft"
 
 # Cache state pair: (key, value) or (conv_state, ssm_state) for stateful
@@ -142,18 +145,35 @@ def _register_kv_cache_outputs(
 
     When the parameters are omitted, output shapes/dtypes are left to the
     shape inference pass that runs during model optimization.
+
+    The present-shape parameters are all-or-nothing by design: pass every one
+    to stamp the explicit type, or none to opt out and infer. A *partial* set
+    is almost always a wiring slip (a caller wired some dims but dropped
+    others), so it is treated conservatively -- the stamp is skipped, the
+    shapes fall back to the known-wrong inference path, and a warning is
+    logged naming the missing parameters so the slip is loud rather than silent.
     """
-    stamp = all(
-        param is not None
-        for param in (
-            batch,
-            num_kv_heads,
-            key_head_dim,
-            value_head_dim,
-            total_seq_len,
-            dtype,
+    params = {
+        "batch": batch,
+        "num_kv_heads": num_kv_heads,
+        "key_head_dim": key_head_dim,
+        "value_head_dim": value_head_dim,
+        "total_seq_len": total_seq_len,
+        "dtype": dtype,
+    }
+    provided = [name for name, value in params.items() if value is not None]
+    stamp = len(provided) == len(params)
+    if provided and not stamp:
+        missing = [name for name in params if params[name] is None]
+        logger.warning(
+            "_register_kv_cache_outputs received a partial set of present-shape "
+            "parameters (provided %s, missing %s); these are all-or-nothing, so "
+            "the explicit present.* stamp is SKIPPED and shapes fall back to "
+            "inference (which mis-derives head_dim for GroupQueryAttention). "
+            "Pass all six parameters to stamp, or none to opt out.",
+            provided,
+            missing,
         )
-    )
     for i, (present_key, present_value) in enumerate(present_key_values):
         if stamp:
             present_key.shape = ir.Shape([batch, num_kv_heads, total_seq_len, key_head_dim])
diff --git a/src/mobius/tasks/_cache_utils_test.py b/src/mobius/tasks/_cache_utils_test.py
index 065bda85..64924ead 100644
--- a/src/mobius/tasks/_cache_utils_test.py
+++ b/src/mobius/tasks/_cache_utils_test.py
@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+import logging
+
 import onnx_ir as ir
 
 from mobius.tasks._base import _make_graph
@@ -72,32 +74,40 @@ def test_distinct_key_value_head_dims(self):
         assert _dims(key)[1] == 16 and _dims(key)[3] == 192
         assert _dims(value)[1] == 16 and _dims(value)[3] == 128
 
-    def test_no_params_leaves_shapes_untouched(self):
+    def test_no_params_leaves_shapes_untouched(self, caplog):
         """Without shape params the helper must not stamp (inference path)."""
         _, builder = _make_graph()
         pairs = [_present_pair("present.0", wrong_head_dim=32)]
 
-        _register_kv_cache_outputs(builder, pairs)
+        with caplog.at_level(logging.WARNING, logger="mobius.tasks._cache_utils"):
+            _register_kv_cache_outputs(builder, pairs)
 
         key, _ = pairs[0]
         # Unchanged: still the (wrong) pre-existing inferred shape.
         assert _dims(key) == ["batch", 32, "seq", 32]
+        # Opting out (zero params) is intentional and must stay silent.
+        assert caplog.text == ""
 
-    def test_partial_params_do_not_stamp(self):
-        """All shape params are required; a partial set falls back to inference."""
+    def test_partial_params_do_not_stamp(self, caplog):
+        """A partial set falls back to inference AND warns about the slip."""
         _, builder = _make_graph()
         pairs = [_present_pair("present.0", wrong_head_dim=32)]
 
-        _register_kv_cache_outputs(
-            builder,
-            pairs,
-            batch=ir.SymbolicDim("batch"),
-            num_kv_heads=32,
-            # key_head_dim / value_head_dim / total_seq_len / dtype omitted
-        )
+        with caplog.at_level(logging.WARNING, logger="mobius.tasks._cache_utils"):
+            _register_kv_cache_outputs(
+                builder,
+                pairs,
+                batch=ir.SymbolicDim("batch"),
+                num_kv_heads=32,
+                # key_head_dim / value_head_dim / total_seq_len / dtype omitted
+            )
 
         key, _ = pairs[0]
         assert _dims(key) == ["batch", 32, "seq", 32]
+        assert "partial set of present-shape parameters" in caplog.text
+        # The warning must name the omitted parameters so the slip is diagnosable.
+        for missing in ("key_head_dim", "value_head_dim", "total_seq_len", "dtype"):
+            assert missing in caplog.text
 
     def test_registers_named_outputs(self):
         """Outputs are registered with the conventional present.{i}.* names."""

From 23564ff012fa2993f59169f40fa258e06aafa942 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Tue, 2 Jun 2026 21:56:57 +0000
Subject: [PATCH 29/38] style: apply lintrunner format + D205 fix for CI

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_dtype_utils_test.py |  3 +-
 src/mobius/components/_attention.py     | 20 +++------
 src/mobius/components/_common.py        |  4 +-
 src/mobius/components/_common_test.py   |  4 +-
 src/mobius/tasks/_task_test.py          |  4 +-
 tests/build_graph_test.py               |  7 +--
 tests/static_cache_decode_test.py       | 58 +++++++------------------
 7 files changed, 29 insertions(+), 71 deletions(-)

diff --git a/src/mobius/_passes/_dtype_utils_test.py b/src/mobius/_passes/_dtype_utils_test.py
index 9bbdf0c6..4267eb7e 100644
--- a/src/mobius/_passes/_dtype_utils_test.py
+++ b/src/mobius/_passes/_dtype_utils_test.py
@@ -39,7 +39,8 @@ def test_const_value_wins_on_disagreement(self, caplog):
         with caplog.at_level(logging.WARNING, logger="mobius._passes._dtype_utils"):
             assert initializer_dtype(v) == ir.DataType.FLOAT16
         assert any(
-            record.levelno == logging.WARNING and "stale type metadata" in record.getMessage().lower()
+            record.levelno == logging.WARNING
+            and "stale type metadata" in record.getMessage().lower()
             for record in caplog.records
         ), "expected a warning when declared dtype disagrees with const_value"
 
diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py
index 801b8155..4073991b 100644
--- a/src/mobius/components/_attention.py
+++ b/src/mobius/components/_attention.py
@@ -66,20 +66,14 @@ def _attend_over_static_cache(
     Returns:
         The attention output for the active phase, shape ``[B, S_q, hidden]``.
     """
-    seq_len = op.Squeeze(
-        op.Shape(query, start=1, end=2), op.Constant(value_ints=[0])
-    )
+    seq_len = op.Squeeze(op.Shape(query, start=1, end=2), op.Constant(value_ints=[0]))
     is_multi_token_step = op.Greater(seq_len, op.Constant(value_int=1))
 
     def _build_attention_branch(name: str, use_causal_mask: bool) -> ir.Graph:
-        branch = ir.Graph(
-            [], [], nodes=[], name=name, opset_imports={"": OPSET_VERSION}
-        )
+        branch = ir.Graph([], [], nodes=[], name=name, opset_imports={"": OPSET_VERSION})
         branch_op = GraphBuilder(branch).op
         attn_mask = (
-            create_static_cache_causal_mask(
-                branch_op, query, key_cache, write_indices
-            )
+            create_static_cache_causal_mask(branch_op, query, key_cache, write_indices)
             if use_causal_mask
             else None
         )
@@ -122,12 +116,8 @@ def _build_attention_branch(name: str, use_causal_mask: bool) -> ir.Graph:
         branch.outputs.append(attn_output)
         return branch
 
-    prefill_branch = _build_attention_branch(
-        "static_cache_prefill", use_causal_mask=True
-    )
-    decode_branch = _build_attention_branch(
-        "static_cache_decode", use_causal_mask=False
-    )
+    prefill_branch = _build_attention_branch("static_cache_prefill", use_causal_mask=True)
+    decode_branch = _build_attention_branch("static_cache_decode", use_causal_mask=False)
     return op.If(
         is_multi_token_step,
         then_branch=prefill_branch,
diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py
index 179c13b0..ef16c5c5 100644
--- a/src/mobius/components/_common.py
+++ b/src/mobius/components/_common.py
@@ -381,9 +381,7 @@ def create_static_cache_causal_mask(
 
     # Scalar S_q (query length) and total KV length (max_seq) for Range.
     q_len = op.Squeeze(op.Shape(query, start=1, end=2), op.Constant(value_ints=[0]))
-    total_kv = op.Squeeze(
-        op.Shape(key_cache, start=1, end=2), op.Constant(value_ints=[0])
-    )
+    total_kv = op.Squeeze(op.Shape(key_cache, start=1, end=2), op.Constant(value_ints=[0]))
 
     # Per-step query offsets 0..S_q-1 and key slot indices 0..max_seq-1.
     q_offsets = op.Range(zero, q_len, one)  # [S_q] int64
diff --git a/src/mobius/components/_common_test.py b/src/mobius/components/_common_test.py
index 93bc2c4b..05d1dad1 100644
--- a/src/mobius/components/_common_test.py
+++ b/src/mobius/components/_common_test.py
@@ -30,9 +30,7 @@ def _run_static_cache_mask(max_seq: int, query_len: int, write_index: int) -> np
     key_cache = create_test_input(
         builder, "key_cache", [1, max_seq, 16], dtype=ir.DataType.FLOAT
     )
-    write_indices = create_test_input(
-        builder, "write_indices", [1], dtype=ir.DataType.INT64
-    )
+    write_indices = create_test_input(builder, "write_indices", [1], dtype=ir.DataType.INT64)
     mask = create_static_cache_causal_mask(op, query, key_cache, write_indices)
     mask.name = "mask"
     graph.outputs.append(mask)
diff --git a/src/mobius/tasks/_task_test.py b/src/mobius/tasks/_task_test.py
index 875618a9..ef81e1c6 100644
--- a/src/mobius/tasks/_task_test.py
+++ b/src/mobius/tasks/_task_test.py
@@ -92,7 +92,9 @@ def test_build_outputs(self):
         assert "present.0.value" in output_names
 
     def test_present_outputs_match_past_inputs(self):
-        """present.{i}.* must declare the same kv_heads/head_dim/dtype as the
+        """Present outputs must match their past-input KV metadata.
+
+        present.{i}.* must declare the same kv_heads/head_dim/dtype as the
         corresponding past_key_values.{i}.* inputs, with total (past+current)
         sequence length. Guards the GQA present-head_dim export bug, where the
         contrib-op shape inference would otherwise mis-declare head_dim.
diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index 3ce03ce4..57ec5f34 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -4441,9 +4441,7 @@ def test_static_cache_has_tensorscatter_and_attention(self):
         op_types = {n.op_type for n in self._walk_nodes(model.graph)}
         assert "TensorScatter" in op_types, "Static cache graph should use TensorScatter"
         assert "Attention" in op_types, "Static cache graph should use Attention"
-        assert "If" in op_types, (
-            "Static cache attention should be phase-split behind an If"
-        )
+        assert "If" in op_types, "Static cache attention should be phase-split behind an If"
 
     def test_static_cache_has_initializers(self):
         """Verify the graph has model parameters."""
@@ -4547,8 +4545,7 @@ def _single_attention(graph):
         """Return the sole ``Attention`` node in an If branch subgraph."""
         attention_nodes = [n for n in graph if n.op_type == "Attention"]
         assert len(attention_nodes) == 1, (
-            f"expected exactly one Attention per If branch, "
-            f"got {len(attention_nodes)}"
+            f"expected exactly one Attention per If branch, got {len(attention_nodes)}"
         )
         return attention_nodes[0]
 
diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py
index 2654b662..7312e89a 100644
--- a/tests/static_cache_decode_test.py
+++ b/tests/static_cache_decode_test.py
@@ -129,9 +129,7 @@ def _build_static_cache_session(
     config = dataclasses.replace(config, dtype=ir_dtype)
     module = registry.get(_MODEL_TYPE)(config)
     task = CausalLMTask(static_cache=True, max_seq_len=_MAX_SEQ_LEN)
-    package = build_from_module(
-        module, config, task=task, execution_provider="default"
-    )
+    package = build_from_module(module, config, task=task, execution_provider="default")
     model = package["model"]
     _fill_random_weights(model, np.random.default_rng(0))
 
@@ -164,9 +162,7 @@ def _empty_caches(
     return feeds
 
 
-def _carry_caches(
-    outputs: dict[str, np.ndarray], num_layers: int
-) -> dict[str, np.ndarray]:
+def _carry_caches(outputs: dict[str, np.ndarray], num_layers: int) -> dict[str, np.ndarray]:
     """Feed the prefill ``updated_*`` caches back in as decode inputs."""
     feeds: dict[str, np.ndarray] = {}
     for layer in range(num_layers):
@@ -194,18 +190,14 @@ def test_static_cache_prefill_and_decode_run_on_cuda():
         # --- Prefill: write N tokens from slot 0 (S_q = N != max_seq). ---
         prefill_len = 4
         prefill_feeds: dict[str, np.ndarray] = {
-            "input_ids": rng.integers(
-                0, vocab, size=(1, prefill_len), dtype=np.int64
-            ),
+            "input_ids": rng.integers(0, vocab, size=(1, prefill_len), dtype=np.int64),
             "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
             "write_indices": np.array([0], dtype=np.int64),
             "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
         }
         prefill_feeds.update(_empty_caches(num_layers, kv_hidden))
 
-        prefill_out = dict(
-            zip(output_names, session.run(output_names, prefill_feeds))
-        )
+        prefill_out = dict(zip(output_names, session.run(output_names, prefill_feeds)))
         prefill_logits = prefill_out["logits"]
         assert prefill_logits.shape == (1, prefill_len, vocab)
         assert np.isfinite(prefill_logits).all(), "prefill logits must be finite"
@@ -219,9 +211,7 @@ def test_static_cache_prefill_and_decode_run_on_cuda():
         }
         decode_feeds.update(_carry_caches(prefill_out, num_layers))
 
-        decode_out = dict(
-            zip(output_names, session.run(output_names, decode_feeds))
-        )
+        decode_out = dict(zip(output_names, session.run(output_names, decode_feeds)))
         decode_logits = decode_out["logits"]
         assert decode_logits.shape == (1, 1, vocab)
         assert np.isfinite(decode_logits).all(), "decode logits must be finite"
@@ -255,9 +245,7 @@ def _executed_attention_events(profile_path: str) -> list[dict]:
         if args.get("op_name") != "Attention":
             continue
         input_shapes = args.get("input_type_shape", [])
-        has_mask = any(
-            len(next(iter(shape.values()))) == 4 for shape in input_shapes
-        )
+        has_mask = any(len(next(iter(shape.values()))) == 4 for shape in input_shapes)
         attention_events.append({"name": event["name"], "has_mask": has_mask})
     return attention_events
 
@@ -303,9 +291,7 @@ def test_static_cache_decode_runs_maskless_on_cuda():
         # Multi-token prefill (S_q = 4): the If must take the masked branch.
         prefill_len = 4
         prefill_feeds: dict[str, np.ndarray] = {
-            "input_ids": rng.integers(
-                0, vocab, size=(1, prefill_len), dtype=np.int64
-            ),
+            "input_ids": rng.integers(0, vocab, size=(1, prefill_len), dtype=np.int64),
             "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
             "write_indices": np.array([0], dtype=np.int64),
             "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
@@ -320,8 +306,7 @@ def test_static_cache_decode_runs_maskless_on_cuda():
 
     # Decode took the maskless branch on every layer (Flash-eligible).
     assert len(decode_events) == num_layers, (
-        f"expected {num_layers} decode-branch Attention executions, "
-        f"got {len(decode_events)}"
+        f"expected {num_layers} decode-branch Attention executions, got {len(decode_events)}"
     )
     assert all(not e["has_mask"] for e in decode_events), (
         "decode-branch Attention must run WITHOUT an attn_mask input so ORT "
@@ -331,8 +316,7 @@ def test_static_cache_decode_runs_maskless_on_cuda():
 
     # Prefill took the masked branch on every layer (memory-efficient path).
     assert len(prefill_events) == num_layers, (
-        f"expected {num_layers} prefill-branch Attention executions, "
-        f"got {len(prefill_events)}"
+        f"expected {num_layers} prefill-branch Attention executions, got {len(prefill_events)}"
     )
     assert all(e["has_mask"] for e in prefill_events), (
         "prefill-branch Attention must carry the explicit causal mask"
@@ -477,9 +461,7 @@ def test_static_cache_decode_selects_flash_kernel_on_cuda(config_overrides):
 
         prefill_len = 4
         prefill_feeds: dict[str, np.ndarray] = {
-            "input_ids": rng.integers(
-                0, vocab, size=(1, prefill_len), dtype=np.int64
-            ),
+            "input_ids": rng.integers(0, vocab, size=(1, prefill_len), dtype=np.int64),
             "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
             "write_indices": np.array([0], dtype=np.int64),
             "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
@@ -530,9 +512,7 @@ def test_static_cache_decode_ignores_keys_beyond_nonpad_on_cuda():
     the logits would diverge.
     """
     with tempfile.TemporaryDirectory() as tmp_dir:
-        session, config = _build_static_cache_session(
-            tmp_dir, ir_dtype=ir.DataType.FLOAT16
-        )
+        session, config = _build_static_cache_session(tmp_dir, ir_dtype=ir.DataType.FLOAT16)
         num_layers = config.num_hidden_layers
         kv_hidden = config.num_key_value_heads * config.head_dim
         vocab = config.vocab_size
@@ -542,17 +522,13 @@ def test_static_cache_decode_ignores_keys_beyond_nonpad_on_cuda():
         # Prefill four real tokens into slots 0..3 to populate the cache.
         prefill_len = 4
         prefill_feeds: dict[str, np.ndarray] = {
-            "input_ids": rng.integers(
-                0, vocab, size=(1, prefill_len), dtype=np.int64
-            ),
+            "input_ids": rng.integers(0, vocab, size=(1, prefill_len), dtype=np.int64),
             "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :],
             "write_indices": np.array([0], dtype=np.int64),
             "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64),
         }
         prefill_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16))
-        prefill_out = dict(
-            zip(output_names, session.run(output_names, prefill_feeds))
-        )
+        prefill_out = dict(zip(output_names, session.run(output_names, prefill_feeds)))
 
         # Decode one token into slot 4; valid keys are slots 0..4 (nonpad=5).
         nonpad = prefill_len + 1
@@ -567,9 +543,7 @@ def test_static_cache_decode_ignores_keys_beyond_nonpad_on_cuda():
             **decode_inputs,
             **_carry_caches(prefill_out, num_layers),
         }
-        baseline = dict(
-            zip(output_names, session.run(output_names, clean_feeds))
-        )
+        baseline = dict(zip(output_names, session.run(output_names, clean_feeds)))
 
         # Poison every cache slot at or beyond ``nonpad`` with large garbage;
         # those positions must never be attended during decode.  Slot 4 (the
@@ -581,9 +555,7 @@ def test_static_cache_decode_ignores_keys_beyond_nonpad_on_cuda():
                 buf[:, nonpad:, :] = np.float16(50.0)
                 poisoned_caches[name] = buf
         poisoned_feeds = {**decode_inputs, **poisoned_caches}
-        perturbed = dict(
-            zip(output_names, session.run(output_names, poisoned_feeds))
-        )
+        perturbed = dict(zip(output_names, session.run(output_names, poisoned_feeds)))
 
     assert np.array_equal(baseline["logits"], perturbed["logits"]), (
         "decode logits changed when cache slots beyond nonpad_kv_seqlen were "

From 96ef1b150e18cd0536513e665860e0baee6a5e14 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Wed, 3 Jun 2026 17:37:13 +0000
Subject: [PATCH 30/38] test: use IR-native ir.save instead of onnx.save in
 _fold_concat_test

Resolves the Copilot review finding flagging a CONTRIBUTING.md
"zero protobuf operations" violation in a test file. Replace the
onnx.save(ir.to_proto(model), ...) call with the IR-native
ir.save(model, model_path) pattern already used elsewhere in the
same file, and drop the now-unused `import onnx`. Behavior is
unchanged: the test still writes the model and loads it in ORT,
asserting the fp16 result shape/dtype and the fp16/fp32 MatMul
type-mismatch regression guard.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_concat_test.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py
index 2c5b309e..3cb237fa 100644
--- a/src/mobius/_passes/_fold_concat_test.py
+++ b/src/mobius/_passes/_fold_concat_test.py
@@ -633,7 +633,6 @@ class TestFoldConcatOrtLoad:
     """
 
     def test_folded_fp16_concat_matmul_loads_in_ort(self, tmp_path):
-        import onnx
         import onnxruntime as ort
 
         # Three fp16 weights with UNSET declared dtype (as after _cast_module_dtype),
@@ -684,7 +683,7 @@ def test_folded_fp16_concat_matmul_loads_in_ort(self, tmp_path):
             del model.graph.initializers[dead]
 
         model_path = tmp_path / "fp16_qkv.onnx"
-        onnx.save(ir.to_proto(model), str(model_path))
+        ir.save(model, model_path)
 
         # Before the fix this raises a fp16/fp32 MatMul type-mismatch at load.
         sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])

From 6fc26441d580e6b8f11e8cf9f21168f4eb6c870a Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Wed, 3 Jun 2026 18:44:53 +0000
Subject: [PATCH 31/38] test: use IR-native graph.all_nodes() and .as_graph()
 in build_graph_test
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses review feedback: replace the hand-rolled recursive _walk_nodes helper with onnx_ir Graph.all_nodes() (coverage-equivalent — recurses into If subgraphs), and use the typed .as_graph() accessor instead of .value for the If then/else branches. Bump onnx_ir floor to >=0.1.2 (all_nodes added in 0.1.2). Node set asserted by the tests is unchanged; TestBuildStaticCacheGraph 11/11 pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 pyproject.toml            |  2 +-
 tests/build_graph_test.py | 33 +++++++--------------------------
 2 files changed, 8 insertions(+), 27 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index fdcf48d2..0ccb5f9b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ license = "MIT"
 dependencies = [
     "huggingface_hub",
     "numpy>=1.24.0",
-    "onnx_ir>=0.1.0",
+    "onnx_ir>=0.1.2",
     "onnx-shape-inference>=0.1.9",
     "onnxscript>=0.7.0",
     "safetensors",
diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index b4112a4d..e6c5f7bc 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -4456,23 +4456,6 @@ class TestBuildStaticCacheGraph:
 
     MAX_SEQ_LEN = 128
 
-    @staticmethod
-    def _walk_nodes(graph):
-        """Yield every node in ``graph``, recursing into subgraphs (e.g. If).
-
-        The static-cache attention is emitted as an ``If`` whose branches
-        each contain an ``Attention`` op, so structural assertions must look
-        inside subgraph attributes rather than only the top-level graph.
-        """
-        for node in graph:
-            yield node
-            for attr in node.attributes.values():
-                if attr.type == ir.AttributeType.GRAPH and attr.value is not None:
-                    yield from TestBuildStaticCacheGraph._walk_nodes(attr.value)
-                elif attr.type == ir.AttributeType.GRAPHS:
-                    for subgraph in attr.value:
-                        yield from TestBuildStaticCacheGraph._walk_nodes(subgraph)
-
     def _build_static_cache_model(self, model_type: str = "qwen2", **config_overrides):
         """Build a model with CausalLMTask(static_cache=True) and return (model, config)."""
         from mobius.tasks import CausalLMTask
@@ -4556,7 +4539,7 @@ def test_static_cache_has_tensorscatter_and_attention(self):
         """
         model, _ = self._build_static_cache_model()
 
-        op_types = {n.op_type for n in self._walk_nodes(model.graph)}
+        op_types = {n.op_type for n in model.graph.all_nodes()}
         assert "TensorScatter" in op_types, "Static cache graph should use TensorScatter"
         assert "Attention" in op_types, "Static cache graph should use Attention"
         assert "If" in op_types, "Static cache attention should be phase-split behind an If"
@@ -4588,9 +4571,7 @@ def test_static_cache_attention_is_causal(self):
         """
         model, config = self._build_static_cache_model()
 
-        attention_nodes = [
-            n for n in self._walk_nodes(model.graph) if n.op_type == "Attention"
-        ]
+        attention_nodes = [n for n in model.graph.all_nodes() if n.op_type == "Attention"]
         # Two branches (prefill + decode) per layer.
         assert len(attention_nodes) == 2 * config.num_hidden_layers
 
@@ -4624,7 +4605,7 @@ def test_static_cache_phase_split_mask_presence(self):
         """
         model, config = self._build_static_cache_model()
 
-        if_nodes = [n for n in self._walk_nodes(model.graph) if n.op_type == "If"]
+        if_nodes = [n for n in model.graph.all_nodes() if n.op_type == "If"]
         # (a) Fail-closed: the phase-split If must exist, one per layer.
         assert len(if_nodes) == config.num_hidden_layers, (
             f"static-cache attention must phase-split via If "
@@ -4633,8 +4614,8 @@ def test_static_cache_phase_split_mask_presence(self):
         )
 
         for if_node in if_nodes:
-            then_branch = if_node.attributes["then_branch"].value
-            else_branch = if_node.attributes["else_branch"].value
+            then_branch = if_node.attributes["then_branch"].as_graph()
+            else_branch = if_node.attributes["else_branch"].as_graph()
             then_attn = self._single_attention(then_branch)
             else_attn = self._single_attention(else_branch)
 
@@ -4677,7 +4658,7 @@ def test_static_cache_is_phase_split_behind_if(self):
         """
         model, config = self._build_static_cache_model()
         op_counts: dict[str, int] = {}
-        for node in self._walk_nodes(model.graph):
+        for node in model.graph.all_nodes():
             op_counts[node.op_type] = op_counts.get(node.op_type, 0) + 1
 
         assert op_counts.get("If", 0) == config.num_hidden_layers, (
@@ -4713,7 +4694,7 @@ def test_static_cache_moe_graph_builds(self):
 
         # Verify TensorScatter (top level) and Attention (inside the
         # phase-split If branches) ops are present.
-        op_types = {n.op_type for n in self._walk_nodes(model.graph)}
+        op_types = {n.op_type for n in model.graph.all_nodes()}
         assert "TensorScatter" in op_types
         assert "Attention" in op_types
 

From a7cbf929f88197cb19d10c5ab66de9b41e434fd8 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Wed, 3 Jun 2026 19:03:49 +0000
Subject: [PATCH 32/38] ci: run static-cache decode runtime tests in
 integration-fast GPU job

The runtime-CUDA decode tests (tests/static_cache_decode_test.py) verify the is_causal=0 phase-split actually runs on GPU (decode stays maskless->Flash, no NOT_IMPLEMENTED). They previously ran in NO CI job: the CPU unit job collects them but they skipif-skip without a CUDAExecutionProvider, and no GPU job invoked them. Add an explicit pytest step to the existing per-PR A10 'Integration (fast)' job (a marker alone would not collect them: that job's step is path-scoped to integration_test.py with a -k allowlist). Closes the runtime-CUDA coverage gap for the static-cache export path. Refs #329 for the remaining numerical-parity coverage follow-up.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .github/workflows/main.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index e1b7e259..70e35c7a 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -292,6 +292,10 @@ jobs:
             -k "smollm-135m or albert-base or t5-small or dinov2-small or wav2vec2-base or whisper-tiny or test_gemma3_multimodal or test_qwen35 or test_qwen3_next or test_deepseek or test_sam_vit or test_ocr2" \
             --cov=src --cov-report=xml --cov-branch --junitxml junit.xml
         timeout-minutes: 15
+      - name: Static-cache decode runtime tests
+        run: |
+          pytest tests/static_cache_decode_test.py -v --tb=short
+        timeout-minutes: 10
       - name: Upload coverage to Codecov
         if: always()
         uses: codecov/codecov-action@v6

From 945e5f50b7fb90bd5d26fd570f49c39ae2af5daa Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Wed, 3 Jun 2026 23:11:30 +0000
Subject: [PATCH 33/38] ci(benchmark): waive intended static-cache phase-split
 node delta + changelog

The Compare-results gate flagged a Blocker regression because static-cache
num_nodes rose (llama/qwen2 58->68, phi3 56->66, +10 each). Those +10 nodes
are the INTENDED per-layer If(Greater(seq_len,1)) + prefill-mask phase-split
introduced by this PR, not a regression. Pin the exact (baseline,current)
node counts in benchmark_compare.EXPECTED_CHANGES so only this precise,
self-cleaning transition is waived (post-merge base=68 no longer matches;
a future 68->78 still blocks). Add regression + display-key-binding guards.
Also add CHANGELOG entries for the phase-split export and the GQA present-KV
shape fix (cf6c5c4).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 CHANGELOG.md                    |  35 +++++++
 tests/benchmark_compare.py      |  40 +++++++-
 tests/benchmark_compare_test.py | 157 ++++++++++++++++++++++++++++++++
 3 files changed, 230 insertions(+), 2 deletions(-)
 create mode 100644 tests/benchmark_compare_test.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dee0e635..a071b045 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,41 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Static-cache Attention Phase-split Export
+
+#### Added
+
+- Causal-LM static-cache exports now emit ONNX `Attention` with `is_causal=0`
+  wrapped in a per-layer phase-split `If(Greater(seq_len, 1))` subgraph: prefill
+  takes an explicit-causal-mask branch while single-token decode takes a maskless
+  branch (which lets onnxruntime route decode to the Flash kernel). This is a
+  user-visible exported-graph change — it adds a fixed +10 top-level nodes per
+  static-cache export (`llama` 58→68, `qwen2` 58→68, `phi3` 56→66), which the
+  benchmark regression comparator recognizes as an intended structural change
+  (see `tests/benchmark_compare.py`). Lands alongside the GQA present-KV shape
+  fix below as part of the same export-correctness PR.
+
+---
+
+### GQA Present KV-Cache Shape Fix
+
+#### Fixed
+
+- GroupQueryAttention exports now declare correct `present.{i}.key` /
+  `present.{i}.value` graph-output shapes and dtype. The GQA contrib op's shape
+  inference mis-derived the present KV `head_dim` (e.g. 32 instead of 96 on
+  `microsoft/Phi-3.5-mini-instruct`), so the present KV-cache outputs declared a
+  `head_dim` inconsistent with the (correct) `past_key_values` inputs. ORT logged
+  `Error merging shape info ... lenient merge` (64 warnings on Phi-3.5) and any
+  consumer that chains `present` → `past` and trusts declared shapes (e.g.
+  `onnxruntime-genai`) saw mismatched past-vs-present KV cache types. This is a
+  metadata / declared-shape correction only — runtime numerics are unchanged
+  (weights byte-identical, next-token parity 20/20). `_register_kv_cache_outputs`
+  now stamps the present KV outputs symmetric to the past inputs. Affects
+  GQA-fusion packed-QKV exports (Phi-3.5, Llama-3.2, Qwen2, Mistral, Phi-3-GQA).
+
+---
+
 ### fp16 GQA Export Fix
 
 #### Fixed
diff --git a/tests/benchmark_compare.py b/tests/benchmark_compare.py
index c9b9fcf9..5d8a2e1b 100644
--- a/tests/benchmark_compare.py
+++ b/tests/benchmark_compare.py
@@ -33,6 +33,30 @@
     "num_nodes": (0.05, 0.10),
 }
 
+# Intended structural changes that would otherwise trip a deterministic-metric
+# blocker, keyed by model display key -> metric -> the exact (baseline, current)
+# values to waive. ABSOLUTE counts (not a relative delta) are pinned on purpose:
+# the waiver fires only for this exact base->current transition, so once the
+# change merges into the base branch (baseline becomes the new value) the entry
+# can NEVER match again and is truly self-cleaning — a later regression that
+# happens to add the same number of nodes still blocks. Only the exact pinned
+# transition is waived; any other base/current pair falls through to the normal
+# threshold logic and can block, so this never masks an accidental regression.
+#
+# PR #328 (static-cache export): the is_causal=0 static-cache path wraps each
+# decoder layer's attention in a phase-split ``If(Greater(seq_len, 1))`` subgraph
+# (prefill takes the explicit-causal-mask branch, decode the maskless branch)
+# plus the prefill mask-build nodes. This intentionally raises num_nodes by a
+# fixed +10 top-level nodes per static-cache export. These are correct, not a
+# regression. Safe to delete this whole table after PR #328 merges. (If an
+# unrelated opset/ORT/exporter change shifts the baseline counts before merge,
+# update the pinned values here — a mismatch fails closed with a RED blocker.)
+EXPECTED_CHANGES: dict[str, dict[str, tuple[int, int]]] = {
+    "llama (static-cache)": {"num_nodes": (58, 68)},
+    "qwen2 (static-cache)": {"num_nodes": (58, 68)},
+    "phi3 (static-cache)": {"num_nodes": (56, 66)},
+}
+
 _GITHUB_REPO_URL = "https://github.com/onnxruntime/mobius"
 
 
@@ -58,7 +82,12 @@ def compare(current_path: str, baseline_path: str) -> tuple[str, bool]:
                 continue
             delta_pct = (curr_val - base_val) / base_val
             warn_t, block_t = THRESHOLDS[metric]
-            if delta_pct > block_t:
+            expected = EXPECTED_CHANGES.get(model, {}).get(metric)
+            if expected is not None and (base_val, curr_val) == expected:
+                # Intended structural change (e.g. the static-cache phase-split
+                # If+mask). Waive only this exact pinned base->current transition.
+                status = "\U0001f7e6"  # blue square: accepted intended change
+            elif delta_pct > block_t:
                 status = "\U0001f534"  # red circle
                 has_blocker = True
             elif delta_pct > warn_t:
@@ -105,7 +134,14 @@ def _sha_link(sha: str) -> str:
     elif any(r[5] == "\u26a0\ufe0f" for r in rows):
         md += "\n> Warning: minor regressions detected. Review flagged metrics.\n"
     else:
-        md += "\n> No performance regressions.\n"
+        md += "\n> No blocking regressions.\n"
+
+    if any(r[5] == "\U0001f7e6" for r in rows):
+        md += (
+            "\n> \U0001f7e6 = intended structural change accepted via "
+            "`EXPECTED_CHANGES` (exact pinned base→current values; see "
+            "`tests/benchmark_compare.py`).\n"
+        )
 
     return md, has_blocker
 
diff --git a/tests/benchmark_compare_test.py b/tests/benchmark_compare_test.py
new file mode 100644
index 00000000..d0ff7b32
--- /dev/null
+++ b/tests/benchmark_compare_test.py
@@ -0,0 +1,157 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for the deterministic-metric regression comparator.
+
+Focuses on the ``EXPECTED_CHANGES`` allowlist used to accept intended
+structural node-count changes (e.g. the PR #328 static-cache phase-split)
+without masking accidental regressions.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+import tests.benchmark_compare as bc
+
+# A blue square marks an accepted intended structural change.
+_ACCEPTED = "\U0001f7e6"
+# A red circle marks a blocking regression.
+_BLOCKER = "\U0001f534"
+
+
+def _write(path: Path, models: dict[str, dict[str, int]]) -> str:
+    path.write_text(json.dumps({"_metadata": {"commit": "deadbeef"}, "models": models}))
+    return str(path)
+
+
+def _run(tmp_path: Path, base: dict, curr: dict) -> tuple[str, bool]:
+    baseline = _write(tmp_path / "baseline.json", base)
+    current = _write(tmp_path / "current.json", curr)
+    return bc.compare(current, baseline)
+
+
+def test_exact_expected_static_cache_delta_is_waived(tmp_path: Path):
+    """The exact intended +10 static-cache node delta must not block."""
+    md, has_blocker = _run(
+        tmp_path,
+        base={"llama (static-cache)": {"num_nodes": 58, "model_size_bytes": 1000}},
+        curr={"llama (static-cache)": {"num_nodes": 68, "model_size_bytes": 1000}},
+    )
+    assert has_blocker is False
+    assert _ACCEPTED in md
+    assert _BLOCKER not in md
+
+
+@pytest.mark.parametrize(
+    "model",
+    ["llama (static-cache)", "qwen2 (static-cache)", "phi3 (static-cache)"],
+)
+def test_all_allowlisted_models_waive_their_change(tmp_path: Path, model: str):
+    base_nodes, curr_nodes = bc.EXPECTED_CHANGES[model]["num_nodes"]
+    md, has_blocker = _run(
+        tmp_path,
+        base={model: {"num_nodes": base_nodes}},
+        curr={model: {"num_nodes": curr_nodes}},
+    )
+    assert has_blocker is False
+    assert _ACCEPTED in md
+
+
+def test_delta_larger_than_expected_still_blocks(tmp_path: Path):
+    """An extra unexpected node beyond the intended delta must still block."""
+    md, has_blocker = _run(
+        tmp_path,
+        base={"llama (static-cache)": {"num_nodes": 58}},
+        curr={"llama (static-cache)": {"num_nodes": 69}},  # +11, not the intended +10
+    )
+    assert has_blocker is True
+    assert _BLOCKER in md
+
+
+def test_non_allowlisted_model_still_blocks(tmp_path: Path):
+    """A model outside the allowlist gets no waiver."""
+    md, has_blocker = _run(
+        tmp_path,
+        base={"llama": {"num_nodes": 58}},
+        curr={"llama": {"num_nodes": 68}},  # +17%, over the 10% block threshold
+    )
+    assert has_blocker is True
+    assert _BLOCKER in md
+
+
+def test_smaller_than_expected_delta_is_not_waived(tmp_path: Path):
+    """Only the exact delta is waived; a smaller change uses normal thresholds."""
+    # +5 on 58 = +8.6%, between warn (5%) and block (10%): a warning, not waived.
+    md, has_blocker = _run(
+        tmp_path,
+        base={"llama (static-cache)": {"num_nodes": 58}},
+        curr={"llama (static-cache)": {"num_nodes": 63}},
+    )
+    assert has_blocker is False
+    assert _ACCEPTED not in md
+    assert "\u26a0\ufe0f" in md  # warning
+
+
+def test_post_merge_zero_delta_is_inert(tmp_path: Path):
+    """Once merged, base==head -> 0 delta: no blocker, no accepted-marker."""
+    md, has_blocker = _run(
+        tmp_path,
+        base={"llama (static-cache)": {"num_nodes": 68}},
+        curr={"llama (static-cache)": {"num_nodes": 68}},
+    )
+    assert has_blocker is False
+    assert _ACCEPTED not in md
+    assert _BLOCKER not in md
+
+
+def test_post_merge_repeat_delta_still_blocks(tmp_path: Path):
+    """Absolute pinning closes the stale-waiver hole.
+
+    After the +10 phase-split merges (baseline becomes 68), a *future* +10
+    regression (68 -> 78) must NOT match the (58, 68) waiver and must still block.
+    """
+    md, has_blocker = _run(
+        tmp_path,
+        base={"llama (static-cache)": {"num_nodes": 68}},
+        curr={"llama (static-cache)": {"num_nodes": 78}},
+    )
+    assert has_blocker is True
+    assert _BLOCKER in md
+    assert _ACCEPTED not in md
+
+
+def test_non_waived_metric_on_allowlisted_model_still_blocks(tmp_path: Path):
+    """A non-waived metric on an allowlisted model still blocks.
+
+    The waiver is scoped per-metric: an allowlisted model whose num_nodes
+    matches the pinned (58, 68) must still block on a DIFFERENT metric
+    (model_size_bytes +30%), guarding against over-broadening the waiver.
+    """
+    md, has_blocker = _run(
+        tmp_path,
+        base={"llama (static-cache)": {"num_nodes": 58, "model_size_bytes": 1000}},
+        curr={"llama (static-cache)": {"num_nodes": 68, "model_size_bytes": 1300}},
+    )
+    assert has_blocker is True
+    assert _BLOCKER in md
+    assert _ACCEPTED in md  # the num_nodes row is still waived
+
+
+def test_expected_changes_keys_are_real_model_display_keys():
+    """Guard against a silent waiver no-op if model keys drift.
+
+    A model_type/task rename would move the comparator's allowlist keys away
+    from the benchmarked display keys, silently re-REDing the intended +10.
+    """
+    from tests.benchmark_build import BENCHMARK_MODELS, _display_key
+
+    valid_keys = {_display_key(e.model_type, e.task_name) for e in BENCHMARK_MODELS}
+    for model_key in bc.EXPECTED_CHANGES:
+        assert model_key in valid_keys, (
+            f"EXPECTED_CHANGES key {model_key!r} is not a benchmarked model "
+            f"display key; the waiver would silently never fire."
+        )

From 110f26b52b20026a4281b130651d8c75d69a21e8 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Wed, 3 Jun 2026 23:31:51 +0000
Subject: [PATCH 34/38] test(fp16): add e2e regression test guarding df203cc
 fold-pass dtype retention
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds an end-to-end regression test (src/mobius/_passes/_fold_dtype_e2e_test.py)
that drives the real fp16 export path (build_from_module + apply_weights) and
asserts packed/transposed initializers keep FLOAT16 through the fold passes,
guarding the df203cc fix at the export level (the existing unit/pass coverage
only exercises hand-built single-pass graphs).

Guards BOTH df203cc mechanisms:
  * FoldConcat/FoldTranspose output-type stamping — the realistic fp16 GQA
    PackQKV export (MatMul(hidden, Transpose(Concat(W_q,W_k,W_v)))) whose Concat
    output carries no declared dtype.
  * initializer_dtype() const_value fallback — reproduced by dropping the
    declared type on the packed-QKV Concat inputs so the fallback is the only
    thing keeping the folded weights fp16.

Includes a serialize->reload-with-external-data round-trip (ir.save + ir.load,
model.onnx + model.onnx.data) asserting the reloaded weights are FLOAT16 with
bytes intact — the ground-truth check for the serialize-time fp16-under-fp32
corruption that an in-memory const_value.numpy() can miss.

3-way revert proof: HEAD/fix -> all pass; full df203cc^ revert -> all fail;
fallback-only revert (initializer_dtype call-sites, type-stamp kept) -> only the
dropped-declared-dtype test fails (pinning the const_value fallback specifically).
Fully synthetic (no HF download, no GPU, no ORT execution) to fit the per-PR CI tier.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_dtype_e2e_test.py | 352 +++++++++++++++++++++
 1 file changed, 352 insertions(+)
 create mode 100644 src/mobius/_passes/_fold_dtype_e2e_test.py

diff --git a/src/mobius/_passes/_fold_dtype_e2e_test.py b/src/mobius/_passes/_fold_dtype_e2e_test.py
new file mode 100644
index 00000000..da25ab7c
--- /dev/null
+++ b/src/mobius/_passes/_fold_dtype_e2e_test.py
@@ -0,0 +1,352 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""End-to-end regression tests for fp16 dtype preservation through the fold passes.
+
+These complement the unit-level coverage in ``_dtype_utils_test.py`` and the
+pass-level coverage in ``_fold_concat_test.py`` / ``_fold_transpose_test.py`` by
+driving the *real* export pipeline (``build_from_module`` + ``apply_weights``)
+rather than hand-built single-pass graphs.
+
+The df203cc fix has two parts, and this module guards both at the export level:
+
+1. **dtype stamping (type-stamping path).** The fold passes stamp the *resolved*
+   fp16 dtype on the new packed / transposed initializer (its ``TensorType`` and
+   its ``LazyTensor``) instead of inheriting the unset type of the rewrite-produced
+   ``Concat`` intermediate or defaulting to ``FLOAT``. Guarded by
+   ``test_packed_qkv_weights_stay_fp16_after_export`` and
+   ``test_all_folded_matmul_weights_stay_fp16_after_export``: in the real fp16 GQA
+   export the PackQKV rewrite emits a ``MatMul(hidden, Transpose(Concat(W_q, W_k,
+   W_v)))`` whose ``Concat`` output carries no declared dtype, so without the
+   stamping the chained FoldConcat -> FoldTranspose widens the packed-QKV weight
+   to fp32.
+
+2. **const_value fallback.** ``_dtype_utils.initializer_dtype`` resolves the dtype
+   from ``const_value`` when an initializer's declared ``type`` was dropped during
+   graph building. Guarded by
+   ``test_dropped_declared_dtype_falls_back_to_fp16_const``, which reproduces that
+   dropped-dtype condition on the packed-QKV weights (as the original regression
+   exhibited) so the const_value fallback is the only thing keeping the folded
+   weights fp16.
+
+Either failure silently widens an fp16 model to fp32 packed weights, which
+onnxruntime rejects at load time with a fp16/fp32 ``MatMul`` type-mismatch
+(``Type parameter (T) of Optype (MatMul) bound to different types``).
+
+The in-memory dtype assertions read the IR ``const_value`` metadata. The
+regression's *ground-truth* symptom, however, is a serialize-time corruption:
+the fold passes write the fp16 *bytes* under a fp32 declaration, so the persisted
+initializer is a fp32-declared tensor backed by half-width data. An in-memory
+``const_value.numpy()`` can stay fp16 and miss this (only the declared dtype
+widens), so ``test_packed_qkv_survives_serialization_roundtrip`` and the
+serialize check in ``test_dropped_declared_dtype_falls_back_to_fp16_const`` drive
+the real export artifact (``ir.save`` with external data — the same
+``model.onnx`` + ``model.onnx.data`` layout the fp16 export produces), reload it,
+and assert the folded weights reload as fp16 with their bytes intact.
+
+The models are tiny and fully synthetic (no HuggingFace download, no GPU and no
+onnxruntime execution — only the ONNX graph is constructed, serialized and
+inspected), so the tests fit the per-PR CI tier.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import onnx_ir as ir
+import pytest
+import torch
+
+from mobius._builder import build_from_module
+from mobius._configs import ArchitectureConfig
+from mobius._optimizations import fold_initializers_after_weights
+from mobius._registry import registry
+from mobius._weight_loading import apply_weights
+
+
+def _make_fp16_llama_config() -> ArchitectureConfig:
+    """A tiny llama config whose Q/K/V weights are packable (no QK norm).
+
+    ``dtype=FLOAT16`` routes the build through the fp16 packed-QKV path that the
+    fold passes must preserve.
+    """
+    return ArchitectureConfig(
+        # Structural invariants that make the fp16 PackQKV fold fire:
+        #   num_attention_heads * head_dim == hidden_size  (4 * 16 == 64), and
+        #   num_key_value_heads < num_attention_heads       (2 < 4, i.e. GQA).
+        # The remaining fields are arbitrary small values kept lightweight.
+        hidden_size=64,
+        intermediate_size=128,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        num_hidden_layers=2,
+        vocab_size=256,
+        max_position_embeddings=128,
+        hidden_act="silu",
+        rms_norm_eps=1e-6,
+        rope_type="default",
+        rope_theta=10000.0,
+        pad_token_id=0,
+        dtype=ir.DataType.FLOAT16,
+    )
+
+
+def _build_fp16_decoder(config: ArchitectureConfig) -> ir.Model:
+    """Build an fp16 llama decoder graph with packed QKV, weights not yet loaded.
+
+    Uses the ``cuda`` execution provider so the GQA + PackQKV rewrites pack
+    Q/K/V into a single ``MatMul(hidden, Transpose(Concat(...)))``. Only the ONNX
+    graph is constructed — no GPU or onnxruntime execution is required. The fold
+    passes have not run yet (they run when weights are loaded).
+    """
+    return build_from_module(
+        registry.get("llama")(config),
+        config,
+        execution_provider="cuda",
+    )["model"]
+
+
+def _fp16_weight_tensors(model: ir.Model) -> dict[str, torch.Tensor]:
+    """fp16 tensors for every uninitialised parameter, mirroring an fp16 checkpoint.
+
+    Weight initializers always carry fully-static integer shapes (they are
+    concrete parameter tensors, never symbolic), so ``int(d)`` is safe here.
+    """
+    return {
+        name: torch.randn(*[int(d) for d in init.shape]).to(torch.float16)
+        for name, init in model.graph.initializers.items()
+        if init.const_value is None
+    }
+
+
+def _matmul_weight_initializers(model: ir.Model) -> list[ir.Value]:
+    """Return every ``MatMul`` second input that is a graph initializer."""
+    initializers = model.graph.initializers
+    weights: list[ir.Value] = []
+    for node in model.graph:
+        if node.op_type != "MatMul":
+            continue
+        weight = node.inputs[1]
+        if weight is not None and weight.name in initializers:
+            weights.append(weight)
+    return weights
+
+
+def _assert_matmul_weights_roundtrip_as_fp16(model: ir.Model, tmp_path: Path) -> None:
+    """Assert every folded MatMul weight reloads as fp16 with its bytes intact.
+
+    Serializes ``model`` through the production external-data path. This is the
+    ground-truth df203cc guard. The regression writes the packed / transposed fp16 *bytes* under a fp32 declaration, a corruption an in-memory
+    ``const_value`` check can mask (the backing array stays fp16 while only the
+    declared dtype widens) but that surfaces on serialize -> reload as a
+    fp32-declared initializer whose bytes no longer match the fp16 source.
+    ``ir.save`` with ``external_data`` mirrors the real fp16 export's
+    ``model.onnx`` + ``model.onnx.data`` layout, exactly where the dtype bug
+    corrupts the weight bytes.
+    """
+    before = {w.name: w.const_value.numpy() for w in _matmul_weight_initializers(model)}
+    assert before, "Expected the export to contain MatMul weight initializers to round-trip"
+
+    model_path = tmp_path / "model.onnx"
+    ir.save(model, model_path, external_data="model.onnx.data")
+    assert (tmp_path / "model.onnx.data").exists(), (
+        "fp16 weights should be externalized, exercising the real export's "
+        "model.onnx + model.onnx.data save path where the dtype bug corrupts bytes"
+    )
+
+    reloaded = ir.load(model_path)
+    for weight in _matmul_weight_initializers(reloaded):
+        assert weight.dtype == ir.DataType.FLOAT16, (
+            f"Reloaded MatMul weight {weight.name!r} serialized as {weight.dtype} "
+            f"(expected FLOAT16): fp16 bytes were written under a fp32 declaration "
+            f"— the df203cc serialize-time corruption."
+        )
+        np.testing.assert_array_equal(
+            weight.const_value.numpy(),
+            before[weight.name],
+            err_msg=(
+                f"Reloaded MatMul weight {weight.name!r} bytes changed across "
+                f"serialize -> reload; the fp16 packed weight was silently corrupted "
+                f"(fp16 data written under a fp32 dtype)."
+            ),
+        )
+
+
+@pytest.fixture(scope="module")
+def fp16_export() -> tuple[ArchitectureConfig, ir.Model]:
+    """A real fp16 packed-QKV export: build + ``apply_weights`` (folds inside).
+
+    Module-scoped so the (read-only) realistic-export assertions share a single
+    build instead of rebuilding per test.
+    """
+    config = _make_fp16_llama_config()
+    model = _build_fp16_decoder(config)
+    # apply_weights assigns the fp16 weights and then runs
+    # fold_initializers_after_weights, folding the Transpose/Concat weight nodes.
+    apply_weights(model, _fp16_weight_tensors(model))
+    return config, model
+
+
+class TestFp16FoldDtypeE2E:
+    def test_packed_qkv_weights_stay_fp16_after_export(
+        self, fp16_export: tuple[ArchitectureConfig, ir.Model]
+    ) -> None:
+        """An fp16 packed-QKV export keeps fp16 packed/transposed weights.
+
+        Guards the df203cc dtype-stamping path: the chained
+        ``FoldConcatInitializersPass`` (the Q/K/V pack) and
+        ``FoldTransposedInitializerPass`` (the weight transpose) must carry the
+        fp16 dtype through the rewrite-produced ``Concat`` intermediate instead
+        of defaulting it to ``FLOAT``.
+        """
+        config, model = fp16_export
+
+        # The packed path must actually have run, otherwise the dtype assertion
+        # below would pass vacuously (e.g. if the build stopped packing QKV).
+        gqa_nodes = [n for n in model.graph if n.op_type == "GroupQueryAttention"]
+        assert len(gqa_nodes) == config.num_hidden_layers, (
+            "Expected one packed GroupQueryAttention per layer; the fp16 GQA "
+            "PackQKV path was not exercised, so this regression guard is vacuous."
+        )
+
+        for gqa in gqa_nodes:
+            matmul = gqa.inputs[0].producer()
+            assert matmul is not None and matmul.op_type == "MatMul", (
+                "Packed GQA projection should be produced by a MatMul"
+            )
+            weight = matmul.inputs[1]
+            assert weight is not None, "Packed projection MatMul is missing its weight input"
+            # FoldConcat + FoldTranspose ran inside apply_weights, so the packed
+            # weight is now a plain folded initializer (no producer node left).
+            assert weight.producer() is None, (
+                "Packed QKV weight should be a folded initializer after export, "
+                "not the output of a residual Transpose/Concat node"
+            )
+            assert weight.const_value is not None
+            assert weight.const_value.dtype == ir.DataType.FLOAT16, (
+                f"Packed QKV weight {weight.name!r} widened to "
+                f"{weight.const_value.dtype} (expected FLOAT16); the fp16 fold "
+                f"dtype fix (df203cc) has regressed."
+            )
+            assert weight.dtype == ir.DataType.FLOAT16
+
+    def test_all_folded_matmul_weights_stay_fp16_after_export(
+        self, fp16_export: tuple[ArchitectureConfig, ir.Model]
+    ) -> None:
+        """Every folded MatMul weight in an fp16 export stays fp16.
+
+        Broader guard than the packed-QKV case: it also covers the standalone
+        ``o_proj`` and MLP weight transposes folded by
+        ``FoldTransposedInitializerPass``. A single fp32-widened weight here is
+        exactly what makes onnxruntime reject the fp16 model at load time.
+        """
+        _, model = fp16_export
+
+        weights = _matmul_weight_initializers(model)
+        assert weights, "Expected the export to contain MatMul weight initializers"
+        # Anti-vacuity: at least one weight must be a folded transposed initializer
+        # (FoldTransposedInitializerPass names them ``..._t``). Without this, the
+        # test could pass simply because no folding ran.
+        assert any(weight.name.endswith("_t") for weight in weights), (
+            "Expected at least one folded transposed weight (``..._t``); the fold "
+            "passes did not run, so this regression guard would be vacuous."
+        )
+
+        # fp32 widening (FLOAT) is the specific documented failure mode of the
+        # df203cc regression — the fold passes defaulting a dropped declared dtype
+        # to FLOAT — so we flag exactly that rather than any non-fp16 dtype.
+        widened = [
+            weight.name
+            for weight in weights
+            if weight.const_value is not None and weight.const_value.dtype == ir.DataType.FLOAT
+        ]
+        assert not widened, (
+            f"fp16 export produced fp32 MatMul weights after folding: {widened}. "
+            f"The fold passes must preserve fp16 (df203cc)."
+        )
+
+    def test_packed_qkv_survives_serialization_roundtrip(
+        self, fp16_export: tuple[ArchitectureConfig, ir.Model], tmp_path: Path
+    ) -> None:
+        """The fp16 export's folded weights survive serialize -> reload as fp16.
+
+        Ground-truth guard for the dtype-stamping path. The in-memory assertions
+        above read the IR ``const_value`` metadata; this drives the real export
+        artifact (``ir.save`` with external data, like the production
+        ``model.onnx`` + ``model.onnx.data``) and reloads it, the only check that
+        catches the df203cc serialize-time fp16-under-fp32 byte corruption — an
+        in-memory ``const_value.numpy()`` can stay fp16 and miss it.
+        """
+        _, model = fp16_export
+        _assert_matmul_weights_roundtrip_as_fp16(model, tmp_path)
+
+    def test_dropped_declared_dtype_falls_back_to_fp16_const(self, tmp_path: Path) -> None:
+        """Folding fp16 weights whose declared dtype was dropped stays fp16.
+
+        Guards the df203cc ``initializer_dtype`` const_value fallback. The
+        original regression arose because graph building left the packed-QKV
+        weights with fp16 ``const_value`` but no declared ``type``; the fold
+        passes then defaulted to ``FLOAT``. Here we reproduce that exact
+        condition on a real fp16 export — load fp16 weights, clear the declared
+        type on the packed-QKV ``Concat`` inputs, then run the real fold
+        orchestration — so the const_value fallback is the *only* thing that can
+        keep the folded weights fp16.
+
+        ``fold_initializers_after_weights`` is invoked directly (rather than via
+        ``apply_weights``) so the dropped-dtype condition can be injected between
+        loading the weights and folding; it is the same orchestration
+        ``apply_weights`` runs internally.
+        """
+        config = _make_fp16_llama_config()
+        model = _build_fp16_decoder(config)
+
+        # Load fp16 weights onto the initializers (mirrors apply_weights' assign
+        # step) without yet folding.
+        for name, tensor in _fp16_weight_tensors(model).items():
+            model.graph.initializers[name].const_value = ir.tensor(tensor.numpy())
+
+        # Reproduce the dropped-declared-dtype condition: clear the declared type
+        # on the packed-QKV weight initializers (the Concat inputs) while keeping
+        # their fp16 const_value.
+        dropped: list[str] = []
+        for node in model.graph:
+            if node.op_type != "Concat":
+                continue
+            for inp in node.inputs:
+                if (
+                    inp is not None
+                    and inp.name in model.graph.initializers
+                    and inp.const_value is not None
+                ):
+                    inp.type = None
+                    dropped.append(inp.name)
+        assert dropped, (
+            "Expected packed-QKV Concat inputs to drop declared dtype on; the "
+            "fp16 PackQKV path was not exercised, so this guard would be vacuous."
+        )
+
+        fold_initializers_after_weights(model)
+
+        gqa_nodes = [n for n in model.graph if n.op_type == "GroupQueryAttention"]
+        assert len(gqa_nodes) == config.num_hidden_layers
+        for gqa in gqa_nodes:
+            matmul = gqa.inputs[0].producer()
+            assert matmul is not None and matmul.op_type == "MatMul", (
+                "Packed GQA projection should be produced by a MatMul"
+            )
+            weight = matmul.inputs[1]
+            assert weight is not None, "Packed projection MatMul is missing its weight input"
+            assert weight.const_value is not None
+            assert weight.const_value.dtype == ir.DataType.FLOAT16, (
+                f"Packed QKV weight {weight.name!r} widened to "
+                f"{weight.const_value.dtype} when its declared dtype was dropped; "
+                f"the initializer_dtype const_value fallback (df203cc) regressed."
+            )
+
+        # Ground-truth check: the const_value fallback must also hold through the
+        # real serialize -> reload path. Reverting only the ``initializer_dtype``
+        # call-sites widens these dropped-dtype weights and corrupts their bytes
+        # here even though the type-stamp keeps the realistic-export tests green.
+        _assert_matmul_weights_roundtrip_as_fp16(model, tmp_path)

From 365e6248891d521a4e9dfe71a54509e9bebb5325 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Thu, 4 Jun 2026 00:10:06 +0000
Subject: [PATCH 35/38] test(fp16): make fp16 fold e2e fixture function-scoped
 for hermetic, xdist-safe runs

The fp16_export fixture was module-scoped and shared across the realistic-export
tests. The serialize-roundtrip test calls ir.save(external_data=...) on that
shared model; on some onnx_ir versions ir.save offloads initializer const_values
to external tensors in place, which can leak mutated/externalized state into the
other tests that read the same model. Under pytest-xdist the tests' execution
order is not guaranteed, so this cross-test contamination is order-dependent and
can flake (a folded weight intermittently observed as FLOAT instead of FLOAT16,
falsely reporting a df203cc regression).

Switching the fixture to function scope gives each test a fresh, hermetic build,
eliminating the cross-test state dependence across all onnx_ir versions at
negligible cost (the synthetic model is tiny). No change to test coverage or
assertions; the four df203cc guards are unchanged.

Verified post-change: 40 test4-alone + 40 full-file serial + 32 full-file
xdist(-n4) fresh-process runs, 0 failures; ruff clean.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_passes/_fold_dtype_e2e_test.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/mobius/_passes/_fold_dtype_e2e_test.py b/src/mobius/_passes/_fold_dtype_e2e_test.py
index da25ab7c..0b98f59c 100644
--- a/src/mobius/_passes/_fold_dtype_e2e_test.py
+++ b/src/mobius/_passes/_fold_dtype_e2e_test.py
@@ -174,12 +174,18 @@ def _assert_matmul_weights_roundtrip_as_fp16(model: ir.Model, tmp_path: Path) ->
         )
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture
 def fp16_export() -> tuple[ArchitectureConfig, ir.Model]:
     """A real fp16 packed-QKV export: build + ``apply_weights`` (folds inside).
 
-    Module-scoped so the (read-only) realistic-export assertions share a single
-    build instead of rebuilding per test.
+    Function-scoped (a fresh build per test) on purpose: the serialize-roundtrip
+    test calls ``ir.save`` on the model, and on some ``onnx_ir`` versions that
+    offloads the initializers' ``const_value`` to external tensors in place. A
+    shared (module-scoped) model could then leak that mutated/externalized state
+    into another test, and under ``pytest-xdist`` the tests' execution order is
+    not guaranteed — making such cross-test contamination order-dependent and
+    flaky. A fresh model per test keeps each test hermetic and deterministic. The
+    build is tiny, so the rebuild cost is negligible.
     """
     config = _make_fp16_llama_config()
     model = _build_fp16_decoder(config)

From 264a94942506e115ac4db148c3b41f65d7d26e4a Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Wed, 3 Jun 2026 23:19:20 +0000
Subject: [PATCH 36/38] feat(graph-diff): recurse into subgraphs so
 control-flow/phase-split changes are visible

The Architecture-Diff tooling collapsed GRAPH-typed attributes (If
then_branch / else_branch, Loop / Scan bodies) to a bare type string and
never recursed into them. As a result the per-layer static-cache
phase-split introduced by PR #328 -- an If(Greater(seq_len, 1)) selecting
a prefill (masked) vs decode (Flash) attention path -- was completely
invisible to the Architecture Diff CI: the top-level op sequence is
unchanged (the If node is present on both sides), so the only signal lives
inside the branch subgraphs that were being discarded.

This recurses GRAPH and GRAPHS attributes into nested canonical forms so
subgraph node structure participates in the comparison, reusing
canonicalize_graph (inner node/value names are ignored the same way
top-level ones are). diff_graphs gains a dedicated
subgraph_structure_change record (MODERATE severity) for structurally
significant subgraph deltas -- a node/branch added, removed, rewired, or a
subgraph interface change -- while a pure inner-attribute tweak stays
changed_attrs (MINOR). Structural significance propagates upward through
nested subgraphs (e.g. an If inside an If). The nested diff detail is
surfaced in the report (e.g. "then_branch: node[0] Concat: axis: 0 -> 1").

Additive and backward-compatible: non-GRAPH attributes are unchanged and
the arch_diff.py consumer (which reads only op_sequence / node counts /
the changes list) is unaffected.

Adds 16 regression tests covering subgraph recursion, the
structural-vs-minor severity boundary (incl. nested), GRAPHS-plural,
op-swap no-double-count, and the readable fallback path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/_graph_diff.py      | 149 +++++++++++--
 src/mobius/_graph_diff_test.py | 375 +++++++++++++++++++++++++++++++++
 2 files changed, 511 insertions(+), 13 deletions(-)

diff --git a/src/mobius/_graph_diff.py b/src/mobius/_graph_diff.py
index f85ea1e3..3cd78ae5 100644
--- a/src/mobius/_graph_diff.py
+++ b/src/mobius/_graph_diff.py
@@ -54,11 +54,23 @@ def _dtype_str(value: ir.Value) -> str:
     return "UNKNOWN"
 
 
+# Sentinel keys marking a recursively-canonicalised subgraph payload inside an
+# attribute's comparable value.  diff_graphs uses these to render subgraph
+# deltas readably instead of dumping a raw nested canonical dict.
+_SUBGRAPH_KEY = "__subgraph__"
+_SUBGRAPHS_KEY = "__subgraphs__"
+
+
 def _attr_to_comparable(attr: ir.Attr) -> Any:
     """Convert an attribute to a JSON-serialisable comparable value.
 
-    Graph and tensor attributes are reduced to their type string so
-    that canonicalisation stays lightweight.
+    Most attributes reduce to their scalar / list value.  GRAPH-typed
+    attributes (``If``'s ``then_branch`` / ``else_branch``, ``Loop`` /
+    ``Scan`` bodies) are *recursively canonicalised* so subgraph structure
+    participates in the diff — without this, per-layer phase-split ``If``
+    subgraphs are invisible to the architecture diff.  Remaining opaque
+    types (TENSOR, SPARSE_TENSOR, TYPE_PROTO, …) are recorded as their type
+    string to keep canonicalisation lightweight.
     """
     simple_types = {
         ir.AttributeType.FLOAT,
@@ -74,7 +86,14 @@ def _attr_to_comparable(attr: ir.Attr) -> Any:
         if isinstance(v, tuple):
             return list(v)
         return v
-    # For TENSOR, GRAPH etc. just record the type
+    # Recurse into subgraphs so their node structure is compared, not collapsed.
+    # Subgraphs reuse canonicalize_graph, so inner node/value names are ignored
+    # the same way top-level ones are (see canonicalize_graph's name-independence).
+    if attr.type == ir.AttributeType.GRAPH:
+        return {_SUBGRAPH_KEY: canonicalize_graph(attr.value)}
+    if attr.type == ir.AttributeType.GRAPHS:
+        return {_SUBGRAPHS_KEY: [canonicalize_graph(g) for g in attr.value]}
+    # For TENSOR, SPARSE_TENSOR, TYPE_PROTO, … just record the type.
     return f"<{attr.type.name}>"
 
 
@@ -203,13 +222,92 @@ def _describe_port_diff(base_port: dict, head_port: dict) -> str:
     return "; ".join(parts) or "changed"
 
 
+def _is_subgraph_payload(value: Any) -> bool:
+    """True if *value* is a recursively-canonicalised subgraph payload."""
+    return isinstance(value, dict) and (_SUBGRAPH_KEY in value or _SUBGRAPHS_KEY in value)
+
+
+def _subgraph_list(value: Any) -> list[dict]:
+    """Extract the list of subgraph canonical forms from an attr payload."""
+    if isinstance(value, dict):
+        if _SUBGRAPH_KEY in value:
+            return [value[_SUBGRAPH_KEY]]
+        if _SUBGRAPHS_KEY in value:
+            return list(value[_SUBGRAPHS_KEY])
+    return []
+
+
+# Nested sub-change types that make a subgraph delta *structurally* significant
+# (a node/branch was added/removed, rewired, or the subgraph interface moved),
+# as opposed to a mere inner-attribute tweak.  Any of these promotes the
+# containing attribute to a subgraph_structure_change → MODERATE.  Note this is
+# uniformly MODERATE: unlike a *top-level* interface_change (MAJOR, an external
+# model-contract break), a subgraph's interface is internal control-flow plumbing,
+# so it stays MODERATE here — a deliberate asymmetry.
+# ``subgraph_structure_change`` is included so structural significance
+# *propagates* upward through nested subgraphs (e.g. an If inside an If).
+_STRUCTURAL_SUB_TYPES = frozenset(
+    {
+        "added_node",
+        "removed_node",
+        "changed_connectivity",
+        "interface_change",
+        "subgraph_structure_change",
+    }
+)
+
+
+def _describe_subgraph_attr_change(key: str, base_val: Any, head_val: Any) -> tuple[str, bool]:
+    """Describe a GRAPH-typed attribute change, recursing into the subgraph(s).
+
+    Returns ``(detail, structural)`` where *detail* is a readable summary
+    that surfaces the nested diff (e.g. ``"then_branch: node[0] Concat:
+    axis: 0 → 1"`` or ``"then_branch: + Mul; - Add"``) and *structural* is
+    True when the nested delta adds/removes a node or branch, rewires
+    connectivity, or changes the subgraph interface — i.e. a change that
+    should outrank a pure inner-attribute tweak.
+    """
+    base_subs = _subgraph_list(base_val)
+    head_subs = _subgraph_list(head_val)
+    count = max(len(base_subs), len(head_subs))
+    structural = False
+    parts: list[str] = []
+    for idx in range(count):
+        bs = base_subs[idx] if idx < len(base_subs) else None
+        hs = head_subs[idx] if idx < len(head_subs) else None
+        # The attribute *key* already names a single subgraph (then_branch,
+        # body, …); only disambiguate by index when there are several (GRAPHS).
+        label = "" if count == 1 else f"subgraph[{idx}]"
+        if bs is None or hs is None:
+            # A whole branch/body was added or removed.
+            structural = True
+            verb = "added" if bs is None else "removed"
+            parts.append(f"{label} {verb}".strip())
+            continue
+        sub_changes = diff_graphs(bs, hs)
+        if not sub_changes:
+            continue
+        if {c["type"] for c in sub_changes} & _STRUCTURAL_SUB_TYPES:
+            structural = True
+        inner = "; ".join(c["details"] for c in sub_changes)
+        parts.append(f"{label}: {inner}" if label else inner)
+    detail = f"{key}: " + ("; ".join(parts) if parts else "subgraph changed")
+    return detail, structural
+
+
 def diff_graphs(base: dict, head: dict) -> list[dict[str, Any]]:
     """Compare two canonical graph representations.
 
     Returns a list of change dicts.  Each dict has a ``"type"`` key with
     one of: ``"added_node"``, ``"removed_node"``, ``"changed_attrs"``,
+    ``"subgraph_structure_change"``, ``"changed_connectivity"``,
     ``"interface_change"``, ``"initializer_change"``.  A ``"details"``
     key carries human-readable information about the change.
+
+    ``subgraph_structure_change`` is emitted for a GRAPH-typed attribute
+    (``If`` branches, ``Loop`` / ``Scan`` bodies) whose nested graph gains
+    or loses a node/branch, is rewired, or changes interface; a subgraph
+    delta that only tweaks an inner attribute stays ``changed_attrs``.
     """
     changes: list[dict[str, Any]] = []
 
@@ -304,19 +402,32 @@ def diff_graphs(base: dict, head: dict) -> list[dict[str, Any]]:
         if bn["attributes"] != hn["attributes"]:
             ba = bn["attributes"]
             ha = hn["attributes"]
-            attr_details: list[str] = []
+            plain_details: list[str] = []
             all_keys = sorted(set(ba) | set(ha))
             for k in all_keys:
                 bv = ba.get(k)
                 hv = ha.get(k)
-                if bv != hv:
-                    attr_details.append(f"{k}: {bv!r} → {hv!r}")
-            changes.append(
-                {
-                    "type": "changed_attrs",
-                    "details": (f"node[{i}] {bn['op_type']}: " + ", ".join(attr_details)),
-                }
-            )
+                if bv == hv:
+                    continue
+                if _is_subgraph_payload(bv) or _is_subgraph_payload(hv):
+                    detail, structural = _describe_subgraph_attr_change(k, bv, hv)
+                    changes.append(
+                        {
+                            "type": (
+                                "subgraph_structure_change" if structural else "changed_attrs"
+                            ),
+                            "details": f"node[{i}] {bn['op_type']}: {detail}",
+                        }
+                    )
+                else:
+                    plain_details.append(f"{k}: {bv!r} → {hv!r}")
+            if plain_details:
+                changes.append(
+                    {
+                        "type": "changed_attrs",
+                        "details": (f"node[{i}] {bn['op_type']}: " + ", ".join(plain_details)),
+                    }
+                )
         if bn["input_ids"] != hn["input_ids"]:
             changes.append(
                 {
@@ -343,7 +454,12 @@ def _change_status(change_list: list[dict[str, Any]]) -> str:
     types = {c["type"] for c in change_list}
     if types & {"interface_change"}:
         return _STATUS_MAJOR
-    if types & {"added_node", "removed_node", "changed_connectivity"}:
+    if types & {
+        "added_node",
+        "removed_node",
+        "changed_connectivity",
+        "subgraph_structure_change",
+    }:
         return _STATUS_MODERATE
     if types & {"changed_attrs", "initializer_change"}:
         return _STATUS_MINOR
@@ -459,6 +575,7 @@ def _sha_link(sha: str) -> str:
             removed = [c for c in change_list if c["type"] == "removed_node"]
             attrs = [c for c in change_list if c["type"] == "changed_attrs"]
             connectivity = [c for c in change_list if c["type"] == "changed_connectivity"]
+            subgraph = [c for c in change_list if c["type"] == "subgraph_structure_change"]
             iface = [c for c in change_list if c["type"] == "interface_change"]
             inits = [c for c in change_list if c["type"] == "initializer_change"]
 
@@ -474,6 +591,12 @@ def _sha_link(sha: str) -> str:
                     lines.append(f"- `{c['details']}`")
                 lines.append("")
 
+            if subgraph:
+                lines.append("**Subgraph structure changes:**")
+                for c in subgraph:
+                    lines.append(f"- `{c['details']}`")
+                lines.append("")
+
             if attrs:
                 lines.append("**Modified attributes:**")
                 for c in attrs:
diff --git a/src/mobius/_graph_diff_test.py b/src/mobius/_graph_diff_test.py
index 2f87842e..bcb94ffd 100644
--- a/src/mobius/_graph_diff_test.py
+++ b/src/mobius/_graph_diff_test.py
@@ -12,6 +12,8 @@
 import onnx_ir as ir
 
 from mobius._graph_diff import (
+    _SUBGRAPH_KEY,
+    _SUBGRAPHS_KEY,
     canonicalize_graph,
     diff_graphs,
     render_markdown,
@@ -611,3 +613,376 @@ def test_structural_change_detected(self) -> None:
         head = _add_relu_graph()
         changes = diff_graphs(canonicalize_graph(base), canonicalize_graph(head))
         assert len(changes) > 0
+
+
+# ------------------------------------------------------------------
+# Subgraph recursion — helpers (graphs containing an If with subgraphs)
+# ------------------------------------------------------------------
+
+
+def _if_graph(
+    *,
+    then_op: str = "Add",
+    else_op: str = "Sub",
+    then_axis: int | None = None,
+    then_domain: str = "",
+    then_node_name: str = "then_op",
+    else_node_name: str = "else_op",
+) -> ir.Graph:
+    """Return a graph: y = If(cond) {then: then_op(x,x)} {else: else_op(x,x)}.
+
+    ``then_axis`` (when given) attaches an ``axis`` attribute to the
+    then-branch node so inner-attribute deltas can be exercised.
+    ``then_domain`` sets the then-branch node's op domain (diff_graphs does
+    not compare domain, so this exercises the 'subgraph changed' fallback).
+    ``*_node_name`` allow renaming inner nodes to prove names are ignored.
+    """
+    x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4]))
+    cond = ir.val("cond", type=ir.TensorType(ir.DataType.BOOL), shape=ir.Shape([]))
+
+    then_attrs = [ir.AttrInt64("axis", then_axis)] if then_axis is not None else []
+    then_node = ir.Node(then_domain, then_op, [x, x], then_attrs, name=then_node_name)
+    then_out = then_node.outputs[0]
+    then_out.name = "then_out"
+    then_g = ir.Graph([], [then_out], nodes=[then_node], name="then_branch")
+
+    else_node = ir.Node("", else_op, [x, x], name=else_node_name)
+    else_out = else_node.outputs[0]
+    else_out.name = "else_out"
+    else_g = ir.Graph([], [else_out], nodes=[else_node], name="else_branch")
+
+    if_node = ir.Node(
+        "",
+        "If",
+        [cond],
+        [ir.AttrGraph("then_branch", then_g), ir.AttrGraph("else_branch", else_g)],
+        name="if_node",
+    )
+    if_out = if_node.outputs[0]
+    if_out.name = "y"
+    return ir.Graph([x, cond], [if_out], nodes=[if_node])
+
+
+# ------------------------------------------------------------------
+# Subgraph recursion — tests (architect D12: descend into subgraphs)
+# ------------------------------------------------------------------
+
+
+def _nested_if_graph(
+    *, inner_then_op: str = "Add", inner_then_axis: int | None = None
+) -> ir.Graph:
+    """Return a graph whose If then-branch itself contains an If (depth 2).
+
+    ``inner_then_axis`` attaches an ``axis`` attribute to the innermost
+    then-branch node, to exercise a *pure inner-attribute* delta nested
+    two levels deep (which must stay MINOR, not promote to structural).
+    """
+    x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4]))
+    cond = ir.val("cond", type=ir.TensorType(ir.DataType.BOOL), shape=ir.Shape([]))
+
+    # Innermost two branches.
+    inner_then_attrs = (
+        [ir.AttrInt64("axis", inner_then_axis)] if inner_then_axis is not None else []
+    )
+    inner_then = ir.Node("", inner_then_op, [x, x], inner_then_attrs, name="inner_then")
+    it_out = inner_then.outputs[0]
+    it_out.name = "it_out"
+    inner_then_g = ir.Graph([], [it_out], nodes=[inner_then], name="inner_then_branch")
+
+    inner_else = ir.Node("", "Sub", [x, x], name="inner_else")
+    ie_out = inner_else.outputs[0]
+    ie_out.name = "ie_out"
+    inner_else_g = ir.Graph([], [ie_out], nodes=[inner_else], name="inner_else_branch")
+
+    inner_if = ir.Node(
+        "",
+        "If",
+        [cond],
+        [
+            ir.AttrGraph("then_branch", inner_then_g),
+            ir.AttrGraph("else_branch", inner_else_g),
+        ],
+        name="inner_if",
+    )
+    inner_if_out = inner_if.outputs[0]
+    inner_if_out.name = "inner_if_out"
+    # Outer then-branch wraps the inner If; outer else-branch is a plain op.
+    outer_then_g = ir.Graph([], [inner_if_out], nodes=[inner_if], name="outer_then")
+
+    outer_else = ir.Node("", "Mul", [x, x], name="outer_else")
+    oe_out = outer_else.outputs[0]
+    oe_out.name = "oe_out"
+    outer_else_g = ir.Graph([], [oe_out], nodes=[outer_else], name="outer_else")
+
+    outer_if = ir.Node(
+        "",
+        "If",
+        [cond],
+        [
+            ir.AttrGraph("then_branch", outer_then_g),
+            ir.AttrGraph("else_branch", outer_else_g),
+        ],
+        name="outer_if",
+    )
+    out = outer_if.outputs[0]
+    out.name = "y"
+    return ir.Graph([x, cond], [out], nodes=[outer_if])
+
+
+def _graphs_attr_graph(*, body_ops: list[str]) -> ir.Graph:
+    """Return a graph with a node carrying a GRAPHS-plural attribute.
+
+    GRAPHS is rare for standard ops (``If`` / ``Loop`` / ``Scan`` bodies are
+    GRAPH singular), so this uses a synthetic op purely to exercise the
+    plural recursion branch in ``_attr_to_comparable``.
+    """
+    x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4]))
+
+    bodies = []
+    for idx, op in enumerate(body_ops):
+        n = ir.Node("", op, [x, x], name=f"body_{idx}")
+        o = n.outputs[0]
+        o.name = f"body_out_{idx}"
+        bodies.append(ir.Graph([], [o], nodes=[n], name=f"body_g_{idx}"))
+
+    multi = ir.Node("", "CustomMulti", [x], [ir.AttrGraphs("bodies", bodies)], name="cm")
+    out = multi.outputs[0]
+    out.name = "y"
+    return ir.Graph([x], [out], nodes=[multi])
+
+
+def _if_graph_branch_output_dtype(*, then_dtype: ir.DataType) -> ir.Graph:
+    """If-graph whose then-branch *output dtype* varies (subgraph interface)."""
+    x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4]))
+    cond = ir.val("cond", type=ir.TensorType(ir.DataType.BOOL), shape=ir.Shape([]))
+
+    cast = ir.Node("", "Cast", [x], [ir.AttrInt64("to", int(then_dtype))], name="cast")
+    then_out = cast.outputs[0]
+    then_out.name = "then_out"
+    then_out.type = ir.TensorType(then_dtype)
+    then_out.shape = ir.Shape([1, 4])
+    then_g = ir.Graph([], [then_out], nodes=[cast], name="then_branch")
+
+    else_node = ir.Node("", "Identity", [x], name="else_node")
+    else_out = else_node.outputs[0]
+    else_out.name = "else_out"
+    else_out.type = ir.TensorType(ir.DataType.FLOAT)
+    else_out.shape = ir.Shape([1, 4])
+    else_g = ir.Graph([], [else_out], nodes=[else_node], name="else_branch")
+
+    if_node = ir.Node(
+        "",
+        "If",
+        [cond],
+        [ir.AttrGraph("then_branch", then_g), ir.AttrGraph("else_branch", else_g)],
+        name="if_node",
+    )
+    if_out = if_node.outputs[0]
+    if_out.name = "y"
+    return ir.Graph([x, cond], [if_out], nodes=[if_node])
+
+
+class TestSubgraphRecursion:
+    """canonicalize_graph / diff_graphs descend into GRAPH-typed attrs."""
+
+    def test_canonicalize_descends_into_subgraphs(self) -> None:
+        """GRAPH attrs are recursively canonicalised, not collapsed to a type."""
+        canon = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        if_attrs = canon["nodes"][0]["attributes"]
+        # Not collapsed to the old "<GRAPH>" placeholder.
+        assert if_attrs["then_branch"] != "<GRAPH>"
+        assert if_attrs["else_branch"] != "<GRAPH>"
+        # Nested canonical form carries the subgraph's op sequence.
+        then_sub = if_attrs["then_branch"][_SUBGRAPH_KEY]
+        else_sub = if_attrs["else_branch"][_SUBGRAPH_KEY]
+        assert then_sub["op_sequence"] == ["Add"]
+        assert else_sub["op_sequence"] == ["Sub"]
+
+    def test_identical_subgraphs_no_diff(self) -> None:
+        """Two structurally identical If-graphs produce no diff."""
+        base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        head = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        assert diff_graphs(base, head) == []
+
+    def test_subgraph_node_names_ignored(self) -> None:
+        """Renaming inner subgraph nodes does not produce a spurious diff."""
+        base = canonicalize_graph(_if_graph(then_node_name="a", else_node_name="b"))
+        head = canonicalize_graph(_if_graph(then_node_name="x", else_node_name="y"))
+        assert diff_graphs(base, head) == []
+
+    def test_then_branch_op_delta_is_structural(self) -> None:
+        """A differing then-branch op is a MODERATE subgraph_structure_change."""
+        base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub"))
+        changes = diff_graphs(base, head)
+        structural = [c for c in changes if c["type"] == "subgraph_structure_change"]
+        assert structural, "op delta inside a branch must be structural"
+        assert any("then_branch" in c["details"] for c in structural)
+        # The op-level subgraph delta is surfaced (added/removed inner node).
+        assert any("Add" in c["details"] and "Mul" in c["details"] for c in structural)
+
+    def test_else_branch_op_delta_is_structural(self) -> None:
+        """A differing else-branch op is detected independently of then-branch."""
+        base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        head = canonicalize_graph(_if_graph(then_op="Add", else_op="Div"))
+        changes = diff_graphs(base, head)
+        structural = [c for c in changes if c["type"] == "subgraph_structure_change"]
+        assert any("else_branch" in c["details"] for c in structural)
+
+    def test_structural_subgraph_change_is_moderate(self) -> None:
+        """render_markdown rates a subgraph structure change as MODERATE (🟡)."""
+        base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub"))
+        changes = diff_graphs(base, head)
+        diffs = {
+            "model": {
+                "sub": {
+                    "changes": changes,
+                    "_base_ops": base["op_sequence"],
+                    "_head_ops": head["op_sequence"],
+                    "_base_node_count": len(base["nodes"]),
+                    "_head_node_count": len(head["nodes"]),
+                }
+            }
+        }
+        md = render_markdown(diffs)
+        assert "🟡" in md
+        assert "Subgraph structure changes" in md
+
+    def test_inner_attribute_delta_is_minor_and_surfaced(self) -> None:
+        """A pure inner-attribute tweak (Concat axis) stays changed_attrs/MINOR.
+
+        The actual nested detail (``axis: 0 → 1``) must be surfaced, not an
+        opaque ``"N sub-change(s)"`` summary.
+        """
+        base = canonicalize_graph(_if_graph(then_op="Concat", else_op="Sub", then_axis=0))
+        head = canonicalize_graph(_if_graph(then_op="Concat", else_op="Sub", then_axis=1))
+        changes = diff_graphs(base, head)
+        # Not promoted to structural — it's an inner-attr-only change.
+        assert not any(c["type"] == "subgraph_structure_change" for c in changes)
+        attr_changes = [c for c in changes if c["type"] == "changed_attrs"]
+        assert attr_changes
+        detail = " ".join(c["details"] for c in attr_changes)
+        assert "then_branch" in detail
+        # The real inner delta is surfaced.
+        assert "axis" in detail and "0" in detail and "1" in detail
+
+    def test_nested_subgraph_delta_detected(self) -> None:
+        """A delta in an If-inside-an-If (depth 2) propagates to the top diff."""
+        base = canonicalize_graph(_nested_if_graph(inner_then_op="Add"))
+        head = canonicalize_graph(_nested_if_graph(inner_then_op="Mul"))
+        changes = diff_graphs(base, head)
+        structural = [c for c in changes if c["type"] == "subgraph_structure_change"]
+        assert structural, "innermost branch op delta must surface at the top"
+        assert any("then_branch" in c["details"] for c in structural)
+
+    def test_nested_pure_attr_delta_stays_minor(self) -> None:
+        """A pure inner-attr tweak nested two levels deep stays changed_attrs/MINOR.
+
+        (Code-review NIT: lock the severity boundary — nesting must not
+        spuriously promote a non-structural change to structural.)
+        """
+        base = canonicalize_graph(_nested_if_graph(inner_then_op="Concat", inner_then_axis=0))
+        head = canonicalize_graph(_nested_if_graph(inner_then_op="Concat", inner_then_axis=1))
+        changes = diff_graphs(base, head)
+        assert not any(c["type"] == "subgraph_structure_change" for c in changes)
+        attr_changes = [c for c in changes if c["type"] == "changed_attrs"]
+        assert attr_changes
+        # The innermost axis delta is still surfaced through both nesting levels.
+        detail = " ".join(c["details"] for c in attr_changes)
+        assert "then_branch" in detail and "axis" in detail
+
+    def test_subgraph_changed_fallback_path(self) -> None:
+        """A subgraph delta invisible to diff_graphs hits the readable fallback.
+
+        (Code-review NIT: lock the ``"subgraph changed"`` fallback.)  The
+        then-branch node's *domain* differs; canonicalize records domain so
+        the attribute payloads differ, but diff_graphs does not compare
+        domain, so the nested diff is empty → the summary falls back to
+        ``"subgraph changed"`` and is classified MINOR (changed_attrs).
+        """
+        base = canonicalize_graph(_if_graph(then_op="Add", then_domain=""))
+        head = canonicalize_graph(_if_graph(then_op="Add", then_domain="custom.domain"))
+        changes = diff_graphs(base, head)
+        assert not any(c["type"] == "subgraph_structure_change" for c in changes)
+        attr_changes = [c for c in changes if c["type"] == "changed_attrs"]
+        assert any("subgraph changed" in c["details"] for c in attr_changes)
+
+    def test_graphs_plural_attr_recursed_and_diffed(self) -> None:
+        """GRAPHS-plural attrs are canonicalised and per-body deltas detected."""
+        base = canonicalize_graph(_graphs_attr_graph(body_ops=["Add", "Add"]))
+        head = canonicalize_graph(_graphs_attr_graph(body_ops=["Add", "Mul"]))
+        # Plural payload is recursed, not collapsed.
+        bodies_attr = base["nodes"][0]["attributes"]["bodies"]
+        assert bodies_attr != "<GRAPHS>"
+        changes = diff_graphs(base, head)
+        structural = [c for c in changes if c["type"] == "subgraph_structure_change"]
+        assert structural
+        # Plural payload genuinely drove the AttributeType.GRAPHS branch.
+        assert _SUBGRAPHS_KEY in bodies_attr
+        assert len(bodies_attr[_SUBGRAPHS_KEY]) == 2
+        # The differing body is identified by index.
+        assert any("subgraph[1]" in c["details"] for c in structural)
+
+    def test_graphs_plural_branch_count_change_is_structural(self) -> None:
+        """Adding/removing a body in a GRAPHS-plural attr is structural."""
+        base = canonicalize_graph(_graphs_attr_graph(body_ops=["Add"]))
+        head = canonicalize_graph(_graphs_attr_graph(body_ops=["Add", "Mul"]))
+        changes = diff_graphs(base, head)
+        structural = [c for c in changes if c["type"] == "subgraph_structure_change"]
+        assert structural
+        assert any("added" in c["details"] or "removed" in c["details"] for c in structural)
+
+    def test_subgraph_interface_change_is_structural(self) -> None:
+        """A subgraph *interface* (output dtype) change is structural."""
+        base = canonicalize_graph(_if_graph_branch_output_dtype(then_dtype=ir.DataType.FLOAT))
+        head = canonicalize_graph(
+            _if_graph_branch_output_dtype(then_dtype=ir.DataType.FLOAT16)
+        )
+        changes = diff_graphs(base, head)
+        structural = [c for c in changes if c["type"] == "subgraph_structure_change"]
+        assert structural, "subgraph interface change must be structural"
+        assert any("then_branch" in c["details"] for c in structural)
+
+    def test_subgraph_delta_does_not_perturb_top_level_ops(self) -> None:
+        """The top-level op sequence is unchanged when only a subgraph differs."""
+        base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub"))
+        # Both still have exactly one top-level If node.
+        assert base["op_sequence"] == ["If"] == head["op_sequence"]
+        changes = diff_graphs(base, head)
+        # No spurious added/removed at the top level.
+        assert not any(c["type"] in {"added_node", "removed_node"} for c in changes)
+
+    def test_top_level_op_swap_not_double_counted(self) -> None:
+        """An op_type swap at the same position is structural-only, never MINOR.
+
+        (Reviewer 6def2895 (a): pick one classification — structural wins.)
+        """
+
+        def _single(op: str) -> ir.Graph:
+            x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4]))
+            n = ir.Node("", op, [x, x], name="n")
+            o = n.outputs[0]
+            o.name = "y"
+            return ir.Graph([x], [o], nodes=[n])
+
+        changes = diff_graphs(
+            canonicalize_graph(_single("Add")), canonicalize_graph(_single("Mul"))
+        )
+        types = {c["type"] for c in changes}
+        assert types == {"added_node", "removed_node"}
+        assert "changed_attrs" not in types
+
+    def test_subgraph_op_swap_is_structural_only(self) -> None:
+        """MEA↔Flash-style op swap *inside* a branch is structural, not MINOR.
+
+        A same-position op_type swap within a subgraph must classify as
+        subgraph_structure_change only — not additionally as changed_attrs.
+        """
+        base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub"))
+        head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub"))
+        changes = diff_graphs(base, head)
+        # Exactly one change for the swapped branch, classified structural.
+        assert [c["type"] for c in changes] == ["subgraph_structure_change"]
+        assert not any(c["type"] == "changed_attrs" for c in changes)

From dab98b98da639ec54b612a713e8c44824bf56f97 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Thu, 4 Jun 2026 22:07:07 +0000
Subject: [PATCH 37/38] docs(changelog): note folded-in Architecture-Diff
 subgraph recursion under Developer Tooling

Adds an Internal / Developer Tooling subsection so the arch-diff recursion
change folded in from #330 is not orphaned in an export-correctness PR. The
tool now recurses into If/Loop/Scan subgraphs (with a subgraph_structure_change
MODERATE severity) and is landed here because this PR introduces the first
control-flow/subgraph change (the static-cache per-layer phase-split) the old
top-level-only diff could not see into. Developer-tooling only; no exported-graph
or runtime impact.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 CHANGELOG.md | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a071b045..51bee326 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,6 +23,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ---
 
+### Developer Tooling
+
+#### Internal
+
+- The Architecture-Diff CI tool (`src/mobius/_graph_diff.py`) now recurses into
+  `If` / `Loop` / `Scan` `GRAPH`-typed attributes instead of collapsing them to a
+  bare type string, so changes *inside* control-flow subgraphs are visible in the
+  diff. A dedicated `subgraph_structure_change` (MODERATE) severity is emitted for
+  structural in-branch deltas (node added/removed/rewired, branch added/removed),
+  while pure inner-attribute tweaks stay `changed_attrs` (MINOR). Added here
+  because this PR introduces the static-cache per-layer phase-split
+  `If(Greater(seq_len, 1))` — the first control-flow/subgraph change the old
+  top-level-only diff could not see — so the tool that guards architectural
+  changes must be able to see it. Developer-tooling only; no exported-graph or
+  runtime impact.
+
+---
+
 ### GQA Present KV-Cache Shape Fix
 
 #### Fixed

From 856b9860a9f4b3e78acec11897210507d5142777 Mon Sep 17 00:00:00 2001
From: titaiwang <titaiwang@microsoft.com>
Date: Thu, 4 Jun 2026 22:09:40 +0000
Subject: [PATCH 38/38] =?UTF-8?q?docs(changelog):=20correct=20arch-diff=20?=
 =?UTF-8?q?note=20=E2=80=94=20old=20diff=20sees=20the=20added=20If,=20not?=
 =?UTF-8?q?=20in-branch=20changes?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The previous wording over-claimed that #328 introduces a change "the old
top-level diff could not see". The old diff DOES report the newly-added
per-layer If phase-split nodes. Reword to the accurate future-proofing
framing: #328 adds the first control-flow subgraph structure; once those If
nodes exist, FUTURE changes inside their branches would be invisible to the
old (subgraph-collapsing) diff, so the recursion is landed now to keep
in-branch changes visible going forward.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 CHANGELOG.md | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 51bee326..7fefb3d2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,15 +29,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - The Architecture-Diff CI tool (`src/mobius/_graph_diff.py`) now recurses into
   `If` / `Loop` / `Scan` `GRAPH`-typed attributes instead of collapsing them to a
-  bare type string, so changes *inside* control-flow subgraphs are visible in the
+  bare type marker, so changes *inside* control-flow subgraphs are visible in the
   diff. A dedicated `subgraph_structure_change` (MODERATE) severity is emitted for
   structural in-branch deltas (node added/removed/rewired, branch added/removed),
-  while pure inner-attribute tweaks stay `changed_attrs` (MINOR). Added here
-  because this PR introduces the static-cache per-layer phase-split
-  `If(Greater(seq_len, 1))` — the first control-flow/subgraph change the old
-  top-level-only diff could not see — so the tool that guards architectural
-  changes must be able to see it. Developer-tooling only; no exported-graph or
-  runtime impact.
+  while pure inner-attribute tweaks stay `changed_attrs` (MINOR). This PR introduces
+  the first control-flow/subgraph *structure* into the exported graphs — the
+  static-cache per-layer phase-split `If(Greater(seq_len, 1))`. The previous
+  top-level diff already reports those newly-added `If` nodes; what it could not do
+  is see *inside* a branch. Once these `If` nodes exist, any future change within
+  their branches would be invisible to the old diff (which collapsed each subgraph
+  to a type marker), so the recursion is added now to keep such in-branch changes
+  visible going forward. Developer-tooling only; no exported-graph or runtime impact.
 
 ---