Flex attention + refactor

[ArthurZucker]

Opening this to add support for all models following #34282

Lets bring support for flex attention to more models! 🤗

• Gemma2

It would be great to add the support for more architectures such as

• Qwen2
• Llama
• Gemma
• QwenVl
• Mistral
• Clip

... and many more

For anyone who wants to contribute just open a PR and link it to this issue, and ping me for a review!! 🤗

[dame-cell]

@ArthurZucker I'll try to open a pr for some architecture maybe llama and gemma

[farrosalferro]

I will do the Clip and Llama

[OmarManzoor]

@ArthurZucker I would like to try out Mistral

[mayankagarwals]

Picking this up for gpt2 and moshi!

[dame-cell]

Picking this up for gpt2 and moshi

@mayankagarwals hey there since moshi actually copies some code form Gemma is it ok if I handle it

[mayankagarwals]

I've already started some efforts but happy to work on it together if you are interested too? @dame-cell . Shot you a dm on kaggle.

[dame-cell]

I've already started some efforts but happy to work on it together if you are interested too? @dame-cell . Shot you a dm on kaggle.

@mayankagarwals
Thanks for reaching out
You are right i think it's best if you do it :)