fp16/GQA export fixes + capture-safe static-cache attention (supersedes #328) - #340
Closed
titaiwangms wants to merge 30 commits into
Closed
fp16/GQA export fixes + capture-safe static-cache attention (supersedes #328)#340titaiwangms wants to merge 30 commits into
titaiwangms wants to merge 30 commits into
Conversation
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> (cherry picked from commit df203cc)
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> (cherry picked from commit 71e84b3)
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> (cherry picked from commit 35d08f6)
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> (cherry picked from commit 4e9bb5a)
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> (cherry picked from commit b3b08cc)
…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> (cherry picked from commit ec5afb9)
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> (cherry picked from commit 96ef1b1)
…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> (cherry picked from commit f74d812)
…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> (cherry picked from commit 110f26b)
…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> (cherry picked from commit 365e624)
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>
(cherry picked from commit cf6c5c4)
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> (cherry picked from commit 24fec65)
…gelog Re-authored from the abandoned #328 branch, carrying ONLY the If/phase-split-independent fp16 and GQA export guidance: - New skill `mobius-onnx-export-gotchas` documenting the fp16 GQA fold-pass fp32-corruption fix (df203cc), VALUE-based packed-QKV weight verification, and the GQA `present.*` head_dim shape fix (cf6c5c4). - CHANGELOG entry for the fp16 GQA fold-pass dtype fix (df203cc). No static-cache phase-split content is included (that work is dropped in the pivot). The exported attention path is unchanged from main. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Re-lint at the end of the salvage instead of cherry-picking the original branch's combined lintrunner commit (which spanned dropped phase-split files). Pure formatting / docstring-summary fixes; no behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… cherry-pick SHAs Triple-review doc-nits (all reviews PASSED, doc-only, in-scope): - CHANGELOG: re-add the "GQA Present KV-Cache Shape Fix" #### Fixed entry. The branch ships the present-KV head_dim fix (be84ece/98352ff, tasks/_cache_utils.py + _causal_lm.py) and SKILL.md section 6 documents it, but the salvage changelog previously carried only the fp16 entry. Metadata / declared-shape correction only; runtime numerics unchanged. No phase-split content introduced. - SKILL.md: de-anchor sections 3/5/6 from the volatile #328 cherry-pick SHAs (df203cc, cf6c5c4 — re-authored here as 7aaff4c/be84ece, and they change again on rebase/merge). Reference the fixes by name instead ("the fp16 GQA fold-fix" / "the GQA present-KV shape 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> (cherry picked from commit 140afae)
…tion Y) Re-expresses the intent of the dropped maskless-nonpad guard (b094af3) for the always-masked Option-Y design (is_causal=0 + explicit create_static_cache_causal_mask). Under Y the explicit mask is the ONLY thing bounding decode attention, so this test asserts decode at offset>0 attends only to slots within the causal frontier j <= write_indices + t (equivalently j < nonpad_kv_seqlen): - Negative (out-of-range) control: poisoning slots at/beyond the frontier leaves the decode logits bit-identical — the mask zeroes them. - Positive (in-range) control: poisoning an in-frontier slot changes the logits — proving decode genuinely attends in-range, so the negative control is non-vacuous. Does NOT port b094af3 as-is: that asserted a maskless decode reading exactly nonpad keys, which contradicts Y's masked semantics. Verified passing on CUDA. (Incidental: lintrunner reflow of the 140afae code in the same file, re-linted here since the original lintrunner commit 23564ff was not cherry-picked into the salvage.) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Re-lint at the end of the Option-Y addendum: 140afae's _common.py / _common_test.py predate the salvage branch's re-lint, and the original combined lintrunner commit (23564ff) spanned dropped phase-split files so was not cherry-picked. Pure ruff-format line reflow; no logic change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…trol flow) Owner-directed, durable guidance for anyone constructing model graphs in mobius (broader than the static-cache gotcha). New top-level section 7: - Principle: exported graphs must run under CUDA Graph capture (CUDA EP enable_cuda_graph, DML always-captures, NvTensorRtRtx, genai past_present_share_buffer), not just eager. - Hard constraint: ORT hard-FAILs session init when a capture-enabled EP loads a model with control-flow nodes — cites inference_session.cc HasControlflowNodes (If/Loop/Scan) and the FAIL message; notes it is branch-agnostic. - Empirically confirmed (CUDA enable_cuda_graph=1 fails on per-layer-If export, same model loads in eager); measurement caveat (eager per-node ~0 cost is not representative of the capture path). - Rule: avoid in-graph If/Loop/Scan; use host-side dispatch or branchless forms. Static-cache masking must be is_causal=0 + explicit offset-aware mask. - Cautionary example: the abandoned If(Greater(S_q,1)) phase-split vs the branchless is_causal=0 + explicit-mask path this PR ships. Also surfaces the capture rule in the skill frontmatter description. Factual, cites the ORT mechanism so a future reader can verify; no internal benchmark ratios or paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ndings 1/2/5) Triple-review doc findings on the Option-Y attention change: - CHANGELOG (readability Minor): add "### Static-cache Attention Causal-Mask Fix" #### Fixed entry — static cache now uses is_causal=0 + explicit offset-aware mask instead of is_causal=1 (ORT opset-24 causal_cross_no_past guard raised NOT_IMPLEMENTED on is_causal=1 + nonpad + S_q!=S_kv + no-past); branchless so capture-compatible; decode/prefill on MEA. The PR previously under-documented this shipped attention change. - SKILL.md (readability Minor): add the causal_cross_no_past gotcha->remedy to the static-cache section (§2) with a cross-reference to §7 (graph-capture), instead of duplicating §7's content. - _attention.py (readability Nit): one-clause comment in the static path noting the always-masked formulation deliberately trades decode-on-Flash for graph-capture compatibility (no If phase-split). Comment-only code change; static path stays branchless (0 control-flow ops). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ndings 3/4) - Nit (code review): fix wording that conflated the causal frontier with the padding boundary. The single-token decode guard exercises the *padding* side (j >= nonpad); reword its docstring/comment/assertion accordingly. - Minor (critical review): add test_static_cache_prefill_causal_mask_blocks_future_keys_within_nonpad_on_cuda, which isolates the *causal* side that a single-token decode cannot probe. It seeds slots 0..3, then runs a 2-token block at positions 1,2 (write_indices=1, nonpad=4 so slots 0..3 are all valid, not padding). Slot 3 is a written, within-nonpad key in the future of both query rows: poisoning it must NOT change the logits (only the causal mask, not the padding bound, can exclude a within-nonpad slot). Positive control: poisoning slot 0 (causal past, carried) must change the logits, so the guard is non-vacuous. Passes on CUDA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…causal_via_explicit_mask The test locks the static path to is_causal=0 + an explicit attn_mask (causality supplied via GreaterOrEqual), so the old name misleads a future reader into assuming is_causal=1. Semantics unchanged; name now matches what it asserts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
The author of this PR, titaiwangms, is not an activated member of this organization on Codecov. |
🏗️ Architecture Diff
falcon / model — 1 change(s)Op summary: 66 → 66 nodes No op-sequence changes. Interface changes:
gemma2 / model — 1 change(s)Op summary: 107 → 107 nodes No op-sequence changes. Interface changes:
gemma4_text / model — 1 change(s)Op summary: 129 → 129 nodes No op-sequence changes. Interface changes:
gpt2 / model — 1 change(s)Op summary: 53 → 53 nodes No op-sequence changes. Interface changes:
llama / model — 1 change(s)Op summary: 61 → 61 nodes No op-sequence changes. Interface changes:
llama (static-cache) / model — 17 change(s)Op summary: 58 → 74 nodes --- base
+++ head
@@ -1,3 +1,19 @@
+Constant
+Constant
+Shape
+Constant
+Squeeze
+Shape
+Constant
+Squeeze
+Range
+Range
+Unsqueeze
+Unsqueeze
+Add
+Unsqueeze
+Unsqueeze
+GreaterOrEqual
Gather
Gather
GatherAdded nodes:
Initializer changes:
phi3 / model — 1 change(s)Op summary: 59 → 59 nodes No op-sequence changes. Interface changes:
phi3 (static-cache) / model — 19 change(s)Op summary: 56 → 72 nodes --- base
+++ head
@@ -1,3 +1,19 @@
+Constant
+Constant
+Shape
+Constant
+Squeeze
+Shape
+Constant
+Squeeze
+Range
+Range
+Unsqueeze
+Unsqueeze
+Add
+Unsqueeze
+Unsqueeze
+GreaterOrEqual
Gather
Gather
GatherAdded nodes:
Connectivity changes:
Initializer changes:
qwen / model — 1 change(s)Op summary: 61 → 61 nodes No op-sequence changes. Interface changes:
qwen (static-cache) / model — 17 change(s)Op summary: 58 → 74 nodes --- base
+++ head
@@ -1,3 +1,19 @@
+Constant
+Constant
+Shape
+Constant
+Squeeze
+Shape
+Constant
+Squeeze
+Range
+Range
+Unsqueeze
+Unsqueeze
+Add
+Unsqueeze
+Unsqueeze
+GreaterOrEqual
Gather
Gather
GatherAdded nodes:
Initializer changes:
qwen2 / model — 1 change(s)Op summary: 61 → 61 nodes No op-sequence changes. Interface changes:
qwen2 (static-cache) / model — 17 change(s)Op summary: 58 → 74 nodes --- base
+++ head
@@ -1,3 +1,19 @@
+Constant
+Constant
+Shape
+Constant
+Squeeze
+Shape
+Constant
+Squeeze
+Range
+Range
+Unsqueeze
+Unsqueeze
+Add
+Unsqueeze
+Unsqueeze
+GreaterOrEqual
Gather
Gather
GatherAdded nodes:
Initializer changes:
qwen2_moe / model — 1 change(s)Op summary: 224 → 224 nodes No op-sequence changes. Interface changes:
qwen2_moe (static-cache) / model — 90 change(s)Op summary: 214 → 230 nodes --- base
+++ head
@@ -1,3 +1,19 @@
+Constant
+Constant
+Shape
+Constant
+Squeeze
+Shape
+Constant
+Squeeze
+Range
+Range
+Unsqueeze
+Unsqueeze
+Add
+Unsqueeze
+Unsqueeze
+GreaterOrEqual
Gather
Gather
GatherAdded nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
qwen3 / model — 1 change(s)Op summary: 73 → 73 nodes No op-sequence changes. Interface changes:
qwen3 (static-cache) / model — 17 change(s)Op summary: 70 → 86 nodes --- base
+++ head
@@ -1,3 +1,19 @@
+Constant
+Constant
+Shape
+Constant
+Squeeze
+Shape
+Constant
+Squeeze
+Range
+Range
+Unsqueeze
+Unsqueeze
+Add
+Unsqueeze
+Unsqueeze
+GreaterOrEqual
Gather
Gather
GatherAdded nodes:
Initializer changes:
qwen3_moe / model — 1 change(s)Op summary: 202 → 202 nodes No op-sequence changes. Interface changes:
qwen3_moe (static-cache) / model — 70 change(s)Op summary: 192 → 208 nodes --- base
+++ head
@@ -1,3 +1,19 @@
+Constant
+Constant
+Shape
+Constant
+Squeeze
+Shape
+Constant
+Squeeze
+Range
+Range
+Unsqueeze
+Unsqueeze
+Add
+Unsqueeze
+Unsqueeze
+GreaterOrEqual
Gather
Gather
GatherAdded nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
This was referenced Jun 5, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates mobius’ ONNX export pipeline to (1) make static-cache Attention capture-safe by removing in-graph control flow and expressing causality via an explicit offset-aware mask, and (2) fix multiple fp16/GQA export correctness issues (dtype preservation through folding and correct present-KV output metadata), with accompanying regression tests and documentation.
Changes:
- Static-cache Attention export now uses
is_causal=0plus an explicit offset-aware causal mask (create_static_cache_causal_mask) instead of relying onis_causal=1. - Fold passes now preserve fp16 dtype correctly when folding packed/transposed initializers (including when declared dtype metadata is missing), and remove dead pre-pack weights reliably.
- Dynamic-cache
present.*KV outputs are explicitly stamped to matchpast_key_values.*shapes/dtypes to avoid GQA shape-inference head_dim mis-declaration.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/static_cache_decode_test.py | New CUDA EP end-to-end regression tests for static-cache prefill/decode and mask correctness. |
| tests/build_graph_test.py | Updates static-cache graph-build assertions to require is_causal=0 and an explicit causal attn_mask. |
| src/mobius/tasks/_task_test.py | Adds a regression test ensuring present.* outputs match past_key_values.* metadata. |
| src/mobius/tasks/_causal_lm.py | Computes explicit KV head dims and stamps present.* output metadata for dynamic-cache builds. |
| src/mobius/tasks/_cache_utils.py | Extends _register_kv_cache_outputs to optionally stamp explicit output shapes/dtypes (all-or-nothing) with warning on partial wiring. |
| src/mobius/tasks/_cache_utils_test.py | New unit tests for _register_kv_cache_outputs stamping + warning behavior. |
| src/mobius/components/_common.py | Adds create_static_cache_causal_mask (branchless offset-aware bool mask for static-cache Attention). |
| src/mobius/components/_common_test.py | Adds structural and value-level CPU-evaluated tests for the static-cache causal mask. |
| src/mobius/components/_attention.py | Static-cache attention path now passes explicit causal mask and sets is_causal=0. |
| src/mobius/_passes/_fold_transpose.py | Preserves dtype for folded transposed initializers via initializer_dtype() (prevents silent fp16→fp32 widening). |
| src/mobius/_passes/_fold_transpose_test.py | Adds a regression test for folding transpose when declared dtype metadata is missing. |
| src/mobius/_passes/_fold_dtype_e2e_test.py | New end-to-end fp16 export tests covering folding + serialization round-trips. |
| src/mobius/_passes/_fold_concat.py | Preserves packed dtype via initializer_dtype(), skips folding before weights are loaded, and removes Concat nodes with safe=True to enable DCE. |
| src/mobius/_passes/_fold_concat_test.py | Adds regression tests for DCE detachment, dtype fallback from const_value, ORT-load, and value-integrity gates (incl. negative control). |
| src/mobius/_passes/_dtype_utils.py | New shared helper initializer_dtype() to resolve dtype from declared type or const_value (warn on mismatch). |
| src/mobius/_passes/_dtype_utils_test.py | Unit tests for initializer_dtype() behavior (declared-only, const-only, mismatch warning, none). |
| CHANGELOG.md | Documents the static-cache mask change, fp16 GQA fold fix, and GQA present-KV metadata fix. |
| .agents/skills/mobius-onnx-export-gotchas/SKILL.md | New export “gotchas” skill doc including graph-capture constraint and validation guidance. |
Contributor
Author
Review synthesis (4-reviewer team: readability, code, critical, deep)Verdict: Strong PR, no Critical issues. The causal-mask math and opset-24 Major
Minor
Nit
Praise
Synthesized from a 4-model review team (readability, code, critical, deep reviewers). |
…k contract) Addresses the in-scope items from the #340 triple-review assessment: M3 (must-fix): update stale is_causal=1 docstrings in _causal_lm.py — the static-cache input-doc and the Falcon-ALiBi limitations block now describe the shipped is_causal=0 + explicit offset-aware mask path. M1 (contract clarity): document the compact/right-trimmed cache invariant in create_static_cache_causal_mask — the mask is purely positional and never consumes nonpad_kv_seqlen, so correctness relies on nonpad == write_indices + S_q (padding right-side only). Ragged/left/interior padding is documented as out of scope. No defensive mask term added. M4 (coverage gap): add test_present_outputs_match_past_inputs_mla_distinct_head_dims exercising an MLA (DeepSeek-style) config where key_head_dim != value_head_dim, guarding the present-KV head_dim stamping fix for the case it was written for. Minors: - rename test_static_cache_has_no_tensorscatter_left_unmasked -> test_static_cache_graph_contains_causal_mask_ops (name now matches the Range/GreaterOrEqual assertion). - drop internal "Option Y" jargon from the decode-test docstring. - replace np.array_equal on CUDA logits with np.testing.assert_allclose / not np.allclose (fp tolerance) across the four poison controls. - add test_static_cache_decode_mask_is_per_batch_on_cuda: a batch>1 value-level guard that a shared slot in one row's frontier but the other's future only changes the correct row, proving the 4D mask is built per-batch. - static path now requests _outputs=1 from op.Attention (the updated cache comes from TensorScatter, not the op's discarded present outputs); verified ORT accepts it on CUDA. Deferred (separate follow-up PR, per assessment): M2 per-layer mask hoist. Kept (documented choice): _register_kv_cache_outputs partial-params warn+skip. _attention.py remains branchless (0 control-flow ops); lint clean; static-cache + task + cache_utils subsets green on CUDA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Addresses the Copilot-bot findings from the #340 review assessment: - COPILOT #3 (completes the M3 is_causal=1 sweep): build_graph_test.py test_static_cache_graph_inputs had a stale inline comment claiming causal masking is "handled by is_causal=1 on the Attention op". Updated to the shipped reality (is_causal=0 + explicit offset-aware mask from write_indices). The is_causal=1-mention sweep is now complete across _causal_lm.py (:64, Falcon block) + build_graph_test.py (the remaining :4561/:4610 mentions are correct references to the *rejected* old form, not stale claims). - COPILOT #1 + #2 (indentation): the two negative-control assertions in the decode-frontier and prefill-causal tests were dedented outside their `with tempfile.TemporaryDirectory()` block. Moved them back inside. - Negative-control comparison semantics: reverted the three "logits MUST change" negative controls to exact `assert not np.array_equal` (any bit of change proves the slot was attended). The np.array_equal -> assert_allclose conversion correctly applies ONLY to the "logits UNCHANGED" equality/parity checks, which keep the fp-tolerance compare. Applied consistently to the new batch>1 per-batch mask test as well, and clarified the tolerance-constants comment. Lint clean; static-cache decode 4/4 (CUDA) and build_graph static-cache 10/10 green. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Member
|
LMK when this is ready, thanks |
…e assert Two non-logic review nits on top of the #340 review fixes (compact-cache contract stays doc-only per all three reviewers — no defensive key<nonpad mask term added): - static_cache_decode_test.py: note in the tolerance-constants comment that the 50.0 poison value far exceeds the 1e-5 band, so the tolerant "unchanged" equality control still cannot mask a genuine leak into a masked slot. - _task_test.py (MLA present/past test): assert `present.shape is not None` with a clear message before the shape comparison, replacing the cryptic TypeError that dims() would otherwise raise if the present-KV head_dim stamp did not run. Comment / assert-message only; no behavior change. Lint clean; static-cache decode + present/past + build_graph static-cache subsets green. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
Author
Round 2 review (delta
|
The static-cache causal attention mask depends only on S_q, max_seq and write_indices — all identical across decoder layers — so rebuilding it per layer duplicated ~16 nodes per layer. Build the mask once and share the same ir.Value across all layers. - Add `causal_mask: ir.Value | None` field to StaticCacheState; thread the shared value from _make_static_cache_inputs into every layer's state. - _apply_attention consumes static_cache.causal_mask when present and keeps a fallback that builds the mask on demand for direct callers. - Build the mask once in _make_static_cache_inputs using input_ids (S_q) and cache_pairs[0][0] (max_seq); the shared Value guarantees bit-identical logits vs the per-layer build (parity by construction). Also reframe the mask docstring + CHANGELOG honestly: the nonpad_kv_seqlen key-bound is enforced by the ORT Attention kernel itself (external-cache input #6, verified bit-identical on CUDA and CPU when poisoning padding slots), so the causal-only mask must not re-encode it — a `j < nonpad` term would merely duplicate input #6 and add dead nodes. Non-compact / interior padding holes are out of contract (a scalar nonpad cannot express them). Tests: add test_static_cache_mask_built_once (exactly one GreaterOrEqual mask root; all Attention nodes share one mask Value) and assert shared-Value identity in test_static_cache_attention_has_causal_mask_input. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Batched Minor-cleanup follow-up to 36617c9 (no behavior change to shipped Major-2 hoist). Addresses code/critical/readability review nits: - Test the previously-uncovered _apply_attention fallback: add TestApplyAttentionStaticCacheFallback exercising static-cache mode with causal_mask=None. Asserts the fallback builds the mask via the 4-arg create_static_cache_causal_mask path (GreaterOrEqual root), that a hoisted mask Value is consumed by-identity (not rebuilt), and that the fallback graph is op-multiset-equivalent to the hoisted graph. (code-reviewer + critical-reviewer) - Guard zero-layer models in _make_static_cache_inputs: early-return [] when cache_pairs is empty, before indexing cache_pairs[0][0], so a 0-layer config yields [] instead of an opaque IndexError. (critical-reviewer M2) - Fix stale inline comment on the mask root op: "(causal + padding)" -> positional causal bound only; nonpad/padding is kernel-enforced (input #6), matching the reframed causal-only docstring. (readability M-1) - Correct create_static_cache_causal_mask `query` param doc: it accepts a 2D [batch, S_q] Value (e.g. input_ids) since only dim 1 is read; drop the stale "dims 0/1" wording. (readability M-2) Tests: +3 new passing (fallback path); full fast suite 21 failed / 2673 passed / 46 errors — zero new failures vs 36617c9 (the 21 fail + 46 err are pre-existing missing-optional-dep collection issues). lintrunner f clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
_register_kv_cache_outputs took six present-shape parameters that are all-or-nothing by contract (pass all six to stamp explicit present.* types, or none to infer). A partial set (1-5) previously logged a warning and proceeded, falling back to the known-wrong GroupQueryAttention inference path — shipping a structurally-wrong model (mis-derived present head_dim) with only a log line. A partial set is always a wiring slip with no legitimate use, so reject it fail-closed: raise ValueError naming both the provided and the missing parameters. This is stronger and simpler than an opt-in strict flag because no conformant caller passes a partial set — verified: every call site passes 0 params (infer opt-out) or all 6 (_causal_lm.py:199, stamp), so the raise cannot regress any production path. Remove the now-dead `import logging` / `logger` (this was the file's only logger use). Reframe the docstring to state partial sets raise. Rewrite test_partial_params_do_not_stamp into test_partial_params_raise: the exact input that previously passed silently now raises, and the message names all four omitted parameters. Zero ONNX node-count change (graph-output naming/typing only); no interaction with the Major-1/2 static-cache mask region. Closes #341. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Follow-up to 9fafc6e (no behavior change). Addresses readability-review nits: - Docstring: a partial present-shape set is "always a wiring slip ... never legitimate" (was "almost always ... no legitimate use"), removing the hedge that contradicted the now-fail-closed contract and aligning with the CHANGELOG's "always". - test_partial_params_raise: merge the two identical pytest.raises calls into a single `with pytest.raises(...) as exc:` block that asserts both the message pattern and all four omitted parameter names against str(exc.value). - test_no_params_leaves_shapes_untouched: drop the vestigial caplog silence-check (the module no longer has a logger), and the now-unused caplog fixture and `import logging`. The shape-untouched assertions are unchanged. Tests: _cache_utils 5/5 pass; full fast suite zero new failures vs 9fafc6e (21 failed / 2673 passed / 46 errors, all pre-existing missing-dep). lint clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Contributor
Author
Review synthesis (5-reviewer team: readability, code, critical, deep, integration)Verdict: mergeable. No Critical or Major blockers. The core attention-mask math was independently verified against the ONNX opset-24 Attention reference implementation and is derived-correct for prefill, decode, and chunked phases. Correctness — verified (grounded in the opset-24 reference)
The one "Major" — reclassified as follow-up (not a blocker)Incomplete rollout of present-shape stamping. Only Adjudicated: unenforced compact-cache invariantTwo reviewers flagged that correctness relies on Worth addressing in-PR (Minor)
Nits (optional)
Follow-ups to file
🤖 Synthesized from a 5-model review team (Claude, GPT, Gemini). |
This was referenced Jun 8, 2026
Contributor
Author
|
@justinchuby This is ready. But if we think ORT fix can get in pretty quick and static cache attention is not used at the moment, maybe we can wait my ort pr. |
Member
|
Thanks - could you isolate bug fixes into potentially another PR so it is easier to review? I think we can probably fix the bugs first, and then assume the ort patch you referred to to produce simpler graphs here. |
Contributor
Author
|
Closing in favor of #351, which splits out the bug fixes (fp16→fp32 fold widening; GQA The Option-Y static-cache graph workaround from this PR ( |
titaiwangms
added a commit
that referenced
this pull request
Jun 12, 2026
Splits the **bug-fix** half of #340 into its own reviewable, mergeable PR, per @justinchuby's review request ("could you isolate bug fixes into potentially another PR"). ## Fixes 1. **fp16 GQA export emitting fp32 packed weights.** When building fp16 models, `_cast_module_dtype` casts params to fp16 but the folded initializer `Value`s lost their declared `.dtype` (None) while `const_value` stayed fp16. `FoldConcatInitializersPass` / `FoldTransposedInitializerPass` then defaulted the packed/transposed initializer to `FLOAT`, serializing fp32 weights and making ORT reject the model with a fp16/fp32 `MatMul` type-parameter error on both CPU and CUDA EPs. New shared helper `mobius._passes._dtype_utils.initializer_dtype()` resolves the effective dtype from `const_value` when the type annotation was dropped. Also strips dead pre-pack weights via `graph.remove(node, safe=True)` so DCE drops the orphaned q/k/v_proj source initializers. 2. **GQA `present.*` head_dim mis-declaration.** `_register_kv_cache_outputs` now stamps explicit `present.*` KV-cache output shapes/dtypes for GQA instead of relying on the known-wrong shape-inference path. 3. **Fail-closed `_register_kv_cache_outputs`** (closes #341). A partial set of present-shape parameters now **raises `ValueError`** (naming provided + missing params) instead of logging a warning and shipping a structurally-wrong model. All-six (stamp) or none (infer) are unaffected. Includes regression tests for all three (155 tests; an e2e ORT CPU-EP load test reproduces the original fp16/fp32 `MatMul` failure without the fix). ## Scope note The **Option-Y** static-cache graph workaround from #340 (`is_causal=0` + explicit causal mask + `nonpad_kv_seqlen`, forcing MEA) is **dropped, not landed**. The maskless `is_causal=1` + `nonpad_kv_seqlen` end-state (Flash-eligible, fewer nodes) will be emitted directly once onnx/onnx#8068 + microsoft/onnxruntime#28958 ship in a pinnable ORT release and mobius bumps its ORT pin — tracked by #345. Supersedes the bug-fix portion of #340. --------- Signed-off-by: titaiwang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fp16/GQA export fixes + capture-safe static-cache attention (Option-Y)
Purpose
This PR delivers two related improvements to ONNX export:
fp16 / GQA export correctness — fixes constant-folding and grouped-query
attention (GQA) handling so fp16/bf16 models export with correct present-KV
shapes and dtypes.
Capture-safe static-cache attention ("Option-Y") — a static (pre-allocated,
right-padded) KV-cache attention configuration that is CUDA-graph-capture-safe
and actually runnable on ORT today.
The static-cache path uses
is_causal=0+ an explicit causal mask +nonpad_kv_seqlen(opset-24Attentioninput Bump ruff from 0.15.4 to 0.15.6 in /requirements/lintrunner #6). With an explicitattn_maskpresent, ORT routes the op to the Memory-Efficient Attention (MEA) kernel,
which applies both the per-batch
nonpad_kv_seqlenkey bound and the causal mask.This is currently the only CUDA-runnable static-cache configuration: the
opset-24 kernel rejects
is_causal=1together withnonpad_kv_seqlenwhenS_q != total_kv(returnsNOT_IMPLEMENTED), until theis_causal+nonpadFlashpath lands upstream (Attention op: support offset-aware causal masking for KV-cache decode and chunked/mid-cache prefill (nonpad_kv_seqlen) onnx/onnx#8054).
What changed this round
Major-2 — hoist the static-cache causal mask (build once, share across layers)
The static-cache causal mask depends only on shapes +
write_indices, so it islayer-invariant. It is now built once at the task level and threaded as a
single
ir.Valueinto every decoder layer's attention, instead of being rebuiltper layer.
ir.Value(identity), plus aNone-fallback parity test for the non-static path.Major-1 — intentionally NOT adding a graph-level
key_positions < nonpadtermA reviewer suggested adding an explicit
And(causal, key_positions < nonpad)term tothe static-cache mask to defend the padding bound. After investigation we are
intentionally not adding it — it is redundant:
nonpad_kv_seqlenis standardai.onnxopset-24Attentioninput Bump ruff from 0.15.4 to 0.15.6 in /requirements/lintrunner #6 (absent inopset 23). Any conformant opset-24 runtime must mask key positions
j >= nonpad_kv_seqlen[b]as the per-batch key bound, independent of the explicitattn_mask.[nonpad, S_kv)yields
max|Δ| = 0on both CUDA and CPU; non-vacuous positive control changes output),against ORT source (MEA applies both
seqlens_kright-padding andattn_bias; thecutlass key loop hard-stops at
k_end = nonpad[b]), and against the ONNX spec.GreaterOrEqualon positions) is kept — it isload-bearing because
is_causal=0means the op does not apply causality itself.nonpad_kv_seqlencan only express a compact valid prefix[0, nonpad),never ragged/interior holes — so static caches rely on the contiguous-fill invariant
(padding only on the right), documented as a precondition.
The mask-builder docstring and CHANGELOG were reframed to state this honestly (the
padding bound is enforced by the kernel via input #6, not by a redundant mask term).
Major-3a —
_cache_utilspresent-shape stamping is now fail-closed_register_kv_cache_outputs(src/mobius/tasks/_cache_utils.py) takes six"present-shape" parameters that are all-or-nothing: pass all six to stamp explicit
present.*types (required soGroupQueryAttentionexports declare the correcthead_dim), or none to opt out and infer. A partial set (1–5 of 6) is always awiring slip — and previously it merely logged a warning and proceeded, shipping a
structurally-wrong model (mis-derived
present.*head_dim) with only a log line.This is now fail-closed: a partial present-shape set raises
ValueErrornamingthe provided and missing parameters, instead of warn-and-proceed. All current call sites
pass either 0 parameters (intentional infer opt-out) or all 6 (stamp), so there
is no production regression — the raise only fires on a genuine future wiring bug.
A unit test that previously asserted the silent fallback now asserts the raise (the
regression proof).
Known items
Benchmark "Compare results" node-count increase is the intended structural
cost of the explicit Option-Y causal mask on the tiny benchmark model (the explicit
mask + static-cache wiring adds nodes vs the
is_causal=1baseline that does not runon CUDA for this config). This is a deliberate correctness/runnability trade-off, not a
regression — requesting maintainer-accept of the node-count delta at merge.
This increase is temporary. The explicit Option-Y mask exists only because ORT
currently hard-rejects
is_causal=1+nonpad_kv_seqlen(nopast_key) for thisstatic-cache shape, which forces the
is_causal=0+ explicit-mask formulation; theextra
num_nodesover the masklessis_causal=1baseline is therefore known andintended, not a regression. Once Fix Attention is_causal causal-mask alignment + composed is_causal/attn_mask NaN robustness for external (static) KV cache (#8054) onnx/onnx#8068 (the spec change defining bottom-right
causal alignment for this path; implements Attention op: support offset-aware causal masking for KV-cache decode and chunked/mid-cache prefill (nonpad_kv_seqlen) onnx/onnx#8054) and the corresponding ORT
follow-up land, mobius will revert to
is_causal=1+nonpad_kv_seqlenwith noexplicit mask — which removes these nodes, clears the benchmark node-count gate, and
unlocks Flash-attention eligibility for this configuration. The mobius-side revert is
tracked as a follow-up (onnxruntime/mobius#345; see the
Upstream coordination & temporary-workaround status section above).
Testing
ir.Valueidentity acrosslayers,
None-fallback parity for the non-static path, the static-cache causal-maskbehavior tests, and the
_register_kv_cache_outputspartial-set fail-closed test(previously asserted warn-and-proceed; now asserts
ValueError).Closes #341 (Major-3a
_cache_utilsfail-open → fail-closed, fixed in-PR rather thandeferred). Major-3b (
apply_weightsINFO-only unmapped weights) is pre-existing onmainand remains out of scope.