Flash Attention v3
Author: hlky
Pull request: yes

# What does this PR do?

Replaces #33522 to avoid conflicts and allow those using it to continue while we get it updated for https://github.com/huggingface/transformers/pull/35235

Initial commit of this PR adds auxiliary code so we can discuss the core FAv3 integration.

cc @ArthurZucker 

- Integrate FAv3 into `_flash_attention_forward`/`flash_attention_forward` as [before](https://github.com/huggingface/transformers/blob/66eefe571611798b2f2b183e994f1e2b13744134/src/transformers/modeling_flash_attention_utils.py) or create new functions?
- Some models still have `FlashAttention2` classes, is refactoring all models to use the new style planned? Integrate FAv3 as [before](https://github.com/huggingface/transformers/blob/66eefe571611798b2f2b183e994f1e2b13744134/src/transformers/models/bart/modeling_bart.py#L290-L416) or do the refactor in this PR?

Also to check:
- Status of dropout, softcap etc
- Status of FP8
- Packaging

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.


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

cc @bn999 if you're interested about the progress 

@vasqu Yup, I'm following.  Good stuff.

Thanks for the info @vasqu 

Gentle ping @ArthurZucker

- Integrate FAv3 into _flash_attention_forward/flash_attention_forward as [before](https://github.com/huggingface/transformers/blob/66eefe571611798b2f2b183e994f1e2b13744134/src/transformers/modeling_flash_attention_utils.py) or create new functions?
- Some models still have FlashAttention2 classes, is refactoring all models to use the new style planned? Integrate FAv3 as [before](https://github.com/huggingface/transformers/blob/66eefe571611798b2f2b183e994f1e2b13744134/src/transformers/models/bart/modeling_bart.py#L290-L416) or do the refactor in this PR?

Hi @ArthurZucker @hlky @vasqu  @muellerzr  , thanks for the great efforts to integrate **Flash Attention 3** 😁. Do we have any plans to merge this PR?

Hey quick thing here @hlky, if you have FA3 installed but not FA2 (which I believe is a valid way it is used in other repos like TE) you end up failing the `is_flash_attn_2_available` check and get `_flash_attention_forward` is not even thought the check and enable FA3 function passes. Not sure if intentional or a bug, but if intentional a better guard could help tell people both FA2 and FA3 are required?

Hi @sam-h-bean. At the time this PR was started (more specifically, the original PR https://github.com/huggingface/transformers/pull/33522) pad functions were not available in FAv3, therefore FAv2 was required. As per https://github.com/huggingface/transformers/pull/36190#pullrequestreview-2618524634 this is likely no longer required and will be updated when this PR is finished. At the moment we are waiting for comments from a core-maintainer, @ArthurZucker, regarding https://github.com/huggingface/transformers/pull/36190#issuecomment-2677968118. 

@hlky sorry I probably forgot to merge 😓 don't worry we'll push trough and add support!

@ArthurZucker lol it's cool, the PR wasn't finished because I had been waiting for your response, didn't have time in Paris then I was fired so I closed it 🤷‍♂️