Skip GQA key/value repeat when query and key head counts are equal

[Copilot]

torch-nightly now passes enable_gqa=True to scaled_dot_product_attention even when the number of query heads equals the number of key/value heads, tripping the assertion in _attention_repeat_kv_for_group_query (q_num_heads > kv_num_heads) and failing the torch_lib tests in CI.

Change

• aten_scaled_dot_product_attention (onnxscript/function_libs/torch_lib/ops/nn.py): gate the group-query key/value expansion on an actual head-count mismatch, so equal-head inputs bypass the repeat (a no-op) instead of hitting the assertion.

if enable_gqa and query.shape[1] != key.shape[1]:
    key, value = _attention_repeat_kv_for_group_query(query, key, value)

[codecov]

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
12441	1	12440	2419

View the top 1 failed test(s) by shortest run time

onnxscript.rewriter.ort_fusions.mha_test.TestMultiHeadAttention::test_smollm

Stack Traces | 11.5s run time

onnxscript\rewriter\ort_fusions\mha_test.py:34: in test_smollm
    original_outputs = ort_run("original", model, inputs)
onnxscript\rewriter\ort_fusions\_test_utils.py:19: in ort_run
    model_proto.SerializeToString(), options, providers=providers
E   MemoryError

View the full list of 2 ❄️ flaky test(s)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logit_cpu_float16

Flake rate in main: 14.19% (Passed 762 times, Failed 126 times)

Stack Traces | 0.606s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 42 / 125 (33.6%)
E   Greatest absolute difference: 0.8115234375 at index (0, 0, 1) (up to 0.0007 allowed)
E   Greatest relative difference: 2.0 at index (0, 0, 1) (up to 0.1 allowed)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logit_cpu_float32

Flake rate in main: 14.19% (Passed 762 times, Failed 126 times)

Stack Traces | 0.581s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 44 / 125 (35.2%)
E   Greatest absolute difference: 0.8109303712844849 at index (0, 0, 1) (up to 1e-05 allowed)
E   Greatest relative difference: 2.0 at index (0, 0, 1) (up to 1.3e-06 allowed)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}