Validate packed GQA head counts during shape inference

[danielsongmicrosoft]

Description

Packed GroupQueryAttention shape inference divided the query hidden size by num_heads + 2 * kv_num_heads without validating the attributes first. Malformed head counts could therefore cause division by zero or signed overflow while resolving a model.

Validate both counts before the calculation, reject grouped-head overflow, require concrete hidden sizes to divide evenly, and preserve symbolic hidden dimensions.

Testing

• cmake --build build\Windows\Release --config Release --target onnxruntime_graph --parallel
• cmake --build build\Windows\Release --config Release --target onnxruntime_provider_test --parallel
• onnxruntime_provider_test.exe --gtest_filter=GroupQueryAttentionTest.CausalMaskDefaultsToEnabled_CPU

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

Regression coverage for the new validation and symbolic-dimension paths is missing.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Validates packed GroupQueryAttention head counts and hidden-size calculations during shape inference.

Changes:

• Rejects invalid or overflowing head-count combinations.
• Checks concrete packed hidden-size divisibility.
• Preserves symbolic hidden dimensions.

File summaries

File	Description
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds packed GQA shape-inference validation.

Review details

Suppressed comments (1)

onnxruntime/core/graph/contrib_ops/bert_defs.cc:303

• This branch is selected by hasInputShape(ctx, 2), not by whether the optional value input is present. A non-packed node can provide key/value while leaving value's rank/shape unknown; in that case this new check treats the query as packed and rejects valid dimensions (for example num_heads=4, kv_num_heads=1, query hidden size 32, where the non-packed head size is 8 but 32 is not divisible by 6). Distinguish packed mode from input presence (ctx.hasInput(2), as used by the PagedAttention inference below) and only apply the grouped-head validation to the no-value-input case.

      if (query_dims[2].has_dim_value()) {
        const int64_t hidden_size = query_dims[2].dim_value();
        if (hidden_size % grouped_heads != 0) {
          fail_shape_inference("Packed query hidden size must be divisible by the grouped head count.");

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Lite (auto)

Note

Copilot is running an experiment and ran this review at Lite.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Jiajia Qin (qjia7)]

Review frame

• Reviewed head: 241b0e71fce3d33877e4414fda455c462fadea63; base: f38538cd5a4b5945a4c839565a8eebc65e1e2ef8.
• Problem validity: validated. Packed shape inference divides by unchecked head-count arithmetic during Graph::Resolve, before kernel validation. Rejecting malformed counts and non-divisible widths, and leaving an uncomputable derived width unknown, are appropriate.
• Risk/scope: deep. This changes model-loading behavior in shared GQA/SparseAttention inference and introduces a test that must compile across platforms.
• Direction gate: pass. The shared inference function is the correct boundary. The owner solution checks arithmetic before evaluation, distinguishes packing by input presence, and preserves only shape information that can be derived. The implementation needs the corrections below.

Confirmed findings

S1 [P1]: Include the assertion helper explicitly

Location: group_query_attention_shape_inference_test.cc:104.

The new file calls ASSERT_STATUS_OK without including test/util/include/asserts.h. MSVC's precompiled test header supplies it transitively, but onnxruntime_test_pch.cmake:3 enables that header only for MSVC. Consequently, the new test breaks the provider-test build with GCC/Clang.

Compiling the exact new file without the precompiled header reproduced C3861: ASSERT_STATUS_OK: identifier not found. Supplying only test/util/include/asserts.h made that compilation succeed. Please include this header directly and validate the test without MSVC's precompiled header.

C1 [P2]: Apply packed-width validation only when value is absent

Location: bert_defs.cc:302.

The enclosing branch uses hasInputShape(ctx, 2) to distinguish unpacked from packed input. A supplied value tensor with an unknown rank has no shape, so it reaches this new packed-only check.

For example, use num_heads=4, kv_num_heads=1, query shape [1,1,32], key shape [1,1,8], and a supplied float value input whose shape is unknown. This is unpacked GQA with head size 8, but the PR evaluates 32 % 6 and rejects the model with Packed query hidden size must be divisible by the grouped head count. A Graph::Resolve comparison succeeded on the base and failed on the PR. The base already inferred an incorrect width in this case; this PR worsens that existing misclassification into a hard model-load failure.

Please select packing from input presence, using ctx.hasInput(2) as the neighboring PagedAttention inference does, and separately guard reads of the value shape. Keep the subsequent cache-shape branch consistent. Add a regression with a supplied, unshaped value input and verify the output width remains 32.

Clarifications

None outstanding.

Test coverage

Validation used freshly compiled base/PR bert_defs.cc and the new test file, linked against existing Windows Release test objects and dependencies. This was targeted validation, not a clean full build or a Linux/macOS test run.

Scenario	Evidence
Invalid counts, grouped-count overflow, non-divisible packed width, symbolic width	All four new tests fail on the base and pass on the PR.
Valid packed width, packed present-cache outputs, ordinary unpacked input	Three additional graph-resolution controls pass on both versions.
Supplied value with unknown rank	Additional regression passes on the base and fails on the PR: C1.
Existing CPU causal behavior	CausalMaskDefaultsToEnabled_CPU passes on both versions; it exercises unpacked input.
Test compilation without the precompiled header	Fails as submitted; succeeds when the assertion header is supplied: S1.

The overflow guard is evaluated after positivity checks, and the derived output multiplication is bounded by the packed width. The symbolic path correctly leaves the derived dimension anonymous. Negative tests are guarded for ORT_NO_EXCEPTIONS. GitHub currently reports only the passing CLA check, with no build/test CI evidence.

Verdict

The problem is valid and the design direction is appropriate. Request changes for S1 and C1. No clarification requests or cleanup-only findings remain.

[Copilot]

🟢 Approval recommended

The validation is overflow-safe, preserves symbolic shapes, and has focused regression coverage.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0 new
• Review effort level: Balanced