Refactor flash attention implementation in transformers
Author: ArthurZucker
Pull request: yes
Labels: name: run-slow

# What does this PR do?
EDIT: just refactor for now


 Enables us to run transformers model with Ragged Tensors: 

<img width="545" alt="image" src="https://github.com/huggingface/transformers/assets/48595927/12f60475-a4c5-477d-bb54-a5645aaa5360">


One of the goals is also to make it easy for people to re-define the `ExtraKwargs` typedict, to build on top of transformers

Comments:
cc @fxmarty, @LysandreJik and @OlivierDehaene 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31446). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Is this using FA2 or SDPA?

FA2


No more flash attention tests fail here compared to main (on H100).
```
================================================================== short test summary info ==================================================================
FAILED tests/models/bark/test_modeling_bark.py::BarkSemanticModelTest::test_flash_attn_2_from_config - ValueError: Unrecognized configuration class <class 'transformers.models.bark.configuration_bark.BarkSemanticConfig'> for this kind of AutoModel: AutoMo...
FAILED tests/models/bark/test_modeling_bark.py::BarkCoarseModelTest::test_flash_attn_2_from_config - ValueError: Unrecognized configuration class <class 'transformers.models.bark.configuration_bark.BarkCoarseConfig'> for this kind of AutoModel: AutoMode...
FAILED tests/models/dpr/test_modeling_dpr.py::DPRModelTest::test_sdpa_can_dispatch_on_flash - RuntimeError: No available kernel. Aborting execution.
FAILED tests/models/gemma/test_modeling_gemma.py::GemmaIntegrationTest::test_model_2b_flash_attn - OSError: You are trying to access a gated repo.
FAILED tests/models/gemma2/test_modeling_gemma2.py::Gemma2ModelTest::test_flash_attn_2_equivalence - AssertionError: assert False
FAILED tests/models/gemma2/test_modeling_gemma2.py::Gemma2ModelTest::test_sdpa_can_dispatch_on_flash - RuntimeError: No available kernel. Aborting execution.
FAILED tests/models/gpt_neox/test_modeling_gpt_neox.py::GPTNeoXModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/idefics2/test_modeling_idefics2.py::Idefics2ModelTest::test_flash_attn_2_inference_equivalence_right_padding - ValueError: You are attempting to perform batched generation with padding_side='right' this may lead to unexpected behaviour for Flash Attention version...
FAILED tests/models/idefics2/test_modeling_idefics2.py::Idefics2ForConditionalGenerationModelTest::test_flash_attn_2_inference_equivalence_right_padding - ValueError: You are attempting to perform batched generation with padding_side='right' this may lead to unexpected behaviour for Flash Attention version...
FAILED tests/models/jamba/test_modeling_jamba.py::JambaModelTest::test_flash_attn_2_generate_padding_right - AssertionError: ValueError not raised
FAILED tests/models/jamba/test_modeling_jamba.py::JambaModelTest::test_sdpa_can_dispatch_on_flash - RuntimeError: No available kernel. Aborting execution.
FAILED tests/models/m2m_100/test_modeling_m2m_100.py::M2M100ModelTest::test_flash_attn_2_from_config - ValueError: Unrecognized configuration class <class 'transformers.models.m2m_100.configuration_m2m_100.M2M100Config'> for this kind of AutoModel: AutoMo...
FAILED tests/models/m2m_100/test_modeling_m2m_100.py::M2M100ModelIntegrationTests::test_flash_attn_2_seq_to_seq_generation - RuntimeError: FlashAttention only support fp16 and bf16 data type
FAILED tests/models/mixtral/test_modeling_mixtral.py::MixtralModelTest::test_flash_attn_2_generate_padding_right - AssertionError: ValueError not raised
FAILED tests/models/olmo/test_modeling_olmo.py::OlmoModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/phi3/test_modeling_phi3.py::Phi3ModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/qwen2/test_modeling_qwen2.py::Qwen2ModelTest::test_flash_attn_2_generate_padding_right - AssertionError: ValueError not raised
FAILED tests/models/qwen2/test_modeling_qwen2.py::Qwen2ModelTest::test_flash_attn_2_inference_equivalence - AssertionError: assert False
FAILED tests/models/qwen2_moe/test_modeling_qwen2_moe.py::Qwen2MoeModelTest::test_flash_attn_2_generate_padding_right - AssertionError: ValueError not raised
FAILED tests/models/qwen2_moe/test_modeling_qwen2_moe.py::Qwen2MoeModelTest::test_flash_attn_2_inference_equivalence - AssertionError: assert False
FAILED tests/models/stablelm/test_modeling_stablelm.py::StableLmModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/starcoder2/test_modeling_starcoder2.py::Starcoder2ModelTest::test_flash_attn_2_generate_padding_right - AssertionError: ValueError not raised
FAILED tests/models/unispeech/test_modeling_unispeech.py::UniSpeechRobustModelTest::test_flash_attn_2_inference_equivalence - RuntimeError: mat1 and mat2 must have the same dtype, but got Float and BFloat16
FAILED tests/models/unispeech/test_modeling_unispeech.py::UniSpeechRobustModelTest::test_flash_attn_2_inference_equivalence_right_padding - RuntimeError: mat1 and mat2 must have the same dtype, but got Float and BFloat16
FAILED tests/models/unispeech/test_modeling_unispeech.py::UniSpeechRobustModelTest::test_sdpa_can_dispatch_on_flash - RuntimeError: mat1 and mat2 must have the same dtype, but got Float and Half
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperModelTest::test_flash_attn_2_from_config - IndexError: too many indices for tensor of dimension 2
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperModelTest::test_flash_attn_2_inference_equivalence_right_padding - IndexError: too many indices for tensor of dimension 2
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_from_config - IndexError: too many indices for tensor of dimension 2
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_generate_left_padding - IndexError: too many indices for tensor of dimension 2
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference_equivalence - AssertionError: assert False
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference_equivalence_right_padding - AssertionError: assert False
================================= 31 failed, 1064 passed, 1789 skipped, 62671 deselected, 166 warnings in 331.97s (0:05:31) =================================
```

Testing on MI250 for extra safety and good to merge.

edit: all good, can be merged