[Attention] FA2&FA3 support more head sizes, ViT support, make default backend #28763
Conversation
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
MatthewBonanni
requested review from
LucasWilkinson,
WoosukKwon,
mgoin,
tlrmchlsmth and
yewentao256
as code owners
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates FlashAttention to support head sizes required for Vision Transformers (40, 72, 80). This is achieved by updating the dependency to a fork of flash-attention, generalizing the head size check in the FlashAttention backend, and updating tests. The logic for selecting the ViT attention backend is also refactored for clarity. My review has identified two main points. First, a critical issue in cmake/external_projects/vllm_flash_attn.cmake where the dependency points to a personal fork, which must be reverted before merging. Second, a high-severity issue in tests/kernels/attention/test_flash_attn.py where a test case for soft_cap has been removed, potentially hiding a feature regression. The other changes look good.
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
MatthewBonanni
changed the title
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
MatthewBonanni
changed the title
|
Do you know if FA2 is supported too? do you mine testing this on Ampere? I think it should be ok |
|
@LucasWilkinson |
…oject#28763) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: inkcherry <mingzhi.liu@amd.com>
youkaichao
changed the title
…oject#28763) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
…oject#28763) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Bhagyashri <Bhagyashri.Gaikwad2@ibm.com>
…oject#28763) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: LuminolT <lumischen01@gmail.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: jiang1.li <jiang1.li@intel.com>
…oject#28763) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
4 participants
Purpose
This PR is paired with vllm-project/flash-attention#109 (merge that first after CI passes, then I'll update the git tag), which enables FA to support the head sizes required for vision transformers (40, 72, and 80). This PR also updates the selector to make FlashAttention the default backend over xFormers.
Test Plan
pytest tests/kernels/attention/test_flash_attn.py(updated with new head sizes)Test Result
Passes
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.