Repository: pytorch/pytorch

Issue:
Backpropagation of `torch.compile`d module with scaled_dot_product_attention fails unless q, k, v shapes are the same
### 🐛 Describe the bug

When calling `torch.compile` on a module that uses `nn.functional.scaled_dot_product_attention(q, k, v, ...)`, backward passes do not work unless `q`, `k`, and `v` all have the same shape. This error only happens if CUDA is used; CPU appears to be fine.

```python
import torch
from torch import nn

class BadModule(nn.Module):
  def __init__(self):
    super().__init__()
    self.l = nn.Linear(8, 16, bias=False)

  def forward(self, x):
    return torch.nn.functional.scaled_dot_product_attention(x, x, self.l(x), is_causal=True)

m = torch.compile(BadModule().to('cuda'))
x = torch.rand(1, 2, 4, 8).to('cuda') # (batch, num_heads, seq_len, embed_dim)
y = m(x) # forward pass
z = y.mean().backward() # fake backward pass
```

resulting error:

```
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError: Function ScaledDotProductEfficientAttentionBackward0 returned an invalid gradient at index 2 - got [1, 2, 4, 8] but expected shape compatible with [1, 2, 4, 16]
```

It works fine if the call to `torch.compile(...)` is removed or on the CPU backend.

### Versions

```
PyTorch version: 2.1.0+cu118                                                                                                                                                                                      Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 13.1.1 20230429
Clang version: 15.0.7
CMake version: version 3.26.3
Libc version: glibc-2.37

Python version: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:40:35) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.37
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080 Laptop GPU
Nvidia driver version: 516.94
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   48 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          16
On-line CPU(s) list:             0-15
Vendor ID:                       AuthenticAMD
Model name:                      AMD Ryzen 9 5900HX with Radeon Graphics
CPU family:                      25
Model:                           80
Thread(s) per core:              2
Core(s) per socket:              8
Socket(s):                       1
Stepping:                        0
BogoMIPS:                        6587.45
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr arat npt nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload umip vaes vpclmulqdq rdpid fsrm
Virtualization:                  AMD-V
Hypervisor vendor:               Microsoft
Virtualization type:             full
L1d cache:                       256 KiB (8 instances)
L1i cache:                       256 KiB (8 instances)
L2 cache:                        4 MiB (8 instances)
L3 cache:                        16 MiB (1 instance)
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] torch==2.1.0+cu118
[pip3] torchaudio==2.1.0+cu118
[pip3] torchvision==0.16.0+cu118
[pip3] triton==2.1.0
[conda] numpy                     1.24.1                   pypi_0    pypi
[conda] torch                     2.1.0+cu118              pypi_0    pypi
[conda] torchaudio                2.1.0+cu118              pypi_0    pypi
[conda] torchvision               0.16.0+cu118             pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi
```

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @mikaylagawarecki @msaroufim @bdhirsh @anijain2305 @chauhang @wconstab


Discussion:

Resolving patch:
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
--- a/torch/_meta_registrations.py
+++ b/torch/_meta_registrations.py
@@ -4961,6 +4961,7 @@ def meta__scaled_dot_product_efficient_backward(
     num_heads = query.size(1)
     max_q = query.size(2)
     head_dim = query.size(3)
+    head_dim_v = value.size(3)
 
     max_k = key.size(2)
 
@@ -4977,7 +4978,7 @@ def meta__scaled_dot_product_efficient_backward(
         device=key.device,
     )
     grad_v = torch.empty_permuted(
-        (batch_size, num_heads, max_k, head_dim),
+        (batch_size, num_heads, max_k, head_dim_v),
         (0, 2, 1, 3),
         dtype=value.dtype,
         device=value.device,
diff --git a/torch/_subclasses/fake_utils.py b/torch/_subclasses/fake_utils.py
--- a/torch/_subclasses/fake_utils.py
+++ b/torch/_subclasses/fake_utils.py
@@ -146,6 +146,18 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                             r_out, fake_out, check_strides=self.check_strides
                         )
                     except Exception as e:
+                        if (
+                            func is aten._scaled_dot_product_flash_attention.default
+                            and idx in (6, 7)
+                            and "Devices" in repr(e)
+                        ):
+                            continue
+                        if (
+                            func is aten._scaled_dot_product_efficient_attention.default
+                            and idx in (2, 3)
+                            and "Devices" in repr(e)
+                        ):
+                            continue
                         error_message = (
                             f"{context} mismatched tensor metadata: {e}"
                             if len(r_flat) == 1


Test patch:
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
--- a/torch/testing/_internal/common_methods_invocations.py
+++ b/torch/testing/_internal/common_methods_invocations.py
@@ -27,7 +27,7 @@
      toleranceOverride, tol)
 from torch.testing._internal.common_cuda import (
     SM53OrLater, SM60OrLater, SM80OrLater, with_tf32_off, TEST_CUDNN,
-    _get_torch_cuda_version, _get_torch_rocm_version, PLATFORM_SUPPORTS_FUSED_SDPA
+    _get_torch_cuda_version, _get_torch_rocm_version,
 )
 from torch.testing._internal.common_utils import (
     make_fullrank_matrices_with_distinct_singular_values,
@@ -13499,11 +13499,6 @@ def reference_flatten(input, start_dim=0, end_dim=-1):
                          device_type='cpu'),
             # OpInfo was implemented with a lambda
             DecorateInfo(unittest.skip("Skipped!"), 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'),
-            # See [Note] SDPA returns Philox Offset and Seed as tensors that will live on CPU when not in cuda graph capture
-            DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake_crossref_backward_amp',
-                         device_type='cuda', dtypes=(torch.float32,), active_if=PLATFORM_SUPPORTS_FUSED_SDPA),
-            DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake_crossref_backward_no_amp',
-                         device_type='cuda', dtypes=(torch.float32,), active_if=PLATFORM_SUPPORTS_FUSED_SDPA),
             # TODO Need to understand what this is testing and why it doesn't work
             DecorateInfo(unittest.skip("Skipped"), 'TestDecomp', 'test_comprehensive'),
             DecorateInfo(unittest.skip('output is non-deterministic (when dropout_p > 0)'), 'TestCommon', 'test_compare_cpu'),