Repository: pytorch
Issue: Scaled Dot-Product Attention Invalid Configuration Error on Large batch size
Labels: triaged,module: sdpa

# Scaled Dot-Product Attention Invalid Configuration Error on Large batch size

## Summary
The `torch.nn.functional.scaled_dot_product_attention` (sdpa) function is not working as expected when the batch size is large. It causes a `RuntimeError: CUDA error: invalid configuration argument`. This problem affects also the `torch.nn.TransformerEncoderLayer` as it relies on the `scaled_dot_product_attention` function.


## Reproducing the error

The following scripts are run with:
```bash
CUDA_LAUNCH_BLOCKING=1 python3 test.py
```
However, the same results can be obtained without setting `CUDA_LAUNCH_BLOCKING=1`.

### SDPA Example

This code will raise the error:
```python
import torch
x = torch.rand(100000,1,10,16,dtype=torch.float32,device="cuda:0")
a = torch.nn.functional.scaled_dot_product_attention(x,x,x)
```
```bash
Traceback (most recent call last):
  File "/path/to/repo/test.py", line 3, in <module>
    a = torch.nn.functional.scaled_dot_product_attention(x,x,x)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: invalid configuration argument
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

### TransformerEncoderLayer Example
This code will raise the error:
```python
import torch
transformer_encoder = torch.nn.TransformerEncoderLayer(
    d_model         = 16      ,
    nhead           = 1       ,
    dim_feedforward = 64      ,
    device          = "cuda:0",
    batch_first     = True    ,
)
x = torch.rand(100000,10,16,dtype=torch.float32,device="cuda:0")
transformer_encoder(x)
```
```bash
Traceback (most recent call last):
  File "path/to/repo/test.py", line 10, in <module>
    transformer_encoder(x)
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/modules/transformer.py", line 904, in forward
    + self._sa_block(x, src_mask, src_key_padding_mask, is_causal=is_causal)
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/modules/transformer.py", line 918, in _sa_block
    x = self.self_attn(
        ^^^^^^^^^^^^^^^
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/modules/activation.py", line 1368, in forward
    attn_output, attn_output_weights = F.multi_head_attention_forward(
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "path/to/repo/venv/lib/python3.12/site-packages/torch/nn/functional.py", line 6278, in multi_head_attention_forward
    attn_output = scaled_dot_product_attention(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: invalid configuration argument
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

## Cause of The Error
I searched a little bit for the cause of the error. Here is what I found:

The error is caused in file `pytorch/aten/src/ATen/native/transformers/cuda/attention.cu`. In particular the `launchKernel` lambda inside `_efficient_attention_forward` is responsible for the error. To my understanding this is the code snippet that triggers the invalid configuration:
```cuda
    auto blocks = p.getBlocksGrid();
    if (blocks.x * blocks.y * blocks.z == 0 || key.size(1) == 0) {
      res.zero_();
      return;
    }
    Kernel::check_supported(p);
    kernel_fn<<<blocks, p.getThreadsGrid(), smem_bytes, stream>>>(p);
```
It appears that, with a large batch size this kernel is run with a number of blocks that exceeds the maximum number of kernel blocks, `65,535`. This triggers the invalid configuration error. As a matter of fact, running the snippet
```python
import torch
x = torch.rand(65535,1,10,16,dtype=torch.float32,device="cuda:0")
a = torch.nn.functional.scaled_dot_product_attention(x,x,x)
```
will not raise the error. Doing the same with a batch size of `65536` will raise the error.

Note that also the number of heads contributes to the number of cuda blocks. So also an increased number of heads will trigger the error. 

## SDPA Workaround
The error can be avoided by simply removing the number of heads dimension in the input tensor. For example,
```python
import torch
#x = torch.rand(100000,1,10,16,dtype=torch.float32,device="cuda:0")
x = torch.rand(100000,10,16,dtype=torch.float32,device="cuda:0")
a = torch.nn.functional.scaled_dot_product_attention(x,x,x)
```
Note that, the pytorch documentation does not mention that it is possible to remove the number of heads dimension in the `torch.nn.functional.scaled_dot_product_attention`. However, by flattening the batch dimension with the heads dimension resolves the issue. To my understading this operation should not affect the result as demonstrated in the following code snippet:
```python
mport torch
x0 = torch.rand(1000,2,10,16,dtype=torch.float32,device="cuda:0")
x1 = x0.clone().flatten(0,1)
a = torch.nn.functional.scaled_dot_product_attention(x0,x0,x0)
b = torch.nn.functional.scaled_dot_product_attention(x1,x1,x1)
print(a.flatten().isclose(b.flatten(), atol=1e-4).all()) # >>> tensor(True, device='cuda:0')
```

## TransformerEncoderLayer Workaround
The workaround for the `torch.nn.TransformerEncoderLayer` requires to modify the internal python implementation. In particular, we need to modify the `_sa_block` method in the `torch.nn.modules.transformer.TransformerEncoderLayer` class. The following code snippet demonstrates the workaround. I comment only on the modified lines:
```python 
def _sa_block(
    self,
    x: Tensor,
    attn_mask: Optional[Tensor],
    key_padding_mask: Optional[Tensor],
    is_causal: bool = False,
) -> Tensor:
    shape = x.shape ### added line
    x = x.flatten(0,1) ### added line
    x = self.self_attn(
        x,
        x,
        x,
        attn_mask=attn_mask,
        key_padding_mask=key_padding_mask,
        need_weights=False,
        is_causal=is_causal,
    )[0]
    x = x.view(shape) ### added line
    return self.dropout1(x)
```

This resolves the transformer issue. It simply flattens the batch dimension with heads dimension and restores the original shape after the attention operation. I do not comment on the efficiency of doing this as it may have some performance implications. As a more general solution, one could check if the number of blocks used would exceed the maximum number of blocks and if so, flatten the batch dimension with the heads dimension.

## Environment
I have tested this issue on torch 2.5.1 version. I have also tested the issue on the main branch (at the time of writing 2.6.0a0+git65c2086). 

Result from `collect_env.py` using the 2.5.1:
```
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: Arch Linux (x86_64)
GCC version: (GCC) 14.2.1 20240910
Clang version: 18.1.8
CMake version: Could not collect
Libc version: glibc-2.40

Python version: 3.12.7 (main, Oct  1 2024, 11:15:50) [GCC 14.2.1 20240910] (64-bit runtime)
Python platform: Linux-6.12.1-arch1-1-x86_64-with-glibc2.40
Is CUDA available: True
CUDA runtime version: 12.6.85
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4080
Nvidia driver version: 565.57.01
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:                        39 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               16
On-line CPU(s) list:                  0-15
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Core(TM) i7-10700K CPU @ 3.80GHz
CPU family:                           6
Model:                                165
Thread(s) per core:                   2
Core(s) per socket:                   8
Socket(s):                            1
Stepping:                             5
CPU(s) scaling MHz:                   92%
CPU max MHz:                          5100.0000
CPU min MHz:                          800.0000
BogoMIPS:                             7602.45
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp vnmi pku ospke md_clear flush_l1d arch_capabilities
Virtualization:                       VT-x
L1d cache:                            256 KiB (8 instances)
L1i cache:                            256 KiB (8 instances)
L2 cache:                             2 MiB (8 instances)
L3 cache:                             16 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-15
Vulnerability Gather data sampling:   Mitigation; Microcode
Vulnerability Itlb multihit:          KVM: Mitigation: Split huge pages
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds:                  Mitigation; Microcode
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==2.1.3
[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

```