RuntimeError: FlashAttention only supports Ampere GPUs or newer.
Author: bilalghanem
Pull request: no

### System Info

I am trying to run the following code:

```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Configs
device = "cuda:7"
model_name = "openchat/openchat_3.5"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, 
                                             bnb_4bit_compute_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')

```

I can load the model completely fine, but when I want to generate, I get this error:

> ---------------------------------------------------------------------------
> RuntimeError                              Traceback (most recent call last)
> Cell In[3], [line 76](vscode-notebook-cell:?execution_count=3&line=76)
>      [74](vscode-notebook-cell:?execution_count=3&line=74) model_input_text = template.format(start, html_, end)
>      [75](vscode-notebook-cell:?execution_count=3&line=75) model_inputs = tokenizer([model_input_text], return_tensors="pt", padding=False).to(device)
> ---> [76](vscode-notebook-cell:?execution_count=3&line=76) generated_ids = model.generate(**model_inputs, do_sample=True, top_p=1.0, temperature=0.8, top_k=50, max_new_tokens=1024)
>      [77](vscode-notebook-cell:?execution_count=3&line=77) model_outputs_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
>      [78](vscode-notebook-cell:?execution_count=3&line=78) print(model_outputs_text[model_input_text.rindex("GPT4 Correct Assistant:")+10:])
> 
> File [~/PATH/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py:115](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py:115), in context_decorator.<locals>.decorate_context(*args, **kwargs)
>     [112](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py:112) @functools.wraps(func)
>     [113](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py:113) def decorate_context(*args, **kwargs):
>     [114](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py:114)     with ctx_factory():
> --> [115](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py:115)         return func(*args, **kwargs)
> 
> File [~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1764](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1764), in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
>    [1756](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1756)     input_ids, model_kwargs = self._expand_inputs_for_generation(
>    [1757](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1757)         input_ids=input_ids,
>    [1758](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1758)         expand_size=generation_config.num_return_sequences,
>    [1759](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1759)         is_encoder_decoder=self.config.is_encoder_decoder,
>    [1760](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1760)         **model_kwargs,
>    [1761](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1761)     )
>    [1763](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1763)     # 13. run sample
> -> [1764](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1764)     return self.sample(
>    [1765](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/transformers/generation/utils.py:1765)         input_ids,
> ...
>      [58](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/flash_attn/flash_attn_interface.py:58)     None,
>      [59](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/flash_attn/flash_attn_interface.py:59) )
>      [60](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a22544e414970726d676d743031227d.vscode-resource.vscode-cdn.net/PATH/notebooks/~/PATH/venv/lib/python3.8/site-packages/flash_attn/flash_attn_interface.py:60) return out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state
> 
> RuntimeError: FlashAttention only supports Ampere GPUs or newer.


I am working on Ubuntu 20.04 with NVIDIA Quadro RTX 5000.
Cuda version: 12.2
NVIDIA-SMI 535.129.03
torch==2.1.2
transformers==4.36.2


### Who can help?

@SunMarc @younesbelkada

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Loading an LLM model with enabling fast attention.

### Expected behavior

Generate text.

Comments:
I got the same error on kaggle's notebook.

Hey! Not sure I understand the issue? 
> RuntimeError: FlashAttention only supports Ampere GPUs or newer. 

Means that flash attention implementation that you install does not support your GPU yet! (either too old or too new). 
I would rather look into the flash attention repo for the support to specific hardware not here! 🤗 

I hit the same issue on Kaggle notebook. I'm using FlashAttention2 and the official repo README.md said

```
FlashAttention-2 currently supports:

Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). Support for Turing GPUs (T4, RTX 2080) is coming soon, please use FlashAttention 1.x for Turing GPUs for now.
Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).
All head dimensions up to 256. Head dim > 192 backward requires A100/A800 or H100/H800.
```

https://github.com/Dao-AILab/flash-attention

So, we cannot use it on Kaggle environment, especially with b16.

Is there a way to check through python if flash attention is supported?
I just want my code to use this parameter when deployed in a machine where flash attention is supported and vice versa.


@ahassaine If a models supports flash attention, it will have the private attribute `_supports_flash_attn_2` set to `True` e.g. [like here for bark](https://github.com/huggingface/transformers/blob/39c3c0a72af6fbda5614dde02ff236069bb79827/src/transformers/models/bark/modeling_bark.py#L487).

> @ahassaine If a models supports flash attention, it will have the private attribute `_supports_flash_attn_2` set to `True` e.g. [like here for bark](https://github.com/huggingface/transformers/blob/39c3c0a72af6fbda5614dde02ff236069bb79827/src/transformers/models/bark/modeling_bark.py#L487).

I think he means, to see if the gpu supports flash attention imp.

@bilalghanem My mistake! Read too quickly. I don't think there's an easy way to do this as [the checks are in cpp](https://github.com/Dao-AILab/flash-attention/blob/197f2083a2f0953af9319cf4ce32d0bf2aae4bd8/csrc/flash_attn/flash_api.cpp#L303). 

@amyeroberts 
That's actually very helpful. So in python, that would be:
```
import torch
def supports_flash_attention(device_id):
    """Check if a GPU supports FlashAttention."""
    major, minor = torch.cuda.get_device_capability(device_id)
    
    # Check if the GPU architecture is Ampere (SM 8.x) or newer (SM 9.0)
    is_sm8x = major == 8 and minor >= 0
    is_sm90 = major == 9 and minor == 0

    return is_sm8x or is_sm90
```
with `device_id` being `0` for the first gpu, `1` for the second...

@ahassaine Right! You can certainly translate to get a python equivalent. One area of difficulty would be making sure the custom validation is up-to-date with the same checks in the FA2 library, but it will still give you a minimal set of compatible hardware. 

Getting this error with `meta-llama/Meta-Llama-3-8B-Instruct` model on RTX4060ti (which is of compute capability 8.9).
PyTorch 2.2.2, CUDA 12.1

@Archie3d If you think this is a mistake, then the best place to raise is on the [flash attention repo](https://github.com/Dao-AILab/flash-attention), where the checks are being performed

> [@ahassaine](https://github.com/ahassaine) If a models supports flash attention, it will have the private attribute `_supports_flash_attn_2` set to `True` e.g. [like here for bark](https://github.com/huggingface/transformers/blob/39c3c0a72af6fbda5614dde02ff236069bb79827/src/transformers/models/bark/modeling_bark.py#L487).

Does this mean that I have to check that out when working in my Google Colab notebook? I had set my runtime type to a T4 GPU... Nvm, I'll try `pip install flash-attn==2.5.7` before coming up further

Check Your GPU Architecture. Run nvidia-smi to check your GPU model. Ampere GPUs include:

NVIDIA A100, RTX 3090/4090, H100, etc.

Older GPUs like V100, T4, or GTX 10xx/20xx will not work with this version of FlashAttention

> Check Your GPU Architecture. Run nvidia-smi to check your GPU model. Ampere GPUs include:
> 
> NVIDIA A100, RTX 3090/4090, H100, etc.
> 
> Older GPUs like V100, T4, or GTX 10xx/20xx will not work with this version of FlashAttention

Thank you, now I get it. My laptop doesn't have an Nvidia GPU so I settled for a runtime in Colab. I was using the T4 there