[`Flash Attention`] Fix flash attention integration
Author: vasqu
Pull request: yes
Labels: name: for patch

The current flash attention implementation has several issues:
- `test_flash_attn_2_equivalence` was failing in all models (I think)
- Fa kwargs no longer had `max_length_q/k`
- Varlen has been handled incorrectly leading to errors by always opting to this path even if this was not the case (pos ids but no attention mask)
- Compiling was failing due to the usage of `globals()`
- `fa_kwargs` did not get the correct dtype as cat recasted to long apparently (fixed by #39843)
- `test_flash_attn_2_inference_equivalence` fails on a lot of models (e.g. Bart) due to some refactor on the test which is not compatible with those models it seems
- fullgraph compile did not work before
- Small fixes on #40029 according to Raushan's comments at https://github.com/huggingface/transformers/pull/40002#discussion_r2269873268

To consider:
- Supports top left mask should be deprecated but some models still depend on it so leaving it with a TODO
- ^ same for is_flash_attention available (we know lazy load on setting the attention implementation)

Discovered internally as well in #38188 for example

Comments:
Pretty important one so ccing a bit more than usual

cc @ArthurZucker @Cyrilvallez @gante @zucchini-nlp @manueldeprada 

Failing tests look unrelated to me, mistral tokenizer smthn smthn

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

The tests should be the common modeling tests of flash attention imo. 

Noticing that Bart still fails but decoder-only models seem fine so investigating. 
Edit: it's only the inference equivalence tests - eager matches fa tests pass
Edit 2: seems like some test rewrite on the fa tests broke the tests themeselves as the old version of the test passes

Fixed the tests on bart - mainly tested fa on bart and llama so that most others should be good as well.

Turns out that the fa test itself has also been broken for some time.

Some other things to note:
- fa2 with fullgraph compile will not work (due to lazy init) --> looking into it if there might be a workaround Edit: works now
- cache lazy init causes recompiles cc @Cyrilvallez @manueldeprada since you worked on the recent cache refactors I think

For reproducing: https://gist.github.com/vasqu/c264d14268fbc7f1eb725e09dfb87331 (without fullgraph)

Compiling prefill is 100% expected to cause recompile, it's not a regresion, always has been the case and always will be as the shape differs 🤗
Also you're using DynamicCache so there is not a real lazy init! (And from your small snippet, you cannot see recompiles, but I assume you tried on something else!)

> Also you're using DynamicCache so there is not a real lazy init! (And from your small snippet, you cannot see recompiles, but I assume you tried on something else!)

Might be also torch version dependent, working on 2.6 to work with fa atm. But glad to hear that's expected! Wasn't sure if I caused something there so rather safe than sorry here :)

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: mxfp4

#39806 changes the `make style` (see a996bd58a70f73fb07b20f069a790f491b15b7da for example), leaving it for now to main state

thanks for the fix @vasqu ❤️ 
For the record: [NVIDIA/kvpress](https://github.com/NVIDIA/kvpress) hooks the attention forward and removes tokens from the cache keys tensors. Probably better not to compute `cu_seq_lens_k` in `generate`, as it causes varlen to OOM if `cu_seq_lens_k[-1] != keys.shape[0]`

Yeah, it was removed previously when fixing an issue for KVPress but it is not a requirement for them. They don't call `generate()` and instead have a manual simplified generation loop. So they will fallback to `not_packed_sequence` path in FA2 utils :)

Everything works now! Only issue left is https://github.com/huggingface/transformers/pull/40002#issuecomment-3179165863 (which is a style issue where I intentionally didn't do anything to be up with main)

Can you take a last check @ArthurZucker, then we're good to go finally :partying_face: 

LGTM as well thanks for fixing will do a patch today / tomorrow!