Fix fp16/GQA static-cache export correctness (split from #340) - #351
Merged
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> Signed-off-by: titaiwang <titaiwang@microsoft.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> Signed-off-by: titaiwang <titaiwang@microsoft.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> Signed-off-by: titaiwang <titaiwang@microsoft.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> Signed-off-by: titaiwang <titaiwang@microsoft.com>
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> Signed-off-by: titaiwang <titaiwang@microsoft.com>
…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> Signed-off-by: titaiwang <titaiwang@microsoft.com>
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> Signed-off-by: titaiwang <titaiwang@microsoft.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> Signed-off-by: titaiwang <titaiwang@microsoft.com>
…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> Signed-off-by: titaiwang <titaiwang@microsoft.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> Signed-off-by: titaiwang <titaiwang@microsoft.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>
Signed-off-by: titaiwang <titaiwang@microsoft.com>
…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> Signed-off-by: titaiwang <titaiwang@microsoft.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> Signed-off-by: titaiwang <titaiwang@microsoft.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> 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>
This was referenced Jun 10, 2026
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR improves correctness of exported fp16 + GQA decoder models by (1) forcing present.* KV-cache outputs to have explicit, symmetric shapes/dtypes to their corresponding past_key_values.* inputs, and (2) preserving fp16 dtypes when folding packed/transposed initializers so models don’t silently widen weights to fp32 (or corrupt bytes at serialization).
Changes:
- Stamp explicit
present.{i}.{key,value}output shapes/dtypes (and fail-closed on partial stamping parameters) to avoid incorrectGroupQueryAttention-inferredhead_dim. - Preserve initializer dtype during
Concat/Transposefolding via a sharedinitializer_dtype()helper (declared-type +const_valuefallback), including safe node removal to enable DCE. - Add unit + end-to-end regression tests covering KV-cache I/O symmetry and fp16 folding/serialization round-trips.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/tasks/_task_test.py | Adds a regression test ensuring present.* outputs match past_key_values.* inputs (dims + dtype). |
| 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 shapes/dtypes and raise on partial parameter sets. |
| src/mobius/tasks/_cache_utils_test.py | Adds focused tests for stamping behavior, partial-parameter failure, and naming. |
| src/mobius/_passes/_fold_transpose.py | Uses initializer_dtype() to keep folded transposed initializers at the correct dtype (e.g., fp16). |
| src/mobius/_passes/_fold_transpose_test.py | Adds a regression test for dtype preservation when declared dtype metadata is missing. |
| src/mobius/_passes/_fold_dtype_e2e_test.py | Adds pipeline-level fp16 regression tests that validate folding + serialization round-trips preserve dtype/bytes. |
| src/mobius/_passes/_fold_concat.py | Adds const-value guard, resolves/stamps packed dtype via initializer_dtype(), and uses safe node removal for DCE. |
| src/mobius/_passes/_fold_concat_test.py | Adds tests for DCE detach, dtype fallback, ORT loadability, and serialized packed-weight value integrity. |
| src/mobius/_passes/_dtype_utils.py | Introduces initializer_dtype() helper for consistent dtype resolution in initializer-producing passes. |
| src/mobius/_passes/_dtype_utils_test.py | Adds unit tests for initializer_dtype() behavior (declared, fallback, disagreement warning). |
| CHANGELOG.md | Documents the KV-cache present-shape fix, fp16 GQA export dtype fix, and fail-closed behavior. |
| .agents/skills/mobius-onnx-export-gotchas/SKILL.md | Adds internal documentation on export “gotchas” and verification steps for fp16/GQA exports. |
|
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: 68 → 68 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: 54 → 54 nodes No op-sequence changes. Interface changes:
llama / model — 1 change(s)Op summary: 62 → 62 nodes No op-sequence changes. Interface changes:
phi3 / model — 1 change(s)Op summary: 60 → 60 nodes No op-sequence changes. Interface changes:
qwen / model — 1 change(s)Op summary: 62 → 62 nodes No op-sequence changes. Interface changes:
qwen2 / model — 1 change(s)Op summary: 62 → 62 nodes No op-sequence changes. Interface changes:
qwen2_moe / model — 1 change(s)Op summary: 224 → 224 nodes No op-sequence changes. Interface changes:
qwen3 / model — 1 change(s)Op summary: 74 → 74 nodes No op-sequence changes. Interface changes:
qwen3_moe / model — 1 change(s)Op summary: 202 → 202 nodes No op-sequence changes. Interface changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
approved these changes
Jun 12, 2026
Address @justinchuby's non-blocking review suggestions: - P1: initializer_dtype() now raises ValueError on a genuine declared-vs- const_value dtype contradiction (names both dtypes + the initializer), consistent with the export pipeline's fail-closed contract. The declared-is-None fallback (the core fp16 fix path) stays non-raising. - P3: build the test fixtures via the ir.Value constructor instead of the ir.val() factory, which validates and refuses the degenerate fixtures. - Update _fold_concat_test fixtures to declare per-array dtypes consistent with their const_value so they exercise the fold path, not the new raise. P2 (fix the dtype-drop at its source) is intentionally not included here: the declared dtype is dropped by the GQA PackQKV rewrite's Concat/Transpose intermediates (src/mobius/rewrite_rules/_group_query_attention.py), which is outside this PR's file scope and multi-call-site; the cast site hypothesised in review (_cast_module_dtype) already stamps the declared dtype. The existing two-layer defense (fold-pass dtype stamping + const_value fallback) keeps fp16 weights correct. Tracked as a follow-up. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Contributor
Author
|
Thanks @justinchuby for the review and approval — the suggestions all push toward the root P1 — raise on a declared-vs- P2 — which process drops the dtype ( That rewrite is outside the scope of this bug-fix PR, spans multiple call sites, and we P3 — test fixtures ( On the automated "PR description describes Olive vision metrics" comment. Thanks — we Thanks again for the careful review. |
The new test docstring added in 3b0fa13 tripped the repo's enforced pydocstyle gate (ruff "D" select, google convention): D205 (blank line required after the summary line), D209 (closing quotes on their own line), and the resulting ruff format diff, failing the PR Lint check. Rewrite the docstring as a single-line summary followed by a blank line and a wrapped description. Lint-only change; no test logic or assertions are touched (22 dtype/fold-concat tests still pass). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
titaiwangms
enabled auto-merge (squash)
June 12, 2026 17:20
Copilot AI
added a commit
that referenced
this pull request
Jul 30, 2026
The PackQKV rewrite built the packed QKV weight with `op.Concat` / `op.Transpose`, whose output values carry no declared type. When those intermediates were folded into initializers, there was no declared dtype to inherit, so fp16 weights could be widened to fp32 (mitigated downstream in #351 by const_value fallback + fold-pass stamping). Verified the issue still reproduces on main (packed Concat/Transpose outputs had `dtype is None`) and fixed it at the source: propagate the projection weight dtype (and bias dtype for the biased path) onto the new intermediates in both PackQKV rewrite sites. The downstream mitigation is kept as defense-in-depth.
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.
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
fp16 GQA export emitting fp32 packed weights. When building fp16 models,
_cast_module_dtypecasts params to fp16 but the folded initializerValues lost their declared.dtype(None) whileconst_valuestayed fp16.FoldConcatInitializersPass/FoldTransposedInitializerPassthen defaulted the packed/transposed initializer toFLOAT, serializing fp32 weights and making ORT reject the model with a fp16/fp32MatMultype-parameter error on both CPU and CUDA EPs. New shared helpermobius._passes._dtype_utils.initializer_dtype()resolves the effective dtype fromconst_valuewhen the type annotation was dropped. Also strips dead pre-pack weights viagraph.remove(node, safe=True)so DCE drops the orphaned q/k/v_proj source initializers.GQA
present.*head_dim mis-declaration._register_kv_cache_outputsnow stamps explicitpresent.*KV-cache output shapes/dtypes for GQA instead of relying on the known-wrong shape-inference path.Fail-closed
_register_kv_cache_outputs(closes Static-cache: add opt-in strict mode for partial present-shape params in _register_kv_cache_outputs #341). A partial set of present-shape parameters now raisesValueError(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
MatMulfailure 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 masklessis_causal=1+nonpad_kv_seqlenend-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.