Repository: pytorch
Issue: performance bug: flex attention much slower than dense attention
Labels: triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher,module: flex attention

### 🐛 Describe the bug

With a 2D spatial neighborhood pattern, flash attention is orders of magnitude slower than dense attention:

hlc=2
seq_length : 192
flex attention : 0.0015106382369995117 [s]
dense attention : 3.8884878158569336e-05 [s]
 
 hlc=3
seq_length : 768
flex attention : 0.0015071055889129639 [s]
dense attention : 3.041529655456543e-05 [s]
 
hlc=4
seq_length : 3072
flex attention : 0.020486905336380003 [s]
dense attention : 3.140068054199219e-05 [s]

hlc is a parameter that, essentially, controls the number of cells. The sparsity pattern is 1-ring neighborhood of cells.



```
import code
import time

import warnings
import numpy as np
import torch
from torch.nn.attention.flex_attention import flex_attention, create_mask, create_block_mask

import astropy_healpix as hp


hlc = 3
num_healpix_cells = 12 * 4**hlc 
print( f'hlc={hlc}')
print( f'seq_length : {num_healpix_cells}')
num_heads = 8
dim_embed = 128
bs = 4

q = torch.ones( bs, num_heads, num_healpix_cells, dim_embed, dtype=torch.float16, device='cuda')
k = torch.ones( bs, num_heads, num_healpix_cells, dim_embed, dtype=torch.float16, device='cuda')
v = torch.ones( bs, num_heads, num_healpix_cells, dim_embed, dtype=torch.float16, device='cuda')

with warnings.catch_warnings(action="ignore"):
    nbours= hp.neighbours( np.arange(num_healpix_cells), 2**hlc, order='nested').transpose()
# build adjacency matrix (smarter ways to do it ...)
nbours_mat = torch.zeros( (num_healpix_cells,num_healpix_cells), dtype=torch.bool, device='cuda')
for i in range(num_healpix_cells) :
    for j in nbours[i] :
        nbours_mat[i,j] = True if j>=0 else False

# create sparse block matrix for flex attention
def sparse_mask(b, h, q_idx, kv_idx):
    # return ddkv_idx in nbours[q_idx]
    return nbours_mat[q_idx,kv_idx]
block_mask = create_block_mask( sparse_mask, B=None, H=None, Q_LEN=dim_embed, KV_LEN=dim_embed)

# experiments

# warmup
for i in range( 10):
  qp = flex_attention( q, k, v, block_mask=block_mask)

t_start = time.time()
for i in range( 1000):
  qp = flex_attention( q, k, v, block_mask=block_mask)
print( f'flex attention : {(time.time() - t_start) / 1000.} [s]', flush=True)

# warmup
for i in range( 10):
  with torch.nn.attention.sdpa_kernel( torch.nn.attention.SDPBackend.FLASH_ATTENTION) :
    qp = torch.nn.functional.scaled_dot_product_attention( q, k, v)

t_start = time.time()
for i in range( 1000):
  with torch.nn.attention.sdpa_kernel( torch.nn.attention.SDPBackend.FLASH_ATTENTION) :
    qp = torch.nn.functional.scaled_dot_product_attention( q, k, v)
print( f'dense attention : {(time.time() - t_start) / 1000.} [s]', flush=True)


```

### Versions

Collecting environment information...
PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux release 8.8 (Ootpa) (x86_64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-18)
Clang version: 15.0.7 (Red Hat 15.0.7-1.module+el8.8.0+17939+b58878af)
CMake version: version 3.20.2
Libc version: glibc-2.28

Python version: 3.11.10 (main, Sep 27 2024, 08:55:04) [GCC 8.5.0 20210514 (Red Hat 8.5.0-18)] (64-bit runtime)
Python platform: Linux-4.18.0-477.43.1.el8_8.x86_64-x86_64-with-glibc2.28
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:
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              256
On-line CPU(s) list: 0-255
Thread(s) per core:  2
Core(s) per socket:  64
Socket(s):           2
NUMA node(s):        8
Vendor ID:           AuthenticAMD
CPU family:          23
Model:               49
Model name:          AMD EPYC 7H12 64-Core Processor
Stepping:            0
CPU MHz:             2600.000
CPU max MHz:         2600.0000
CPU min MHz:         1500.0000
BogoMIPS:            5200.23
Virtualization:      AMD-V
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            16384K
NUMA node0 CPU(s):   0-15,128-143
NUMA node1 CPU(s):   16-31,144-159
NUMA node2 CPU(s):   32-47,160-175
NUMA node3 CPU(s):   48-63,176-191
NUMA node4 CPU(s):   64-79,192-207
NUMA node5 CPU(s):   80-95,208-223
NUMA node6 CPU(s):   96-111,224-239
NUMA node7 CPU(s):   112-127,240-255
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 nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es

Versions of relevant libraries:
[pip3] numpy==2.1.2
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] torch==2.5.1
[pip3] triton==3.1.0
[conda] Could not collect

cc @ezyang @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh @yf225 @Chillee @drisspg @yanboliang @BoyuanFeng