Refactor flash attention implementation in transformers

[ArthurZucker]

What does this PR do?

EDIT: just refactor for now

Enables us to run transformers model with Ragged Tensors:

One of the goals is also to make it easy for people to re-define the ExtraKwargs typedict, to build on top of transformers

[ArthurZucker]

cc @fxmarty, @LysandreJik and @OlivierDehaene

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Thanks so much @fxmarty for going to the end of this!

[fxmarty]

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