[mgoin]

What does this PR do?

Fixes large memory consumption and performance issues with unfused attention in Pixtral. The usage comes from the vLLM implementation in vllm-project/vllm#9597

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

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.

[Rocketknight1]

cc @zucchini-nlp - should we keep the eager implementation around as well rather than making SDPA the only option?

[zucchini-nlp]

@mgoin cool, thanks for adding thsi! Yes, lets add SDPA as another option and run the tests (including slow) to make sure everything works. We'll also need to set the _supports_sdpa flag to True so it can dispatch to it by default

[ArthurZucker]

Thanks for this! Pixtral now goes through the shared attention interface, which already provides SDPA (and other backends), so hand-wired SDPA calls in the model file aren't needed anymore. Closing, thank you!