Repository: pytorch
Issue: Memory efficient attention for tensors where the last dimension is not divisible by 8
Labels: triaged,module: sdpa

### 🚀 The feature, motivation and pitch

Currently, using `scaled_dot_product_attention` and the memory efficient kernel requires that the last dimension of the inputs is divisible by 8. Typically, this corresponds to the dimension per head in multihead attention, for example when using the `[batch, head, seq, dim]` convention.

Using inputs that do not conform to this requirement results in a `RuntimeError: No available kernel.  Aborting execution.` and a warning: `UserWarning: Mem efficient attention requires last dimension of inputs to be divisible by 8.`

It would be great if this requirement could be relaxed, for example by only being divisible by 2. The [TPU implementation associated with the paper](https://github.com/google-research/google-research/tree/master/memory_efficient_attention) appears to work with arbitrary dimensions, but this might not be the case for GPUs.

It would also be helpful if these requirements would be documented (the [documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) appears to be missing in this regard).

### Alternatives

The Flash attention kernel supports this feature, but it is missing some others, e.g. attention masks.

### Additional context

A minimal example:

```python
import torch
import torch.nn.functional as F

qkv_size = (10, 128, 123, 2)

Q = torch.rand(size=qkv_size, device='cuda', dtype=torch.bfloat16)
K = torch.rand(size=qkv_size, device='cuda', dtype=torch.bfloat16)
V = torch.rand(size=qkv_size, device='cuda', dtype=torch.bfloat16)

with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
    O = F.scaled_dot_product_attention(Q, K, V, attn_mask=None, dropout_p=0)
```

The output
```
[/tmp/ipykernel_16779/975066207.py:2](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_16779/975066207.py:2): UserWarning: Memory efficient kernel not used because: (Triggered internally at [/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:350](https://file+.vscode-resource.vscode-cdn.net/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:350).)
  O = F.scaled_dot_product_attention(Q, K, V, attn_mask=None, dropout_p=0)
[/tmp/ipykernel_16779/975066207.py:2](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_16779/975066207.py:2): UserWarning: Mem efficient attention requires last dimension of inputs to be divisible by 8. Got Query.size(-1): 2, Key.size(-1): 2, Value.size(-1): 2 instead. (Triggered internally at [/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:128](https://file+.vscode-resource.vscode-cdn.net/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:128).)
  O = F.scaled_dot_product_attention(Q, K, V, attn_mask=None, dropout_p=0)
[/tmp/ipykernel_16779/975066207.py:2](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_16779/975066207.py:2): UserWarning: Flash attention kernel not used because: (Triggered internally at [/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:352](https://file+.vscode-resource.vscode-cdn.net/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:352).)
  O = F.scaled_dot_product_attention(Q, K, V, attn_mask=None, dropout_p=0)
[/tmp/ipykernel_16779/975066207.py:2](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_16779/975066207.py:2): UserWarning: Flash attention has been runtime disabled. (Triggered internally at [/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/sdp_utils_cpp.h:439](https://file+.vscode-resource.vscode-cdn.net/opt/conda/conda-bld/pytorch_1696146114277/work/aten/src/ATen/native/transformers/sdp_utils_cpp.h:439).)
  O = F.scaled_dot_product_attention(Q, K, V, attn_mask=None, dropout_p=0)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[34], line 2
      1 with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
----> 2     O = F.scaled_dot_product_attention(Q, K, V, attn_mask=None, dropout_p=0)

RuntimeError: No available kernel.  Aborting execution.
```

This is using PyTorch 2.2.0.dev20231001, CUDA 11.8, and an Ampere GPU.

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @mikaylagawarecki