Fix Phi-3.5 ONNX Attention (static-cache) + GQA fp16 export bugs by titaiwangms · Pull Request #328 · onnxruntime/mobius

titaiwangms · 2026-06-02T00:02:22Z

Fix opset-24 Attention GQA export (phase-split decode) + long-context decode win

Branch: fix/phi35-onnx-attention-gqa-export · Head: 23564ff

Performance numbers below are re-derived to the digit from the source CSVs (see Sources) and
are consistent with the companion PERF_REPORT.md and GQA_GENERALIZATION_REPORT.md.
Llama-3.2-1B generalization is QA-signed (@b5d02a20).

Part 1 — Correctness (what this PR ships)

This PR makes the opset-24 ONNX Attention export actually run on the CUDA EP with an
external KV-cache, and brings it to numerical parity with the HF reference. Four fixes:

(a) Phase-split decode: `is_causal=0` + explicit causal mask + `If`-op

Broken path: the opset-24 Attention CUDA kernel returns NOT_IMPLEMENTED for
is_causal=1 combined with an external (pre-allocated) KV-cache — so the natural
"causal Attention" export simply will not load/run on CUDA.

Fix: export is_causal=0 and supply causality explicitly, split by phase behind an
If op keyed on whether we are prefilling or decoding:

decode (else-branch): a single maskless Flash Attention — at decode the query is one
token attending to all cached keys, so no mask tensor is needed.
prefill (then-branch): a 13-node causal-mask build feeding a masked Attention.

This sidesteps the unimplemented kernel path and, as a bonus, gives the decode branch the
fast maskless Flash kernel (the source of the Part 2 performance win).

(b) `df203cc` — GQA fp16 dtype fix

GQA op received a dtype-mismatched input under fp16 export; this corrects the dtype so the
grouped-query path produces correct fp16 results (precondition for parity).

(c) `cf6c5c4` — present-KV metadata

Fixes the present/updated KV-cache output metadata so the exported graph advertises the
correct cache tensors (shapes/names) for the runtime to thread through decode steps.

(d) `71e84b3` — lean strip

Strips dead/unused graph structure left over from the export, reducing the model to the
nodes actually exercised.

Parity result

Phi-3.5-mini: primary gate attn_dynamic ↔ gqa PASS, 20/20 tokens, no divergence
(parity_gqa_canonical_199.json).
Llama-3.2-1B: 20/20 for all variants (gqa, attn_static, attn_static_mea) vs the
attn_dynamic reference (parity_llama.json).

Both models reproduce the HF reference token stream (19–20/20 vs reference across the
variant set), so the export is numerically correct, not just runnable.

Part 2 — Measured value (why the phase-split is worth it)

The phase-split decode path is not just a workaround — it delivers a long-context decode
performance win: the static-cache Attention decode crosses over and beats the GQA op at
long context.

Phi-3.5-mini — MHA, 32Q/32KV, head_dim 96 (QA-signed, safe to publish)

Per-token decode latency, median ms (source: profile_4variant_matrix.csv,
profile_extended_8320.csv):

seqlen	static (ms)	gqa (ms)	ratio static/gqa
128	10.398	8.095	1.28×
512	10.620	8.613	1.23×
2048	11.007	10.361	1.06×
4096	11.300	12.529	0.90×
8192	14.190	16.698	0.850×

Static crosses 1.00 between seq 2048 and 4096 and peaks at 0.850× @8192 (8320 matched
buffer). GQA wins only short context (≤512), where the static TensorScatter buffer-tax
dominates.

Llama-3.2-1B — true GQA, 32Q/8KV (group-4), head_dim 64 (QA-signed @b5d02a20)

Per-token decode latency, median ms (source: profile_llama_matrix.csv):

seqlen	static (ms)	gqa (ms)	ratio static/gqa
128	3.569	3.245	1.100×
512	3.861	3.690	1.046×
2048	3.959	4.117	0.962×
4096	4.043	4.343	0.931×
8192	4.131	4.818	0.857×

The crossover holds on true GQA: 0.857× @8192 is nearly identical to Phi MHA's 0.850×,
and it arrives ~one band earlier (between seq 512 and 2048 vs Phi's 2048–4096) because
Llama-1B's absolute decode latencies are ~3× smaller, so the fixed buffer-tax amortizes
sooner. This proves the win is architecture-robust — across MHA and true-GQA, and
across two head_dims (96 and 64) — even though GQA's structural 4× KV-cache reduction is in
play. (The ~2048 crossover is conservative: static's ≤4096 cells run on the over-provisioned
4160 buffer paying the full TensorScatter tax; right-sizing the buffer would push it earlier.)

Mechanism — a fixed-cost vs grows-with-context tradeoff (not a kernel-quality gap)

The crossover is not "one variant uses a better kernel." CUPTI confirms both the static
and GQA decode paths run genuine Flash kernels (flash_fwd[_splitkv]). So the difference is
purely in how each path's cost scales with context:

ONNX Attention (static cache) is context-FLAT: +16% over a 64× context increase
(Llama 3.569 → 4.131 ms). Per step it Flash-reads a fixed-size, pre-allocated KV buffer and
writes the new token with a constant ~2.2 ms TensorScatter cache update — neither cost
depends on the actual sequence length, so the curve is nearly flat.
GQA op (dynamic cache) GROWS: +48% over the same range (Llama 3.245 → 4.818 ms). Its KV
read scales with the actual seqlen, so per-step cost rises as context grows.
The crossover is where these two lines meet. Short context → static's fixed buffer-tax
dominates and GQA wins (≤512). Long context → GQA's grown read-cost overtakes static's flat
line and ONNX Attention wins (0.857× @8192). The flat-vs-rising geometry is why the win
widens with context rather than being a single lucky point.
The MEA control never crosses (monotonically worst at every seqlen). MEA shares the static
cache mechanics but decodes on the MEA kernel instead of maskless Flash — so its failure to
cross isolates the phase-split maskless-Flash-on-decode branch as the active ingredient
behind the win, not the static cache alone.
Conservative bound: static's ≤2048 cells are measured on the over-provisioned 4160 buffer
(buffer ≫ real context), so they pay the full TensorScatter buffer-tax. Right-sizing the
buffer to real context would reduce that tax and push the crossover even earlier — the
reported ~2048 Llama crossover is therefore a conservative lower bound.

Export recipes — how to generate each attention variant (and which to pick)

All three shippable variants come from the same mobius build invocation differing only in
--execution-provider and --static-cache (flag strings verified against this session's export
provenance, owner-attested @798f0e27):

Variant	mobius export flags	Attention op	KV-cache	Decode kernel	Performance regime
`attn_static`	`--ep default --static-cache --max-seq-len {N}`	opset-24 ONNX `Attention` + phase-split `If`	fixed pre-allocated buffer (in-graph `TensorScatter`)	maskless Flash	Long-context winner. Context-FLAT (+16% over 64× ctx); crosses over and beats GQA at seq ≥ 2048 (0.857× @8192). Slightly slower at short ctx (fixed-buffer tax).
`gqa`	`--ep cuda` (no `--static-cache`)	contrib `com.microsoft::GroupQueryAttention`	dynamic	Flash (`flash_fwd_splitkv`; @128 = non-split `flash_fwd`)	Short-context winner (≤512). Decode grows +48% with ctx. On a true-GQA model also gets the structural 4× KV-cache reduction.
`attn_dynamic`	`--ep default` (no `--static-cache`)	opset-24 ONNX `Attention` (dynamic shapes)	dynamic concat-grow	MEA (memory-efficient; no maskless-Flash path)	Most portable/flexible, slowest decode. Grows with ctx (the MEA control, 2.88× @8192). Serves as the parity baseline + flexible-deploy option.

--max-seq-len sizes the static buffer: 4160 holds contexts ≤ 4096; use 8320 for 8192.

Key takeaway: the --static-cache flag is what unlocks the long-context crossover win — it
enables the phase-split → maskless-Flash decode on a fixed buffer. Pick attn_static for
long-context decode, gqa for short-context or when you want true KV-reduction, and
attn_dynamic for maximum portability.

Honest tradeoffs

+17% node count is intentional, not bloat. The benchmark graph grows ~17% because of the
phase-split If-op structure (the 13-node prefill mask-build branch + the maskless decode
branch). That structure is exactly what buys the crossover — it routes decode to the fast
maskless Flash kernel. It is a deliberate correctness+performance feature, not accidental
expansion.
May need a node-count re-baseline / waiver. Any CI guard that asserts on graph node count
will trip on this +17%. Recommend re-baselining the expected node count (or granting a waiver)
with a note pointing at the phase-split rationale.

Testing / validation

Parity gates: Phi-3.5 attn_dynamic ↔ gqa PASS 20/20; Llama-3.2-1B 20/20 all variants
vs attn_dynamic reference. Token streams match the HF reference.
Kernel verification (CUPTI, fail-closed): every decode row's actual launched kernel was
confirmed — static/gqa = Flash (flash_fwd_splitkv); mea/dynamic = MEA
(attention_kernel_batched_impl). No silent fallback, so the normalized ratios are valid.
Precision note: gqa decode @128 routes to the non-split flash_fwd_kernel (standard
short-KV Flash heuristic), switching to flash_fwd_splitkv at seq ≥ 512 — still a genuine
Flash kernel, and @128 is not crossover-critical (gqa wins there), so the conclusion is
unaffected.
QA sign-off: Phi-3.5 matrix QA-signed (published PERF_REPORT.md); Llama-3.2-1B matrix
QA-signed @b5d02a20 (PASS — five-item reconciliation: ratios 0.962/0.931/0.857 exact,
crossover @2048 robust with static p90 < gqa median at every seq ≥ 2048, kernel provenance
fail-closed, parity 20/20, phase-split delta).

Sources

Phi-3.5: /datadisks/disk1/titaiwang/phi35_bench/results/ — MATRIX_REPORT.md,
profile_4variant_matrix.csv, profile_extended_8320.csv, parity_gqa_canonical_199.json.
Llama-3.2-1B: /datadisks/disk1/titaiwang/llama32_bench/results/ —
MATRIX_REPORT_LLAMA.md, profile_llama_matrix.csv, parity_llama.json.
Companion reports: PERF_REPORT.md (Phi-3.5 decode-latency report, published),
GQA_GENERALIZATION_REPORT.md (two-model generalization, QA-signed).
Build: mobius HEAD 23564ff (+ df203cc packed-QKV/dtype fix).

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>

…4 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>

…sh/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>

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>

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>

…ty) 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>

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>

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>

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>

…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>

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>

…uard) 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>

…reement 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>

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>

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>

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>

…n 5) 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>

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>

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>

…lt 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>

…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>

…oss fixes 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>

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>

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>

codecov · 2026-06-02T00:05:12Z

The author of this PR, titaiwangms, is not an activated member of this organization on Codecov.
Please activate this user on Codecov to display this PR comment.
Coverage data is still being uploaded to Codecov.io for purposes of overall coverage calculations.
Please don't hesitate to email us at support@codecov.io with any questions.

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>

github-actions · 2026-06-02T00:06:46Z

🏗️ Architecture Diff

Comparing 8d492d9 → 856b986

Model	Sub-model	Changes	Status
bert (feature-extraction)	model	0	⚪
falcon	model	1	🔴
gemma2	model	1	🔴
gemma4 (gemma4)	decoder	0	⚪
gemma4 (gemma4)	embedding	0	⚪
gemma4 (gemma4)	vision_encoder	0	⚪
gemma4_text	model	1	🔴
gpt2	model	1	🔴
llama	model	1	🔴
llama (static-cache)	model	14	🟡
mamba (ssm-text-generation)	model	0	⚪
phi3	model	1	🔴
phi3 (static-cache)	model	20	🟡
qwen	model	1	🔴
qwen (static-cache)	model	14	🟡
qwen2	model	1	🔴
qwen2 (static-cache)	model	14	🟡
qwen2_moe	model	1	🔴
qwen2_moe (static-cache)	model	24	🟡
qwen3	model	1	🔴
qwen3 (static-cache)	model	15	🟡
qwen3_5_moe (hybrid-text-generation)	model	0	⚪
qwen3_5_text (hybrid-text-generation)	model	0	⚪
qwen3_5_vl (hybrid-qwen-vl)	decoder	0	⚪
qwen3_5_vl (hybrid-qwen-vl)	embedding	0	⚪
qwen3_5_vl (hybrid-qwen-vl)	vision_encoder	0	⚪
qwen3_moe	model	1	🔴
qwen3_moe (static-cache)	model	20	🟡
qwen3_next (hybrid-text-generation)	model	0	⚪
t5 (seq2seq)	decoder	0	⚪
t5 (seq2seq)	encoder	0	⚪
whisper (speech-to-text)	decoder	0	⚪
whisper (speech-to-text)	encoder	0	⚪

falcon / model — 1 change(s)

Op summary: 66 → 66 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

gemma2 / model — 1 change(s)

Op summary: 107 → 107 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

gemma4_text / model — 1 change(s)

Op summary: 129 → 129 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

gpt2 / model — 1 change(s)

Op summary: 53 → 53 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

llama / model — 1 change(s)

Op summary: 61 → 61 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

llama (static-cache) / model — 14 change(s)

Op summary: 58 → 68 nodes

--- base
+++ head
@@ -12,7 +12,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add
@@ -38,7 +43,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add

Added nodes:

+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If
+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If

Removed nodes:

- Attention
- Attention

phi3 / model — 1 change(s)

Op summary: 59 → 59 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

phi3 (static-cache) / model — 20 change(s)

Op summary: 56 → 66 nodes

--- base
+++ head
@@ -12,7 +12,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add
@@ -37,7 +42,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add

Added nodes:

+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If
+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If

Removed nodes:

- Attention
- Attention

Connectivity changes:

node[50] Transpose: input_ids [26] → [23]
node[51] MatMul: input_ids [84, 85] → [79, 80]
node[52] Add: input_ids [76, 86] → [62, 81]
node[53] RMSNormalization: input_ids [87, 27] → [82, 24]
node[54] Transpose: input_ids [28] → [25]
node[55] MatMul: input_ids [88, 89] → [83, 84]

qwen / model — 1 change(s)

Op summary: 61 → 61 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

qwen (static-cache) / model — 14 change(s)

Op summary: 58 → 68 nodes

--- base
+++ head
@@ -12,7 +12,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add
@@ -38,7 +43,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add

Added nodes:

+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If
+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If

Removed nodes:

- Attention
- Attention

qwen2 / model — 1 change(s)

Op summary: 61 → 61 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

qwen2 (static-cache) / model — 14 change(s)

Op summary: 58 → 68 nodes

--- base
+++ head
@@ -12,7 +12,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add
@@ -38,7 +43,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add

Added nodes:

+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If
+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If

Removed nodes:

- Attention
- Attention

qwen2_moe / model — 1 change(s)

Op summary: 224 → 224 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

qwen2_moe (static-cache) / model — 24 change(s)

Op summary: 214 → 224 nodes

--- base
+++ head
@@ -15,7 +15,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add
@@ -119,7 +124,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add

Added nodes:

+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If
+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If

Removed nodes:

- Attention
- Attention

Modified attributes:

node[129] Constant: value_int: None → 1, value_ints: [2] → None

Connectivity changes:

node[63] Mul: input_ids [130, 131] → [124, 129]
node[79] Mul: input_ids [146, 147] → [140, 145]
node[95] Mul: input_ids [162, 163] → [156, 161]
node[101] Transpose: input_ids [37] → [35]
node[102] MatMul: input_ids [89, 170] → [92, 168]
node[110] Add: input_ids [178, 40] → [163, 176]
node[170] Mul: input_ids [239, 241] → [210, 237]
node[186] Mul: input_ids [255, 257] → [210, 253]
node[202] Mul: input_ids [271, 273] → [210, 269]

qwen3 / model — 1 change(s)

Op summary: 73 → 73 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

qwen3 (static-cache) / model — 15 change(s)

Op summary: 70 → 80 nodes

--- base
+++ head
@@ -18,7 +18,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add
@@ -50,7 +55,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add

Added nodes:

+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If
+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If

Removed nodes:

- Attention
- Attention

Connectivity changes:

node[47] Reshape: input_ids [84, 18] → [79, 15]

qwen3_moe / model — 1 change(s)

Op summary: 202 → 202 nodes

No op-sequence changes.

Interface changes:

output[1]: shape [] → ['?', '2', '?', '16']; output[2]: shape [] → ['?', '2', '?', '16']; output[3]: shape [] → ['?', '2', '?', '16']; output[4]: shape [] → ['?', '2', '?', '16']

qwen3_moe (static-cache) / model — 20 change(s)

Op summary: 192 → 202 nodes

--- base
+++ head
@@ -18,7 +18,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add
@@ -111,7 +116,12 @@
 RotaryEmbedding
 TensorScatter
 TensorScatter
-Attention
+Shape
+Constant
+Squeeze
+Constant
+Greater
+If
 Transpose
 MatMul
 Add

Added nodes:

+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If
+ Shape
+ Constant
+ Squeeze
+ Constant
+ Greater
+ If

Removed nodes:

- Attention
- Attention

Modified attributes:

node[121] Constant: value_int: None → 1, value_ints: [2] → None

Connectivity changes:

node[66] Mul: input_ids [125, 126] → [119, 124]
node[82] Mul: input_ids [141, 142] → [135, 140]
node[108] Reshape: input_ids [167, 18] → [162, 15]
node[162] Mul: input_ids [223, 225] → [194, 221]
node[178] Mul: input_ids [239, 241] → [194, 237]

Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

…control 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>

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>

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>

Copilot

Pull request overview

This PR addresses two export/runtime correctness gaps in mobius’ ONNX attention stack: (1) static-cache (TensorScatter) Attention on opset-24 CUDA where is_causal=1 is rejected and decode must remain Flash-eligible, and (2) fp16 GQA export issues where folded packed-QKV / transposed initializers could end up with incorrect/unknown dtype metadata and/or prevent dead-weight stripping, plus explicit present-KV metadata stamping to avoid shape-inference mismatches.

Changes:

Add a phase-split static-cache attention path (prefill masked vs decode maskless) and comprehensive structural + CUDA runtime regression tests to guard kernel routing and correctness.
Introduce a shared initializer_dtype() helper and wire it into fold passes to preserve fp16 dtype through folded initializers; add regression tests including serialization round-trips and ORT loadability.
Stamp explicit present.* KV-cache output shapes/dtypes (opt-in) symmetric to past inputs to work around GQA contrib-op shape inference; add targeted tests and update changelog/docs.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/static_cache_decode_test.py	New CUDA end-to-end regression suite for static-cache Attention (runnability, maskless decode, Flash vs MEA kernel selection, and nonpad bounds semantics).
tests/build_graph_test.py	Adds recursive node-walk + fail-closed structural assertions ensuring static-cache Attention is behind per-layer `If` with masked prefill / maskless decode and `is_causal=0`.
src/mobius/tasks/_task_test.py	Adds a graph-level test ensuring `present.*` KV outputs match past-input metadata (dims + dtype).
src/mobius/tasks/_causal_lm.py	Computes explicit key/value head dims and stamps present KV output metadata via `_register_kv_cache_outputs(...)` to avoid GQA present head-dim mismatches.
src/mobius/tasks/_cache_utils.py	Extends `_register_kv_cache_outputs` with optional all-or-nothing explicit present shape/dtype stamping + warning on partial parameter sets.
src/mobius/tasks/_cache_utils_test.py	New unit tests covering present KV stamping behavior (override wrong inference, opt-out silence, partial-params warning, naming).
src/mobius/components/_common.py	Adds `create_static_cache_causal_mask()` producing a 4D per-batch causal mask for static-cache attention.
src/mobius/components/_common_test.py	Adds value-level CPU tests validating the static-cache causal mask structure and semantics.
src/mobius/components/_attention.py	Introduces `_attend_over_static_cache()` and updates static-cache attention path to `is_causal=0` with `If` phase split (prefill masked / decode maskless) to preserve decode Flash routing.
src/mobius/_passes/_fold_transpose.py	Uses `initializer_dtype()` to ensure folded-transpose initializers preserve fp16 dtype when declared dtype metadata is missing.
src/mobius/_passes/_fold_transpose_test.py	Adds regression test for dtype resolution when declared dtype is missing but `const_value` is fp16.
src/mobius/_passes/_fold_concat.py	Uses `initializer_dtype()`, skips folding when weights aren’t loaded, and removes Concat nodes with `safe=True` so DCE can prune dead pre-pack weights.
src/mobius/_passes/_fold_concat_test.py	Adds multiple regression tests for dtype preservation, DCE detachment, ORT loadability, and serialization value integrity (note: currently introduces an `onnx` dependency).
src/mobius/_passes/_dtype_utils.py	New shared helper `initializer_dtype()` for pass authors to resolve effective initializer dtype from declared metadata and/or `const_value`.
src/mobius/_passes/_dtype_utils_test.py	Tests for `initializer_dtype()` behavior including disagreement warning and missing metadata fallback.
CHANGELOG.md	Documents the fp16 GQA folded-initializer dtype fix and its user-visible impact.
.agents/skills/mobius-onnx-export-gotchas/SKILL.md	Adds an internal skill doc summarizing ONNX export pitfalls and verification guidance for fp16 GQA and static-cache.

titaiwangms · 2026-06-03T17:26:54Z

Review synthesis (4-reviewer fan-out: readability, code, critical, deep)

Overall: Solid, well-tested change. Core math/spec was verified against the installed opset-24 Attention schema and found sound — mask polarity (True=attend), 4D rank, the <= causal relation, present/past KV symmetry, and phase-split equivalence are all correct. No Critical issues. The substantive findings cluster around API robustness / silent-failure modes rather than the happy path.

Major

1. Decode is maskless and trusts nonpad_kv_seqlen with no backstop — src/mobius/components/_attention.py (decode branch)
Decode correctness depends entirely on the runtime supplying nonpad_kv_seqlen == write_indices + S_q. Prefill is robust (the explicit mask bounds it regardless); decode has no such guard. If a consumer (e.g. onnxruntime-genai) passes nonpad = allocated cache length, or interprets it as "valid length before this call," decode silently attends to padding/unwritten slots → wrong logits, no error. This is the highest-priority finding because the failure is silent and strictly worse than prefill's.
→ Derive the bound in-graph where possible (write_indices + Shape(query)[1]), or at minimum document the contract loudly at the decode branch. (Filed as a question pending confirmation the genai driver guarantees the contract.)

2. Ragged batched prefill can attend to padded cache slots — create_static_cache_causal_mask, src/mobius/components/_common.py
The mask keeps j <= write_indices[b] + t but does not combine nonpad_kv_seqlen. The docstring's claim that padding slots are "always greater" only holds when every batch row has exactly S_q real tokens. With prompt lengths [2,4] padded to S_q=4, row 0's padded queries attend to slots written from padded input.
→ AND the mask with j < nonpad_kv_seqlen[b] to fail closed. (Matters only if batched ragged prefill is a supported use case — worth confirming scope.)

3. Partial present-shape stamp fails open into the known-wrong path — src/mobius/tasks/_cache_utils.py (flagged independently by two reviewers)
A partial param set only warns, then falls back to GQA inference that mis-declares head_dim (the exact bug this PR fixes). A future caller dropping one of six params ships bad metadata with only a log line.
→ raise ValueError on provided and not stamp instead of warning.

Minor

Per-layer branch prefixes are constant (static_cache_prefill_ / static_cache_decode_) across all layers — potential name collision if downstream tooling flattens subgraph namespaces. Consider a per-call unique suffix.
_capture_attention_kernel_log (tests/static_cache_decode_test.py) restores logger severity to a hardcoded WARNING instead of the prior value → order-dependent tests.
_outputs=3 on the static-cache Attention requests present_* outputs alongside nonpad_kv_seqlen, which the opset-24 schema prose says shouldn't be combined; the present outputs are discarded anyway. Pre-existing; consider _outputs=1 if ORT's external-cache kernel doesn't require 3.
Flash-kernel CUDA test is version-gated to ORT 1.27.x → no kernel-selection coverage on 1.28+. The structural maskless test remains version-robust, so acceptable, but note the gap.

Resolved by tie-breaker: one reviewer argued initializer_dtype()'s "const_value wins on disagreement" should raise rather than warn. The deep/spec review ruled it safe — const_value is the data actually serialized, so matching the declared type to it is correct, and the disagreement path now warns. Recommend keeping as-is.

Nits (readability)

stamp → stamp_explicit_shape; trim the over-long ordering comment in _build_attention_branch (delegate to rename_subgraph_values docstring); replace raw commit SHAs in SKILL.md headers with #328/version; dead float | None on softcap; redundant kv_ prefix on locals; opaque has_mask one-liner; exact-count GreaterOrEqual assertion is brittle.

Praise

Phase-split-maskless-decode is the right perf design; the 4D mask correctly handles per-batch write_indices (a 3D mask would be a batched-correctness trap); tests verify mask values by running the subgraph on ORT rather than asserting structure alone, and include negative controls; initializer_dtype() centralization is a clean fix that reduces drift across passes.

_{🤖 Generated by a multi-agent review fan-out (readability + code + critical + deep reviewers).}

justinchuby · 2026-06-03T17:37:13Z

+### 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.
+
+### 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.
+


When these are fixed, do we still need them in the skill?

Fair point — mostly yes, but they should be reframed. A skill's job is to prevent future regressions, not to track the status of a fixed bug, so rather than delete these we'd split them:

Compress the bug-narrative: the 'Status: FIXED as of <commit>' banners and pre-fix symptom traces shrink to a one-line historical pointer (they retain a little value for interpreting pre-fix artifacts, but that's it).

Keep + promote the convention: the rule that any pass materializing an initializer must use initializer_dtype() (never value.dtype or FLOAT) is the actual regression guard — it still applies and even flags sibling passes that aren't yet covered.

Keep the value-based weight-equivalence check (correlation/norm, not dtype/count) as enduring methodology, and the 'known remaining' item (still unfixed).

Alternatively, if you'd rather the skill hold only live gotchas, we could relocate the durable conventions into a CONTRIBUTING/_passes doc and trim the skill accordingly — happy to go either way.

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>

titaiwangms · 2026-06-03T17:37:29Z

Addressed the CONTRIBUTING.md "zero protobuf operations" review finding in src/mobius/_passes/_fold_concat_test.py.

New SHA: 96ef1b1 (fast-forward on 5de4345)

What changed (test-only, behavior-preserving):

Replaced onnx.save(ir.to_proto(model), str(model_path)) with the IR-native ir.save(model, model_path) pattern already used elsewhere in the same file.
Removed the now-unused import onnx.

The test still writes the model and loads it in ONNX Runtime, asserting the fp16 result shape/dtype and the fp16/fp32 MatMul type-mismatch regression guard. Verified: lintrunner clean and pytest _fold_concat_test.py → 18 passed.

justinchuby · 2026-06-03T17:38:59Z

+        branch.outputs.append(attn_output)
+        return branch
+
+    prefill_branch = _build_attention_branch("static_cache_prefill", use_causal_mask=True)


What does this mean for performance? Should we really look into having separate prefill and decode models?

Good question — we profiled this, and the phase-split If is intentional and net-positive, which is why we don't think separate prefill/decode models are needed.

What the If does: If(Greater(seq_len, 1)) routes per step — prefill takes the explicit-causal-mask branch (MEA kernel), decode takes the maskless branch (Flash). ORT runs only the taken branch and the predicate is a host-side scalar compare, so it's not a per-token GPU cost. That's the same prefill/decode specialization two separate models would give, but in one graph.

Is it worth it? Yes — measured. Our control (static cache with the mask forced on, i.e. no split) is the slowest variant at every length (the phase-split is ~27% faster at 4096; equivalently the always-masked control is ~37% slower). The maskless-Flash decode the split unlocks is the source of a long-context win: with the static KV buffer sized to the real max context, static Attention decode crosses over and beats the contrib GroupQueryAttention op at long context (~0.90x at 4096, ~0.85x at 8192 on Phi-3.5-mini MHA; reproduced on a true-GQA model, Llama-3.2-1B, at 0.857x @8192). Honest caveat: this depends on buffer sizing — over-provisioning the static buffer erases the win, and the crossover point is model/buffer-dependent (interpolated N*~2.7K on Phi, ~2K on Llama), so it's not 'static always wins past 4K.' GQA still wins short context (<=2048 on Phi, <=512 on Llama).

Why one graph, not two models: prefill and decode share the same pre-allocated key_cache/value_cache buffers, updated in place via TensorScatter at write_indices (0 for prefill, current position for decode). One graph preserves that in-place KV-cache I/O contract — you hand the same OrtValues straight from prefill into each decode step. Two models would duplicate weights (~2x) and force cross-session buffer re-binding, for no kernel-level benefit. (A fuller decode-latency write-up is forthcoming.)

…test 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>

justinchuby · 2026-06-03T18:59:08Z

I think we need to update the architectural diff to look at subgraphs. can be separate.

justinchuby

LGTM overall. Would need some more eyes preferably.

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>

…hangelog 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>

github-actions · 2026-06-03T23:13:29Z

Performance Comparison

Comparing 8d492d9 → 856b986

Model	Metric	Baseline	Current	Delta
bert (feature-extraction)	model_size_bytes	359 KB	359 KB	+0.0%	⚪
bert (feature-extraction)	num_nodes	60	60	+0.0%	⚪
falcon	model_size_bytes	364 KB	364 KB	+0.0%	⚪
falcon	num_nodes	66	66	+0.0%	⚪
gemma2	model_size_bytes	428 KB	428 KB	+0.0%	⚪
gemma2	num_nodes	107	107	+0.0%	⚪
gpt2	model_size_bytes	388 KB	388 KB	+0.0%	⚪
gpt2	num_nodes	53	53	+0.0%	⚪
llama	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama	num_nodes	61	61	+0.0%	⚪
llama (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama (static-cache)	num_nodes	58	68	+17.2%	🟦
mamba (ssm-text-generation)	model_size_bytes	296 KB	296 KB	+0.0%	⚪
mamba (ssm-text-generation)	num_nodes	98	98	+0.0%	⚪
phi3	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3	num_nodes	59	59	+0.0%	⚪
phi3 (static-cache)	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3 (static-cache)	num_nodes	56	66	+17.9%	🟦
qwen2	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2	num_nodes	61	61	+0.0%	⚪
qwen2 (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2 (static-cache)	num_nodes	58	68	+17.2%	🟦
qwen3_5_moe (hybrid-text-generation)	model_size_bytes	506 KB	506 KB	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	num_nodes	275	275	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	model_size_bytes	458 KB	458 KB	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	num_nodes	129	129	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	model_size_bytes	977 KB	977 KB	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	num_nodes	413	413	+0.0%	⚪
t5 (seq2seq)	model_size_bytes	836 KB	836 KB	+0.0%	⚪
t5 (seq2seq)	num_nodes	166	166	+0.0%	⚪
whisper (speech-to-text)	model_size_bytes	1008 KB	1008 KB	+0.0%	⚪
whisper (speech-to-text)	num_nodes	128	128	+0.0%	⚪

No blocking regressions.

🟦 = intended structural change accepted via EXPECTED_CHANGES (exact pinned base→current values; see tests/benchmark_compare.py).

…retention 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>

…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>

…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>

…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>

… not in-branch changes 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>

titaiwangms · 2026-06-05T21:24:30Z

Closing this PR — it is superseded by #340.

Why: #328's static-cache attention used a per-layer If(Greater(S_q, 1))
phase-split to route between prefill and decode. An in-graph control-flow If
makes the exported model fail to load under ONNX Runtime graph capture: ORT
hard-fails session initialization when the graph contains control-flow nodes
(inference_session.cc, HasControlflowNodes). This gate is
execution-provider-agnostic — empirically confirmed on CUDA
(enable_cuda_graph=1 → session-init FAIL; the same model loads fine in eager),
and DirectML (always captures) / TensorRT-RTX are rejected by the same
EP-agnostic gate (by the same control-flow check, not separately runtime-tested).
Since mobius ships models for onnxruntime / onnxruntime-genai
consumers that run under graph capture, an eager-only model isn't viable.

What carries forward in #340: all of this PR's If-independent fixes — the
fp16 GQA packed-weight dtype fold-fix (+ e2e tests), the GQA present-KV
head_dim metadata fix, the fold-concat value-integrity tests, dead pre-pack
stripping, and the dtype-warning test. #340 also replaces the phase-split
with a branchless, capture-safe static-cache attention: is_causal=0 + an
explicit offset-aware causal mask (decode and prefill run on the
Memory-Efficient Attention kernel; maskless-Flash decode is deferred to a future
host-split export).

Thanks to everyone who reviewed here — continuing in #340.

Copilot AI and others added 24 commits June 1, 2026 22:24

titaiwangms and others added 3 commits June 2, 2026 00:13

Merge branch 'main' into fix/phi35-onnx-attention-gqa-export

5de4345

titaiwangms marked this pull request as ready for review June 2, 2026 23:16

titaiwangms requested review from a team and Copilot June 2, 2026 23:16

Copilot started reviewing on behalf of titaiwangms June 2, 2026 23:17 View session

Copilot AI reviewed Jun 2, 2026

View reviewed changes

Comment thread src/mobius/_passes/_fold_concat_test.py

Comment thread src/mobius/_passes/_fold_concat_test.py

justinchuby reviewed Jun 3, 2026

View reviewed changes

Comment thread tests/build_graph_test.py Outdated

justinchuby reviewed Jun 3, 2026

View reviewed changes

Comment thread tests/build_graph_test.py Outdated

titaiwangms mentioned this pull request Jun 3, 2026

CI: add numerical parity coverage for the static-cache export path #329

Closed

justinchuby approved these changes Jun 3, 2026

View reviewed changes

titaiwangms and others added 2 commits June 3, 2026 19:03

titaiwangms mentioned this pull request Jun 3, 2026

Architecture diff: recurse into subgraphs (make control-flow/phase-split changes visible) #330

Closed

titaiwangms and others added 5 commits June 3, 2026 23:31

titaiwangms mentioned this pull request Jun 5, 2026

fp16/GQA export fixes + capture-safe static-cache attention (supersedes #328) #340

Closed

titaiwangms closed this Jun 5, 2026

Uh oh!

Conversation

titaiwangms commented Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Fix opset-24 Attention GQA export (phase-split decode) + long-context decode win

Part 1 — Correctness (what this PR ships)

(a) Phase-split decode: is_causal=0 + explicit causal mask + If-op

(b) df203cc — GQA fp16 dtype fix

(c) cf6c5c4 — present-KV metadata

(d) 71e84b3 — lean strip

Parity result

Part 2 — Measured value (why the phase-split is worth it)

Phi-3.5-mini — MHA, 32Q/32KV, head_dim 96 (QA-signed, safe to publish)

Llama-3.2-1B — true GQA, 32Q/8KV (group-4), head_dim 64 (QA-signed @b5d02a20)

Mechanism — a fixed-cost vs grows-with-context tradeoff (not a kernel-quality gap)

Export recipes — how to generate each attention variant (and which to pick)

Honest tradeoffs

Testing / validation

Sources

Uh oh!

codecov Bot commented Jun 2, 2026

Uh oh!

github-actions Bot commented Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🏗️ Architecture Diff

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

titaiwangms commented Jun 3, 2026

Review synthesis (4-reviewer fan-out: readability, code, critical, deep)

Major

Minor

Nits (readability)

Praise

Uh oh!

justinchuby Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

titaiwangms Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

titaiwangms commented Jun 3, 2026

Uh oh!

justinchuby Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

titaiwangms Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Jun 3, 2026

Uh oh!

justinchuby left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance Comparison

Uh oh!

titaiwangms commented Jun 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

titaiwangms commented Jun 2, 2026 •

edited

Loading

(a) Phase-split decode: `is_causal=0` + explicit causal mask + `If`-op

(b) `df203cc` — GQA fp16 dtype fix

(c) `cf6c5c4` — present-KV metadata

(d) `71e84b3` — lean strip

github-actions Bot commented Jun 2, 2026 •

edited

Loading

titaiwangms Jun 3, 2026 •

edited

Loading

github-actions Bot commented Jun 3, 2026 •

edited

Loading