[G. Ramalingam (gramalingam)]

Summary

• Honor nested GQA mask-match truthiness and consistently bind window_size; reject unsupported sliding-window and structurally different masks with useful diagnostics instead of discarding their semantics.
• Preserve supported pure causal, legacy extra-column/slice, and known-zero-cache mask forms. Verify captured current/past sequence lengths against the actual query/cache, rejecting unproven or broadcastable unrelated sources. Retain the existing ORT GenAI attention-mask/position assumptions and leave ORT defaults unchanged.
• Share FusedMatMul division validation for matching supported floating dtypes, singleton finite/nonzero divisors, and representable alpha scaling. Conservatively decline signed zero, nonfinite divisors/alpha, overflow, underflow, subnormal scales, and double scales not exactly representable by the float32 attribute.
• Preserve output rank: reject rank-expanding singleton divisors and non-scalar divisors when MatMul output rank cannot be established; keep scalar divisors valid with unknown operand ranks.
• Add regression coverage and document the intentional conservative behavior. Independent review identified the captured-length and rank-expansion cases; both repros were fixed and independently reverified before approval of head ac47c2d2e6210131c174afddf0ca508592311c0c.

Validation

• python -m pytest onnxscript\rewriter\ort_fusions\gqa_test.py onnxscript\rewriter\ort_fusions\gqa_extended_test.py onnxscript\rewriter\ort_fusions\gqa_packed_qkv_test.py onnxscript\rewriter\ort_fusions\gqa_packed_qkv_extended_test.py onnxscript\rewriter\ort_fusions\fused_matmul_rule_sets_test.py -q — 48 passed, 1 existing memory-related skip; independently rerun by reviewer.
• GQA causal positives and FusedMatMul float/double shape/value comparisons run with ORT graph optimizations disabled. Rejected candidates assert zero applications, unchanged graphs, and explanatory match reasons.
• lintrunner -a run before each commit; reports unrelated baseline TID251/RUF036 warnings in examples/framework/generated-opset files. No unrelated files changed.
• lintrunner -a -m 0c3ae10481d53cbba2e19979e41a53b92704e34b — passes with no lint issues, using repository-pinned Ruff 0.15.8. git diff --check passes.
• Limitation: no GPU runtime available. Float16/bfloat16 positive rewrites are checked structurally; execution parity is covered for the CPU float/double FusedMatMul kernels.

This change is independent of #3133 and does not implement other rewriter-readiness workstreams.

[codecov]

Codecov Report

❌ Patch coverage is 99.31507% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.03%. Comparing base (0c3ae10) to head (51acfc2).
⚠️ Report is 8 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/rewriter/ort_fusions/gqa.py	95.55%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3134      +/-   ##
==========================================
+ Coverage   73.24%   74.03%   +0.78%     
==========================================
  Files         267      270       +3     
  Lines       33158    34130     +972     
  Branches     3176     3285     +109     
==========================================
+ Hits        24288    25269     +981     
- Misses       7801     7812      +11     
+ Partials     1069     1049      -20     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[Copilot]

🟢 Approval recommended

There are no active review issues for the attention-mask and FusedMatMul fusion guard changes.

No open issues 🎉

What changed in this PR

This PR strengthens ONNX Runtime-specific rewrite guards to preserve attention-mask semantics and prevent unsafe FusedMatMul division fusions.

Changes:

• Validate GQA mask matches and sequence lengths while retaining supported causal-mask forms and rejecting unsupported sliding windows.
• Share FusedMatMul division checks for compatible dtypes, safe scaling, and output-rank preservation.
• Add regression coverage and document conservative fusion boundaries.

File	Description
onnxscript/​rewriter/​ort_fusions/​gqa.py	Validates causal-mask structure, window bindings, and query/cache lengths.
onnxscript/​rewriter/​ort_fusions/​gqa_test.py	Covers supported masks, rejected variants, diagnostics, and numerical equivalence.
onnxscript/​rewriter/​ort_fusions/​fused_matmul_rule_sets.py	Centralizes divisor, scaling, dtype, and output-rank validation.
onnxscript/​rewriter/​ort_fusions/​fused_matmul_rule_sets_test.py	Tests division boundaries, rank preservation, and supported rewrites.
docs/​api/​rewriter.md	Documents supported fusion contracts and conservative rejection behavior.

_{Review effort: Balanced}