I've been looking into the current implementation of scaled_dot_product_attention for MPS, and I'd like to lay out some of the problems with it and my plans for improving it. These are not necessarily in the order that they will actually be completed in.

Problem - Messy call site of _scaled_dot_product_attention_math_for_mps

In the place where the MPS impl is called, there is a bunch of extra code that should be moved into the MPS impl: https://github.com/pytorch/pytorch/blob/0cc07dbe82d222bcf97266aa978542e2733814f3/aten/src/ATen/native/transformers/attention.cpp#L795-L825

Invoking the MPS impl should be as simple as calling one function.

Problem - No backward pass

Currently we don't have a dedicated MPS impl of the backward pass. Instead, the backward pass is implemented using implicit composite autograd from the device-agnostic "math" backend function _scaled_dot_product_attention_math. The forward pass has to call _scaled_dot_product_attention_math in order to build the graph for the backward pass. This means that if someone wants to run a backward pass of SDPA for MPS, they can't use the dedicated MPS impl of the forward pass. So not only is the backward pass slower than it could be, but the forward pass is also slower, since both are using a device-agnostic impl.

So a dedicated MPS backward pass needs to be implemented.

Problem - Unused second output tensor

As some issues (#176730, #175873) have mentioned, _scaled_dot_product_attention_math_for_mps returns a second output tensor that is currently not used by anything. #173943 proposes completely removing the code that generates the second output, and just return an empty tensor, to improve performance. But this is not the right solution.

The second output is required for calculating the backward pass. So when we have the backward implemented in the future, we will need this second output.

I propose that we conditionally calculate the second output only if one of the input tensors requires grad. In the case where we don't need to calculate it, we should return a null tensor (at::Tensor()) so we can avoid the overhead of creating an empty tensor.

Problem - Math vs. flash attention

While there isn't really documentation about this, it appears that the "math" backend is supposed to just call the generic device-agnostic _scaled_dot_product_attention_math function. Then if the user wants to use a dedicated GPU implementation, they should select one of the other backends, such as the "flash" backend. This is true on CUDA platforms. But for MPS platforms, the dedicated MPS impl gets called if the "math" backend is selected and the inputs don't require grad and some other odd conditions: https://github.com/pytorch/pytorch/blob/19cda9b00702cd65bc9c46d467f0f5e928991f9f/aten/src/ATen/native/transformers/attention.cpp#L811-L814

If my assumption that the "math" backend is supposed to always call the device-agnostic impl is correct, then we should move the dedicated MPS impl to the "flash" backend.

Problem - Generic Metal implementation is never called

There appears to be a Metal kernel implementation of SDPA in sdpa_full_attention_mps, but it isn't called at all.

Presumably, a Metal kernel is more reliable and potentially faster than the current MPSGraph impl, so this codepath should be revived and fixed if it is broken. The MPSGraph impl can be removed if it is categorically worse.

cc @jerryzh168 @ezyang @albanD @gqchen @nikitaved @soulitzer @Varal7 @bobrenjc93 @kulinseth @malfet @DenisVieriu97 @jhavukainen @aditvenk @drisspg @liangel-02 @howardzhang-cv @manuelcandales