Repository: pytorch
Issue: FlashAttention IMA
Labels: high priority,triaged,module: sdpa

# Summary

Recently a very large number of FlashAttention tests have become flaky. I am having a very hard time reproducing this in my local environment.
Example Tests:
- https://github.com/pytorch/pytorch/issues/131240
- https://github.com/pytorch/pytorch/issues/131227
- https://github.com/pytorch/pytorch/issues/131221

I am still having a hard time reproducing this. However I ran the test locally under compute sanitizer. And I am seeing an error reported for flash_attentions' split k kernel.

```
 PYTORCH_NO_CUDA_MEMORY_CACHING=1  compute-sanitizer --tool memcheck python test/test_transformers.py -k TestSDPACudaOnlyCUDA.test_flash_attention_vs_math_ref_grads_batch_size_1_seq_len_q_1024_seq_len_k_1024_head_dim_16_is_causal_False_dropout_p_0_0_bfloat16_scale_l1_cuda_bfloat16 > ima.txt
```

PyTorch Built with Debug = 1 produces a very large output

```
========= COMPUTE-SANITIZER
========= Invalid __global__ write of size 4 bytes
=========     at void pytorch_flash::flash_fwd_splitkv_kernel<pytorch_flash::Flash_fwd_kernel_traits<(int)32, (int)64, (int)256, (int)4, (bool)0, (bool)0, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<(int)32, (int)64, (int)256, (int)4, cutlass::bfloat16_t>>, (bool)0, (bool)0, (bool)0, (bool)0, (bool)0, (bool)1, (bool)0>(pytorch_flash::Flash_fwd_params)+0x10d40
=========     by thread (32,0,0) in block (0,0,0)
=========     Address 0x7f86621a4040 is out of bounds
=========     and is inside the nearest allocation at 0x7f86621a4000 of size 262,144 bytes
=========     Saved host backtrace up to driver entry point at kernel launch time
=========     Host Frame: [0x304fd2]
=========                in /lib64/libcuda.so.1
=========     Host Frame: [0x156e4]
=========                in /usr/local/cuda-12.3/lib64/libcudart.so.12
=========     Host Frame:cudaLaunchKernel [0x74631]
=========                in /usr/local/cuda-12.3/lib64/libcudart.so.12
=========     Host Frame:cudaError cudaLaunchKernel<char>(char const*, dim3, dim3, void**, unsigned long, CUstream_st*) in /usr/local/cuda-12.3/include/cuda_runtime.h:216 [0xf6874d2]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:__device_stub__ZN13pytorch_flash24flash_fwd_splitkv_kernelINS_23Flash_fwd_kernel_traitsILi32ELi64ELi256ELi4ELb0ELb0EN7cutlass10bfloat16_tENS_19Flash_kernel_traitsILi32ELi64ELi256ELi4ES3_EEEELb0ELb0ELb0ELb0ELb0ELb1ELb0EEEvNS_16Flash_fwd_paramsE(pytorch_flash::Flash_fwd_params const&) in /tmp/tmpxft_001e1e9e_00000000-6_flash_fwd_split_hdim32_bf16_sm80.cudafe1.stub.c:125 [0xf685470]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:void pytorch_flash::__wrapper__device_stub_flash_fwd_splitkv_kernel<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> >, false, false, false, false, false, true, false>(pytorch_flash::Flash_fwd_params const&) in /tmp/tmpxft_001e1e9e_00000000-6_flash_fwd_split_hdim32_bf16_sm80.cudafe1.stub.c:127 [0xf68548e]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:void pytorch_flash::flash_fwd_splitkv_kernel<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> >, false, false, false, false, false, true, false>(pytorch_flash::Flash_fwd_params) in /home/drisspg/meta/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h:50 [0xf6876e2]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:pytorch_flash::run_flash_splitkv_fwd<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> > >(pytorch_flash::Flash_fwd_params&, CUstream_st*)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const in /home/drisspg/meta/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h:111 [0xf68c084]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:pytorch_flash::run_flash_splitkv_fwd<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> > >(pytorch_flash::Flash_fwd_params&, CUstream_st*)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const in /home/drisspg/meta/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h:111 [0xf68bfe5]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:pytorch_flash::run_flash_splitkv_fwd<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> > >(pytorch_flash::Flash_fwd_params&, CUstream_st*)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const in /home/drisspg/meta/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h:111 [0xf68bf9b]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:pytorch_flash::run_flash_splitkv_fwd<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> > >(pytorch_flash::Flash_fwd_params&, CUstream_st*)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const in /home/drisspg/meta/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h:111 [0xf68bde7]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:pytorch_flash::run_flash_splitkv_fwd<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> > >(pytorch_flash::Flash_fwd_params&, CUstream_st*)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const in /home/drisspg/meta/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h:111 [0xf68bd4d]
=========                in /home/drisspg/meta/pytorch/torch/lib/libtorch_cuda.so
=========     Host Frame:pytorch_flash::run_flash_splitkv_fwd<pytorch_flash::Flash_fwd_kernel_traits<32, 64, 256, 4, false, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<32, 64, 256, 4, cutlass::bfloat16_t> > >(pytorch_flash::Flash_fwd_params&, CUstream_st*)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda()#1}::operator()() const in /home/drisspg/meta/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h:111 [0xf68bcc7]
```

### Smaller Repro:
```Python
import torch
import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend
def test_flash_attention():
    device = "cuda"
    dtype = torch.bfloat16
    batch_size = 1
    n_heads = 1
    seq_len_q = 1
    seq_len_k = 1024
    head_dim = 16
    is_causal = False
    dropout_p = 0.0
    scale = (1 / head_dim)

    # Create input tensors
    query = torch.rand(batch_size, n_heads, seq_len_q, head_dim, device=device, dtype=dtype, requires_grad=True)
    key = torch.rand(batch_size, n_heads, seq_len_k, head_dim, device=device, dtype=dtype, requires_grad=True)
    value = torch.rand(batch_size, n_heads, seq_len_k, head_dim, device=device, dtype=dtype, requires_grad=True)

    with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
        out = F.scaled_dot_product_attention(query, key, value, dropout_p=dropout_p, is_causal=is_causal, scale=scale)

    print("Output shape:", out.shape)
    print("Output sum:", out.sum().item())

if __name__ == "__main__":
    test_flash_attention()

```
## Updates
* Same error Occurs on 2.3.1 , so not sure if the flaky tests are the same as this..
* Trying a PR where we disable splitkv kernel. This solves the ima locally but I want to know if they flaky tests are related: https://github.com/pytorch/pytorch/pull/131267
* Same bug with upstream split_k kernel: https://github.com/Dao-AILab/flash-attention/issues/1076
* Current best guess: https://github.com/pytorch/pytorch/pull/131277

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim