Runtime Error when using Flash Attention

chaoyuaw: <p>I ran into the following runtime error when trying to use Flash Attention for <code>scaled_dot_product_attention</code>. I tried getting PyTorch from “stable (conda)”, “nightly (conda)”, and “compile from source”, and all gave me the same error.</p>
<p>Thanks for any advice in advance!</p>
<p><strong>Minimal code to reproduce:</strong></p>
<pre><code class="lang-auto">import torch
Q = torch.zeros(3, 10, 128).cuda()
K = torch.zeros(3, 10, 128).cuda()
V = torch.zeros(3, 10, 128).cuda()
with torch.backends.cuda.sdp_kernel(
    enable_flash=True,
    enable_math=False,
    enable_mem_efficient=False,
):
    print(torch.nn.functional.scaled_dot_product_attention(Q, K, V).shape)
</code></pre>
<p><strong>Error log:</strong></p>
<pre><code class="lang-auto">test_sdp.py:10: UserWarning: Memory efficient kernel not used because: (Triggered internally at /home/conda/feedstock_root/build_artifacts/pytorch-recipe_1680527322149/work/aten/src/ATen/native/transformers/cuda/sdp_utils.h:527.)
  print(torch.nn.functional.scaled_dot_product_attention(Q, K, V).shape)
Traceback (most recent call last):
  File "/data/home/cywu/test_sdp.py", line 10, in &lt;module&gt;
    print(torch.nn.functional.scaled_dot_product_attention(Q, K, V).shape)
RuntimeError: Torch was not compiled with flash attention.
</code></pre>
<p><strong>Versions</strong><br>
Collecting environment information…<br>
PyTorch version: 2.0.0.post200<br>
Is debug build: False<br>
CUDA used to build PyTorch: 11.2<br>
ROCM used to build PyTorch: N/A</p>
<p>OS: Ubuntu 20.04.5 LTS (x86_64)<br>
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0<br>
Clang version: Could not collect<br>
CMake version: version 3.25.2<br>
Libc version: glibc-2.31</p>
<p>Python version: 3.9.16 | packaged by conda-forge | (main, Feb  1 2023, 21:39:03)  [GCC 11.3.0] (64-bit runtime)<br>
Python platform: Linux-5.15.0-1030-aws-x86_64-with-glibc2.31<br>
Is CUDA available: True<br>
CUDA runtime version: 11.8.89<br>
CUDA_MODULE_LOADING set to: LAZY<br>
GPU models and configuration:<br>
GPU 0: NVIDIA A100-SXM4-80GB<br>
GPU 1: NVIDIA A100-SXM4-80GB<br>
GPU 2: NVIDIA A100-SXM4-80GB<br>
GPU 3: NVIDIA A100-SXM4-80GB<br>
GPU 4: NVIDIA A100-SXM4-80GB<br>
GPU 5: NVIDIA A100-SXM4-80GB<br>
GPU 6: NVIDIA A100-SXM4-80GB<br>
GPU 7: NVIDIA A100-SXM4-80GB</p>
<p>Nvidia driver version: 525.85.12<br>
cuDNN version: Could not collect<br>
HIP runtime version: N/A<br>
MIOpen runtime version: N/A<br>
Is XNNPACK available: True</p>
<p>CPU:<br>
Architecture:                    x86_64<br>
CPU op-mode(s):                  32-bit, 64-bit<br>
Byte Order:                      Little Endian<br>
Address sizes:                   46 bits physical, 48 bits virtual<br>
CPU(s):                          96<br>
On-line CPU(s) list:             0-95<br>
Thread(s) per core:              2<br>
Core(s) per socket:              24<br>
Socket(s):                       2<br>
NUMA node(s):                    2<br>
Vendor ID:                       GenuineIntel<br>
CPU family:                      6<br>
Model:                           85<br>
Model name:                      Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz<br>
Stepping:                        7<br>
CPU MHz:                         1263.993<br>
BogoMIPS:                        5999.99<br>
Hypervisor vendor:               KVM<br>
Virtualization type:             full<br>
L1d cache:                       1.5 MiB<br>
L1i cache:                       1.5 MiB<br>
L2 cache:                        48 MiB<br>
L3 cache:                        71.5 MiB<br>
NUMA node0 CPU(s):               0-23,48-71<br>
NUMA node1 CPU(s):               24-47,72-95<br>
Vulnerability Itlb multihit:     KVM: Mitigation: VMX unsupported<br>
Vulnerability L1tf:              Mitigation; PTE Inversion<br>
Vulnerability Mds:               Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown<br>
Vulnerability Meltdown:          Mitigation; PTI<br>
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown<br>
Vulnerability Retbleed:          Vulnerable<br>
Vulnerability Spec store bypass: Vulnerable<br>
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization<br>
Vulnerability Spectre v2:        Mitigation; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected<br>
Vulnerability Srbds:             Not affected<br>
Vulnerability Tsx async abort:   Not affected<br>
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves ida arat pku ospke</p>
<p>Versions of relevant libraries:<br>
[pip3] numpy==1.24.2<br>
[pip3] pytorch3d==0.7.3<br>
[pip3] torch==2.0.0.post200<br>
[pip3] torchaudio==2.0.0<br>
[pip3] torchvision==0.14.1a0+59d9189<br>
[conda] cudatoolkit               11.3.1               h2bc3f7f_2<br>
[conda] libmagma                  2.7.1                hc72dce7_1    conda-forge<br>
[conda] libmagma_sparse           2.7.1                hc72dce7_2    conda-forge<br>
[conda] magma                     2.7.1                ha770c72_2    conda-forge<br>
[conda] mkl                       2022.2.1         h84fe81f_16997    conda-forge<br>
[conda] numpy                     1.24.2           py39h7360e5f_0    conda-forge<br>
[conda] pytorch                   2.0.0           cuda112py39ha9981d0_200    conda-forge<br>
[conda] pytorch-cuda              11.8                 h7e8668a_3    pytorch-nightly<br>
[conda] pytorch-mutex             1.0                        cuda    pytorch<br>
[conda] pytorch3d                 0.7.3                     dev_0    <br>
[conda] torch                     2.1.0a0+git9a2a6fc           dev_0    <br>
[conda] torchaudio                2.0.0                py39_cu118    pytorch<br>
[conda] torchvision               0.14.1          cuda112py39hb350dc8_1    conda-forge</p>

ptrblck: <p>It seems you are using <a href="https://anaconda.org/conda-forge/pytorch/files">unofficial conda binaries</a> from <code>conda-forge</code> created by <code>mark.harfouche</code>, which do not seem to ship with FlashAttention.<br>
You can see it by the custom tag:</p>
<pre><code class="lang-python">PyTorch version: 2.0.0.post200
...
[conda] pytorch 2.0.0 cuda112py39ha9981d0_200 conda-forge
</code></pre>
<p>as well as an old CUDA runtime, which we are not using anymore for our current builds:</p>
<pre><code class="lang-python">CUDA used to build PyTorch: 11.2
</code></pre>