[V1] Support any head size for FlexAttention backend #20467
Conversation
Signed-off-by: DarkLight1337 <[email protected]>
DarkLight1337
requested review from
WoosukKwon,
hmellor and
ywang96
as code owners
DarkLight1337
added
the
ready
DarkLight1337
requested review from
alexm-redhat,
comaniac,
houseroad,
mgoin,
njhill,
robertgshaw2-redhat,
simon-mo,
tlrmchlsmth and
youkaichao
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
mergify
bot
added
documentation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @DarkLight1337, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request enhances the flexibility of the vLLM library by enabling the FlexAttention backend to support any head size. This change allows previously unsupported models to run in V1 mode and updates error messages to guide users towards using FlexAttention when necessary. The PR also includes updates to the documentation and test configurations to reflect these changes.
Highlights
- FlexAttention Backend: Updated the FlexAttention backend to allow any head size, removing the previous restriction to specific sizes.
- Error Messaging: Modified the unsupported head size message for other backends to guide users to set
VLLM_ATTENTION_BACKEND=FLEX_ATTENTIONinstead ofVLLM_USE_V1=0. - Model Support: Re-enabled affected models in multimodal and quantization tests by leveraging the FlexAttention backend for models with head sizes like 80.
- Documentation: Removed the note in
supported_models.mdabouthead_size = 80not being supported.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request effectively implements support for arbitrary head sizes in the FlexAttention backend, which is a significant improvement for model compatibility. The changes are well-structured, with the introduction of a validate_head_size static method across various attention backends. This centralizes validation logic and provides clear guidance to users about using the FlexAttention backend for unsupported head sizes. Documentation, examples, and tests have been updated appropriately to reflect these changes, ensuring correctness and consistency. The refactoring improves code maintainability and clarity without introducing any new issues.
| Read audio from video pre-processing is currently supported on V0 (but not V1), because overlapping modalities is not yet supported in V1. | ||
| `--mm-processor-kwargs '{"use_audio_in_video": true}'`. | ||
| For Qwen2.5-Omni, reading audio from video pre-processing (`--mm-processor-kwargs '{"use_audio_in_video": true}'`) | ||
| is currently supported on V0 (but not V1), because overlapping modalities is not yet supported in V1. |
There was a problem hiding this comment.
Fix the out-of-date documentation
| REQUIRES_V0_MODELS = [ | ||
| # V1 Test: no way to fall back for head_dim = 80 | ||
| # https://github.com/vllm-project/vllm/issues/14524 | ||
| "qwen_vl", |
There was a problem hiding this comment.
This model actually has head size 128
|
FlexAttention doesn't seem to work in the CI environment: |
|
cc @youkaichao @drisspg I think it's related to pytorch/pytorch#133254 Any ideas on how to fix/work around this in the meantime? |
Signed-off-by: DarkLight1337 <[email protected]>
| # As of this writing, head_size=80 is only supported by FlexAttention in V1 | ||
| if model_name.endswith("-2b"): | ||
| os.environ["VLLM_ATTENTION_BACKEND"] = "FLEX_ATTENTION" |
There was a problem hiding this comment.
Can we add a check in get_attn_backend_cls to automatically fall back to FlexAttn?
Lines 191 to 194 in ffe00ef
The check should be very simple to fit FlashInfer and FlashAttn:
if head_size % 32:
return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"There was a problem hiding this comment.
This check is not complete though. I am not sure whether we can import the attention backend classes to call validate_head_size beforehand.
There was a problem hiding this comment.
From a quick look at the files, it's very likely that importing these files will initialize CUDA. Perhaps we can move validate_head_size to a separate file.
There was a problem hiding this comment.
@WoosukKwon how should we solve this?
There was a problem hiding this comment.
Well, let's see if importing the metadata directly can pass the tests as is first...
There was a problem hiding this comment.
Looks like we can safely import the attention backends
There was a problem hiding this comment.
Please don't forget to update the phi2 test, see: https://github.com/vllm-project/vllm/blob/v0.9.2rc1/tests/lora/test_phi.py#L54
The default FlexAttention config can run out of shared memory on L4, perhaps we need to reduce the M/N block_size when running FlexAttn like the fp32 PR: vllm/vllm/v1/attention/backends/flex_attention.py Lines 466 to 484 in dd6a37a The default FlexAttn config is here: https://github.com/pytorch/pytorch/blob/47c8aa80905f8cc9ea5488b5a3a209bee62d8409/torch/_inductor/kernel/flex_attention.py#L595-L618 |
Signed-off-by: DarkLight1337 <[email protected]>
Signed-off-by: DarkLight1337 <[email protected]>
|
There are still some failing models: https://buildkite.com/vllm/ci/builds/23252/steps/canvas?jid=0197da62-4d56-453a-a8df-506dd4df2085 |
|
It appears that none of the V1 attention backends currently support blocksparse attention which is used in Phi-3-Small. I'll continue to skip that test until FlexAttention supports it then. |
Signed-off-by: DarkLight1337 <[email protected]>
Signed-off-by: DarkLight1337 <[email protected]>
|
Also cc @LucasWilkinson |
|
All tests pass apart from the ones that are already failing on main |
| def validate_head_size(cls, head_size: int) -> None: | ||
| return # FlexAttention supports any head size |
There was a problem hiding this comment.
It would be nice if @drisspg could confirm
There was a problem hiding this comment.
Just seeing this now, yeah we support any head size. You will start to run out of shared memory around 256 + w/ some tricks you can get 512/576 working but for all intents and purposes and accept for some vision models this is true
There was a problem hiding this comment.
How can we calculate the amount of shared memory required? The current way to reduce BLOCK_M and BLOCK_N in L463 is very rough.
Sync to v0.9.2 + remove libsodium + [fix cachetokeziner](neuralmagic/nm-vllm-ent@1423512) git log: ``` commit 7b94527 (HEAD -> sync-v0.9.2, nm-fork/sync-v0.9.2) Merge: 1423512 d07be8a Author: Selbi Nuryyeva <[email protected]> Date: Fri Jul 11 07:03:51 2025 -0400 Merge remote-tracking branch 'nm-fork/main' into sync-v0.9.2 commit 1423512 Author: Isotr0py <[email protected]> Date: Mon Jun 30 18:16:16 2025 +0800 disable using CacheTokenizer for transformers >= 4.53.0 fixes vllm-project#20224 addendum to vllm-project#20244 commit d07be8a (nm-fork/main, nm-fork/HEAD) Merge: bbccdbe 02152ad Author: Daniele <[email protected]> Date: Wed Jul 9 15:18:56 2025 +0200 Dockerfile*.ubi: remove libsodium (opendatahub-io#245) It's not needed anymore https://issues.redhat.com/browse/INFERENG-848 commit 7dd12da Merge: bbccdbe a5dd03c Author: Selbi Nuryyeva <[email protected]> Date: Tue Jul 8 10:08:37 2025 -0400 Merge branch 'v0.9.2-upstream' into sync-v0.9.2 commit a5dd03c (tag: v0.9.2rc2, tag: v0.9.2, upstream/releases/v0.9.2, v0.9.2-upstream, upstream-v0.9.2) Author: simon-mo <[email protected]> Date: Sun Jul 6 14:02:36 2025 -0700 Revert "[V0 deprecation] Remove V0 CPU/XPU/TPU backends (vllm-project#20412)" This reverts commit e202dd2. commit c18b3b8 Author: Cyrus Leung <[email protected]> Date: Mon Jul 7 05:01:48 2025 +0800 [Bugfix] Add `use_cross_encoder` flag to use correct activation in `ClassifierPooler` (vllm-project#20527) Signed-off-by: DarkLight1337 <[email protected]> commit 9528e3a Author: Woosuk Kwon <[email protected]> Date: Sun Jul 6 12:44:52 2025 -0700 [BugFix][Spec Decode] Fix spec token ids in model runner (vllm-project#20530) Signed-off-by: Woosuk Kwon <[email protected]> commit 9fb52e5 Author: Cyrus Leung <[email protected]> Date: Mon Jul 7 00:54:36 2025 +0800 [V1] Support any head size for FlexAttention backend (vllm-project#20467) Signed-off-by: DarkLight1337 <[email protected]> ``` Test: CUDA: https://github.com/neuralmagic/nm-cicd/actions/runs/16218517666 ROCM: https://github.com/neuralmagic/nm-cicd/actions/runs/16218578391
) Signed-off-by: DarkLight1337 <[email protected]>
) Signed-off-by: DarkLight1337 <[email protected]>
) Signed-off-by: DarkLight1337 <[email protected]> Signed-off-by: Jinzhen Lin <[email protected]>
7 participants
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
Currently, models with
head_size in [80, 112, 120]are supported in V0 but not in V1. The reason is that the FlexAttention backend is marked as only supportinghead_size in [32, 64, 80, 96, 112, 120, 128, 192, 256]when it should actually support all head sizes.This PR updates the FlexAttention backend to allow any head size, and updates the unsupported head size message of other backends to ask users to set
VLLM_ATTENTION_BACKEND=FLEX_ATTENTIONinstead ofVLLM_USE_V1=0.Based on the discussion in: https://vllm-dev.slack.com/archives/C087RA55P0D/p1750313174942849
cc @mgoin @WoosukKwon
FIX #14524
Test Plan
Re-enable affected models in multimodal and quantization tests.
Test Result
(Optional) Documentation Update
Removed note about
head_size = 80not being supported.