[Bug]: Clarification regarding bug inside vllm-flash-attn vision module

[vrdn-23]

Your current environment

N/A

🐛 Describe the bug

Hi,
I am looking for a clarification regarding the warning message that pops up when trying to load a Molmo Vision model in vLLM. It seems that the warning message was first introduced in this PR and carried along with many refactors but I couldn't find any discussion or reference to what the actual bug was or any issue tracking it in the vllm-flash-attn repo. It seems anyway that the vllm-flash-attn repo is (mostly) up-to-date with the latest flash-attn, so I'm confused why this bug still exists in the fork but not in the main repo.

"Current `vllm-flash-attn` has a bug inside vision module, so we use xformers backend instead. You can run `pip install flash-attn` to use flash-attention backend."

Can we add an issue tracking this and add it to the warning message or perhaps remove the check if it is perhaps no longer relevant?
cc @ywang96 @mrsalehi @DarkLight1337 @mgoin

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[vrdn-23]

Potential duplicate of #10791

[DarkLight1337]

I remember that this was first introduced by Qwen2-VL PR. (#7905) cc @Isotr0py

[Isotr0py]

Hmmm, we have tried to add FA support on ViT before. However, vllm-flash-attn is only available for certain head_size while most of ViTs have unsupported head_size (For example, Qwen2-VL has a head_size=80 at ViT, which is not supported), and the performance improvement we observed is also very minor (only about 0.2%).

Therefore, we finally reverted it to use current fallback for better maintenance.

[vrdn-23]

Thanks for the quick response @Isotr0py @DarkLight1337
Just so that I can understand the issue a bit more clearly:

However, vllm-flash-attn is only available for certain head_size while most of ViTs have unsupported head_size (For example, Qwen2-VL has a head_size=80 at ViT, which is not supported), and the performance improvement we observed is also very minor (only about 0.2%).

Does this mean that the original flash-attn library is capable of supporting those head sizes (the warning message indicates that installing it from the original package it will help us run it on flash attention) ? If yes, is there any reason why vllm-flash-attn cannot just use those changes from the upstream library? Or is the error message incorrect in saying that the original flash attention package will make the model work on FA?

[Isotr0py]

Does this mean that the original flash-attn library is capable of supporting those head sizes (the warning message indicates that installing it from the original package it will help us run it on flash attention) ?

Yes, the original FA package support all head_size up to 256.

If yes, is there any reason why vllm-flash-attn cannot just use those changes from the upstream library?

I'm not really sure why vllm-flash-attn missed the all head_size support from upstream as well. Perhaps @LucasWilkinson may have some insights?

[LucasWilkinson]

The head size support should be the same in vllm-flash-attn and upstream (we were playing around with reducing this to save on compile time / wheel size but ultimately just left it the same, not sure if that ever even made it to upstream and the timeline doesn't quite match)

Im not aware of the context of this warning, maybe @fyabc can provide more context? would be nice to get rid of it if possible

[vrdn-23]

Thanks for the prompt responses everyone! While we wait for @fyabc to respond, would it make sense for me to raise a PR that effectively removes the warning and the xformers fallback and see if the models does start-up correctly with vllm-flash-attn? My primary concern with not trying this out in the first place, was that the warning was indicating that the "bug" was in regards to a issue in model quality, i.e., the model outputs were different when run on vllm-flash-attn v/s otherwise (which I'm not sure if we have tests set up for) but it seems from the discussion above, the initial concern was more about an implementation/design issue rather than a subtle quality issue.

Let me know if this sounds acceptable @DarkLight1337 @LucasWilkinson @Isotr0py

[LucasWilkinson]

would it make sense for me to raise a PR that effectively removes the warning and the xformers fallback and see if the models does start-up correctly with vllm-flash-attn?

I think this would be wise, alot has changed on the vllm and vllm-flash-attn since this was introduced. Im not super well versed in vision models, maybe @mgoin or @ywang96 knows of some good evals (ideally lm-eval) we could run to check correctness?

[fyabc]

The head size support should be the same in vllm-flash-attn and upstream (we were playing around with reducing this to save on compile time / wheel size but ultimately just left it the same, not sure if that ever even made it to upstream and the timeline doesn't quite match)

Im not aware of the context of this warning, maybe @fyabc can provide more context? would be nice to get rid of it if possible

Hi @LucasWilkinson @vrdn-23 @DarkLight1337, when I implemented Qwen2-VL in #7905, vllm_flash_attn.flash_attn_interface.flash_attn_varlen_func would output NaN values (and I couldn't find the bug), so I used flash_attn.flash_attn_varlen_func instead.
I am not sure if vllm-flash-attn have fixed this bug, and I will run the test on the latest version.

[mertunsall]

any updates on this?

[fyabc]

Hi @LucasWilkinson @vrdn-23, I have tried to replace current implementation with vllm.vllm_flash_attn.flash_attn_interface.flash_attn_varlen_func, and it will throw an error:

It seems that current version of vllm-flash-attn cannot support Qwen2.5-VL model series, how can I fix it?

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!