torch.bfloat16 inference failed with RuntimeError: cutlassF: no kernel found to launch!
Author: VINUK0
Pull request: no

### System Info

***[Environment Information]***

*GPU :* `Nvidia T4 (15GB)`
*Python Version:* `3.10.12`
*Pytorch Version :* `2.1.1 (CUDA 12.1 | 11.8) [Both same results.]`
*Transformers Version :* `4.37.1`
*Accelerate Version :* `0.26.1`
*Optimum Version:* `1.16.2`
*Bitsandbytes:* `0.42.0`

***[Task Information]***

*Model Type:* `Text-Generation`
*Model Architecture:* `llama`
*Attention Implementation:* `sdpa "flash attention 1"`
*Model Load In Float16:* `True`
*Model Load In Bfloat16:* `True`
*Model Generate With Float16:* `True`
*Model Generate With Bfloat16:* `False (RuntimeError: cutlassF: no kernel found to launch!)`




### Who can help?

@ArthurZucker

### Information

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

### Tasks

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

### Reproduction

***[Replicate Requirements]***

*Load a pre trained model with `attention_Implementation="sdpa", torch_dtype=torch.bfloat16` to generate a sequence of tokens. It will show the error.*


### Expected behavior

*(RuntimeError: cutlassF: no kernel found to launch)*

Comments:
Hi @VINUK0, thanks for raising this issue! 

I believe this is a torch and hardware issue, rather than something in the transformers library. 

Other people have reported issues between torch and sdp, when loading in bf16 on T4. See [this issue](https://github.com/Lightning-AI/lit-gpt/issues/327) and [comment](https://github.com/Lightning-AI/lit-gpt/issues/327#issuecomment-1664674460), which others said helped resolve their issue. 

Other related issues:
* https://github.com/pytorch/pytorch/issues/102029#issuecomment-1560071148
* https://github.com/fishaudio/fish-speech/issues/7

I'd suggest trying to load your model in fp16 instead of bf16. 

cc @fxmarty for reference

Thank you, I can reproduce the issue on a T4:

```python
import torch

query = torch.rand(32, 8, 128, 64, dtype=torch.bfloat16, device="cuda")
key = torch.rand(32, 8, 128, 64, dtype=torch.bfloat16, device="cuda")
value = torch.rand(32, 8, 128, 64, dtype=torch.bfloat16, device="cuda")

res = torch.nn.functional.scaled_dot_product_attention(query,key,value)
```

Using `torch.backends.cuda.enable_mem_efficient_sdp(False)` fixes the issue (https://pytorch.org/docs/master/backends.html#torch.backends.cuda.enable_mem_efficient_sdp) and can be used as a temporary workaround. Alternatively, you can load the model with the `attn_implementation="eager"` argument to avoid using SDPA: `model = AutoModel.from_pretrained(..., attn_implementation="eager")`.

@drisspg Do you think this is a bug? Should the auto-dispatch of SDPA automatically fall back on the math path in case it is used with a device not supporting bf16?

This definitely was a bug, but should have been fixed with this PR: https://github.com/pytorch/pytorch/pull/116272
Do you know which version you are on for PyTorch when you repoed above?

This was with torch==2.1.2, and appears to be fixed in torch==2.2.0. @VINUK0 Please update to 2.2.0 or use `torch.backends.cuda.enable_mem_efficient_sdp(False)`. Thank you!