Validate GQA fusion projection shapes

[Akshay Sonawane (apsonawane)]

This pull request strengthens the validation logic for Group Query Attention (GQA) fusion in ONNX Runtime by adding stricter shape checks for projection tensors and guarding against integer overflows in hidden size calculations. It also introduces a new unit test to ensure that the fusion is skipped when projection tensor shapes do not match the expected configuration.

Improved validation and error handling:

• Added the ProjectionTensorShapesMatch function to verify that the Q, K, and V projection tensor shapes match the expected dimensions based on quantization and head attributes. The fusion is skipped if the shapes do not match. [1] [2]
• Added checks to prevent integer overflows when calculating q_hidden_size, kv_hidden_size, and output_hidden_size, ensuring all head attributes are positive and calculations are safe.
• Included <limits> for use of std::numeric_limits in overflow checks.

Testing enhancements:

• Added a new test, GroupQueryAttentionFusionSkipsMismatchedProjectionSizesTest, to verify that fusion is correctly skipped when projection tensor sizes do not match the head attributes.
• Updated the GQA fusion test builder to allow specifying the number of GQA heads, supporting more flexible test scenarios. [1] [2]

[Copilot]

Pull request overview

This PR tightens the GroupQueryAttentionFusion preconditions by validating that Q/K/V projection initializers match the expected hidden sizes derived from GQA head attributes, and by adding overflow/validity guards around hidden size calculations. It also extends the existing GQA fusion test graph builder and adds a new unit test covering a mismatched-projection scenario to ensure fusion is skipped.

Changes:

• Added ProjectionTensorShapesMatch and integrated it into the GQA fusion eligibility checks.
• Added positivity and integer-overflow guards when computing q_hidden_size, kv_hidden_size, and output_hidden_size.
• Updated the rotary-embedding GQA fusion test graph builder to allow custom num_heads, and added a regression test ensuring fusion is skipped when projection sizes don’t match head attributes.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/optimizer/group_query_attention_fusion.cc	Adds projection shape validation and overflow-safe hidden size computations to gate the GQA fusion.
onnxruntime/test/optimizer/graph_transform_test_layernorm.cc	Extends the GQA fusion graph builder to parameterize head count and adds a test for mismatched projection sizes (fusion should be skipped).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.