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