Support enable_gqa in scaled_dot_product_attention

[Yigtwxx]

Summary

PyTorch 2.5 added the enable_gqa argument to torch.nn.functional.scaled_dot_product_attention. It is how grouped-query attention is expressed, where key and value carry fewer heads than query.

The converter never read that argument, so key and value kept their original head count and conversion failed:

class Model(torch.nn.Module):
    def forward(self, query, key, value):
        return torch.nn.functional.scaled_dot_product_attention(
            query, key, value, enable_gqa=True
        )

query = torch.rand(2, 8, 5, 8)        # 8 query heads
key = value = torch.rand(2, 2, 7, 8)  # 2 key / value heads
traced = torch.jit.trace(Model(), (query, key, value))
ct.convert(
    traced,
    inputs=[ct.TensorType(shape=t.shape) for t in (query, key, value)],
    minimum_deployment_target=ct.target.iOS18,
)

ValueError: query, key, value must have a same batch dimension, got
* query batch = (2, 8)
* key batch = (2, 2)
* value batch = (2, 2)

Below iOS18, where the op is decomposed rather than fused, the same model fails with ValueError: Incompatible dim 1 in shapes (2, 8) vs. (2, 2). Both TorchScript and torch.export hit this.

Implementation

enable_gqa is keyword-only in torch, so TorchScript serializes it as the eighth positional input while torch.export leaves it in the node kwargs. Both are now parsed. _scaled_dot_product_flash_attention_for_cpu and coreml::sdpa have no such argument and are unchanged.

Key and value heads are then repeated along dimension -3 with expand_dims + tile + reshape, reproducing torch.repeat_interleave(x, repeats, dim=-3). This is deliberately not a plain tile: for two repeats, repeat-interleave gives head0, head0, head1, head1 while a tile gives head0, head1, head0, head1, and the two produce different attention outputs.

The expansion runs before the fused/decomposed branch, so the iOS18 scaled_dot_product_attention op and the earlier decomposition are both covered by one code path.

Nothing is emitted when enable_gqa is false, or when query and key/value already have the same number of heads, so models that do not use grouped-query attention convert to exactly the same program as before.

The head count has to be known at conversion time, since it decides how many repeats to emit; batch size and sequence lengths may still be dynamic. Head counts that are not divisible, key and value head counts that disagree, a rank below 3, and a non-constant enable_gqa each raise with a specific message.

Tests

Two tests are added to TestScaledDotProductAttention:

• test_enable_gqa, over query/key-value head counts (8, 2), (6, 3), (4, 1) and (4, 4), both compute units, both backends, all frontends, and both None and iOS18 deployment targets, so the fused op and the decomposition are each exercised. It also asserts that one tile is inserted per key and value, and none at all when the head counts already match. (6, 3) is the smallest case where a tile and a repeat-interleave disagree numerically, so it pins the expansion semantics.
• test_enable_gqa_dynamic_shapes, which keeps the head counts static while batch size and both sequence lengths are dynamic.

Both are skipped below torch 2.5.0. The ExecuTorch frontend decomposes scaled_dot_product_attention before it reaches the converter, so there the tests serve as numerical regression coverage rather than exercising this translation.

I do not have an Apple device, so the on-device comparison in run_compare_torch did not run locally. To check the lowering I evaluated the emitted MIL program in NumPy against PyTorch eager for every head count and deployment target above, including is_causal and dynamic shapes, on the TorchScript and torch.export frontends; all matched. The rest of TestScaledDotProductAttention reports the same results before and after this change.

[TobyRoseman]

This change looks good. Nice unit test coverage.

CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/2848162955

@Yigtwxx - Can you give us some background on this change? Does this change allow us to convert a specific model which we previously could not?