Repository: pytorch
Issue: Performance degradation in scaled_dot_product_attention with attn_mask and sequence length multiples of 16
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

There's a performance degradation when using `scaled_dot_product_attention` with the `attn_mask` argument, when the sequence length is a multiple of 16. This issue can be reproduced using the following code snippet.

**Reproducer Code**
```
import torch
from torch._inductor.runtime.benchmarking import benchmarker
from torch.nn import functional as F

def run(seqlen):
    with torch.device("cuda"):
        def f(q, k, v, mask):
            return F.scaled_dot_product_attention(
                q, k, v, attn_mask=mask, dropout_p=0.0
            )
        f_compiled = torch.compile(f)

        # Create inputs
        bsz = 32
        q = torch.randn(bsz, 16, seqlen, 64, dtype=torch.bfloat16)
        k = torch.randn(bsz, 16, seqlen, 64, dtype=torch.bfloat16)
        v = torch.randn(bsz, 16, seqlen, 64, dtype=torch.bfloat16)
        mask = torch.ones([bsz, 1, seqlen, seqlen], dtype=torch.bool)
        inputs = [q, k, v, mask]

        # Benchmark
        time = benchmarker.benchmark_gpu(lambda: f_compiled(*inputs), warmup=5, rep=50)
        return time

for seqlen_start in [1008, 1024, 2048, 4096]:
    for offset in range(-1, 2):
        seqlen = seqlen_start + offset

        torch._dynamo.reset()
        time = run(seqlen)

        print(seqlen, time)
    print()
```

**Output on H100 GPU**
```
1007 1.569983959197998
1008 2.0037760734558105
1009 1.5577600002288818

1023 1.553056001663208
1024 2.000607967376709
1025 1.7111680507659912

2047 6.071455955505371
2048 8.064703941345215
2049 6.349376201629639

4095 23.773408889770508
4096 35.05900955200195
4097 24.331039428710938
```

**Analysis**
The results show that incrementing the sequence length from multiples of 16 (1024, 2048, 4096) to non-multiples of 16 (1025, 2049, 4097) results in up to 1.43x speedup. This is counterintuitive and suggests that the selected/generated kernel for sequence lengths of multiples of 16 could be improved.

**Expected**
I'd expect multiples of 16 to perform equal, or even better than neighboring sizes, because of e.g. better divisibility for tiling.

### Error logs

_No response_

### Versions

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: CentOS Stream 9 (x86_64)
GCC version: (GCC) 11.5.0 20240719 (Red Hat 11.5.0-2)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.34

Python version: 3.10.15 (main, Oct  3 2024, 07:27:34) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.4.3-0_fbk14_hardened_2601_gcd42476b84e9-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA H100
GPU 1: NVIDIA H100
GPU 2: NVIDIA H100
GPU 3: NVIDIA H100

Nvidia driver version: 550.90.07
cuDNN version: Probably one of the following:
/usr/lib64/libcudnn.so.8.8.0
/usr/lib64/libcudnn_adv_infer.so.8.8.0
/usr/lib64/libcudnn_adv_train.so.8.8.0
/usr/lib64/libcudnn_cnn_infer.so.8.8.0
/usr/lib64/libcudnn_cnn_train.so.8.8.0
/usr/lib64/libcudnn_ops_infer.so.8.8.0
/usr/lib64/libcudnn_ops_train.so.8.8.0
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:                      52 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             184
On-line CPU(s) list:                0-183
Vendor ID:                          AuthenticAMD
Model name:                         AMD EPYC 9654 96-Core Processor
CPU family:                         25
Model:                              17
Thread(s) per core:                 1
Core(s) per socket:                 184
Socket(s):                          1
Stepping:                           1
BogoMIPS:                           4792.80
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 rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw perfctr_core invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 clzero xsaveerptr wbnoinvd arat npt lbrv nrip_save tsc_scale vmcb_clean pausefilter pfthreshold v_vmsave_vmload vgif avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm arch_capabilities
Virtualization:                     AMD-V
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          11.5 MiB (184 instances)
L1i cache:                          11.5 MiB (184 instances)
L2 cache:                           92 MiB (184 instances)
L3 cache:                           2.9 GiB (184 instances)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-183
Vulnerability Gather data sampling: Not affected
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:    Vulnerable
Vulnerability Spectre v1:           Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:           Vulnerable, IBPB: disabled, STIBP: disabled, PBRSB-eIBRS: Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==2.2.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 @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov @eellison @yanboliang @ezyang @BoyuanFeng