Repository: pytorch
Issue: [MPS] Inconsistent performance issues
Labels: module: performance,triaged,module: mps

### 🐛 Describe the bug

In some cases certain ops are slow. The failure mode does not always trigger, but its relatively easy to trigger this behavior if run in rapid succession.

Take the following SDPA benchmark as an example

```python
import torch
from torch.profiler import ProfilerActivity, profile

device = "mps"
dtype = torch.bfloat16
n = 19

shape = [8, 64, 1024, 1024]

query = torch.rand(shape, device=device, dtype=dtype)
key = torch.rand(shape, device=device, dtype=dtype)
value = torch.rand(shape, device=device, dtype=dtype)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    for _ in range(n):
        torch.nn.functional.scaled_dot_product_attention(query, key, value)

print(
    prof.key_averages(group_by_input_shape=True).table(
        sort_by="cpu_time_total", row_limit=10
    )
)
prof.export_chrome_trace(f"trace_test_{device}.json")
```

If we run this multiple times we get very different output in terms of performance characteristics.

### Output 1

```
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                                                Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls                                                                      Input Shapes  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                  aten::scaled_dot_product_attention         0.01%     221.958us       100.00%        1.536s      80.858ms            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
    aten::_scaled_dot_product_attention_math_for_mps        99.97%        1.536s        99.99%        1.536s      80.847ms            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
                                         aten::empty         0.01%     200.878us         0.01%     200.878us       5.286us            38                                                          [[], [], [], [], [], []]  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
Self CPU time total: 1.536s
```

### Output 2

```
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                                                Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls                                                                      Input Shapes  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                  aten::scaled_dot_product_attention         0.79%      23.418us       100.00%       2.975ms     156.570us            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
    aten::_scaled_dot_product_attention_math_for_mps        98.38%       2.927ms        99.21%       2.951ms     155.338us            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
                                         aten::empty         0.83%      24.837us         0.83%      24.837us       0.654us            38                                                          [[], [], [], [], [], []]  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
Self CPU time total: 2.975ms
```

### Output 3

```
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                                                Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls                                                                      Input Shapes  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                  aten::scaled_dot_product_attention         0.01%     146.001us       100.00%        1.062s      55.883ms            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
    aten::_scaled_dot_product_attention_math_for_mps        99.97%        1.061s        99.99%        1.062s      55.875ms            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
                                         aten::empty         0.01%     126.417us         0.01%     126.417us       3.327us            38                                                          [[], [], [], [], [], []]  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
Self CPU time total: 1.062s
```

### Output 4

```
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                                                Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls                                                                      Input Shapes  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
                  aten::scaled_dot_product_attention         0.01%     233.747us       100.00%        1.923s     101.205ms            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
    aten::_scaled_dot_product_attention_math_for_mps        99.98%        1.922s        99.99%        1.923s     101.193ms            19  [[8, 64, 1024, 1024], [8, 64, 1024, 1024], [8, 64, 1024, 1024], [], [], [], [],   
                                         aten::empty         0.01%     205.504us         0.01%     205.504us       5.408us            38                                                          [[], [], [], [], [], []]  
----------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------------------  
Self CPU time total: 1.923s
```

In the above examples the total time to run 19 iterations of SDPA varies from 2.975ms to 1.923s. That's a 2-3 orders of magnitude difference in terms of performance.

`n` was set relatively low on purpose. The higher it is the more likely you are to see this issue.

My gut feeling is that this is not necessarily related to the SDPA code itself. Possibly related to #124850.

xref #135778

### Versions

PyTorch version: 2.5.0a0+gitdd47f6f
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.6.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.30.1
Libc version: N/A

Python version: 3.12.4 | packaged by conda-forge | (main, Jun 17 2024, 10:13:44) [Clang 16.0.6 ] (64-bit runtime)
Python platform: macOS-14.6.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M3 Max

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] flake8-bugbear==23.3.23
[pip3] flake8-comprehensions==3.15.0
[pip3] flake8-executable==2.1.3
[pip3] flake8-logging-format==0.9.0
[pip3] flake8-pyi==23.3.1
[pip3] flake8-simplify==0.19.3
[pip3] mypy==1.10.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.0
[pip3] optree==0.12.1
[pip3] torch==2.5.0a0+gitdd47f6f
[pip3] torch-tb-profiler==0.4.3
[pip3] torchvision==0.20.0a0+0d80848
[pip3] triton==3.0.0
[conda] numpy                     1.26.0                   pypi_0    pypi
[conda] optree                    0.12.1                   pypi_0    pypi
[conda] torch                     2.5.0a0+gitdd47f6f           dev_0    <develop>
[conda] torch-tb-profiler         0.4.3                    pypi_0    pypi
[conda] torchfix                  0.4.0                    pypi_0    pypi
[conda] torchvision               0.20.0a0+0d80848           dev_0    <develop>
[conda] triton                    3.0.0                    pypi_0    pypi

cc @msaroufim @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen