Add ALL_ATTENTION_FUNCTIONS compatibility for Pixtral model #37960
Conversation
|
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the |
Disallow using sdpa and output_attentions
|
Thanks @qubvel for the review 🙏! I made the changes you suggested, let me know if anything else is needed. |
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Great addition thanks! We don't really need the position ids, (should be kwargs imo!)
zucchini-nlp
left a comment
zucchini-nlp
left a comment
There was a problem hiding this comment.
Oh cool, I also had a PR for attention in VLMs in #37576 😄
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Thanks everyone for your reviews! Let me know if anything else is needed before merging 😊 |
|
Thanks for the contrib! |
…ace#37960) * Add ALL_ATTENTION_FUNCTIONS compatibility for Pixtral model * Fix invalid operand type * Allow image_sizes to be optional in forward pass to fit tests Disallow using sdpa and output_attentions * Disallow using sdpa with output_attentions * Delete useless comments, use eager attention from smolvlm, use pattern from mistral * add _supports_attention_backend * use kwargs instead of position_ids --------- Co-authored-by: aurelien.lac <[email protected]>
5 participants
What does this PR do?
This PR adds support for
ALL_ATTENTION_FUNCTIONSto the Pixtral model’s attention mechanism. I added and verified compatibility withsdpa,flash_attention_2, andflex_attention. Since Pixtral also serves as the vision tower inMistral 3.1, users can now set the entire model to useflash_attention_2.I tried to follow the implementation pattern of other models using this interface. For
flash_attention_2, I reusedposition_idsbecause the existing attention mask shape isn’t supported. Since Pixtral uses sequence packing and already generatesposition_ids, we leverageprepare_fa2_from_position_idsinstead of a mask.I tested these changes in training and inference: losses match very closely and we observe a 10–25 % throughput improvement depending on the setup.
Who can review?