[ONNX] Use onnx Attention operator for scaled_dot_product_attention

[csoiram]

🚀 The feature, motivation and pitch

ONNX have introduced MHA operator in opset 23 (https://onnx.ai/onnx/operators/onnx__Attention.html#l-onnx-op-attention-23). This could be used when exporting scaled_dot_product_attention to ONNX format. Currently the scaled_dot_product_attention gets broken down to constituent ops when exporting to ONNX which complicates the model and makes it harder to identify the attention block when compiling the network for inference in custom HW backends.

Alternatives

No response

Additional context

No response

[justinchuby]

cc @shubhambhokare1 for multi-opset support

[kkhode]

This can be assigned to me, I was working on a PR for this feature. I should be allowed to post it in a few days.

[yuanyao-nv]

When I try exporting a toy model with aten::scaled_dot_product_attention

class SDPA(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, Q, K, V):
        return torch.nn.functional.scaled_dot_product_attention(Q, K, V)

query = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")

model = SDPA()
model(query, key, value)

ep = torch.export.export(model, (query, key, value),)
print(ep)
torch.onnx.export(
    ep,
    (query, key, value),
    "sdpa.onnx",
    dynamo=True,
    opset_version=23,
)

The resulting ONNX model seems to have num_heads incorrectly inferred?