[Bugfix] Honor --mm_encoder_attn_backend when used

[bradleyhd]

Summary:
In #26104, some changes were made in layer.py that resulted in always trying to switch to FA backend for ViT, even when VLLM_ATTENTION_BACKEND is set.

This broke Meta's internal AMD pipelines as it is not desired nor expected behavior. With this change, the models that were changed in the offending PR can explicitly opt-in to this behavior.

Differential Revision: D84946967

[github-actions]

👋 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 fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

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 ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request makes the FlashAttention backend upgrade for Vision Transformer (ViT) models an opt-in behavior, addressing an issue where it was unconditionally attempted, causing problems on AMD platforms. The change is implemented by introducing a try_switch_to_fa flag in maybe_get_vit_flash_attn_backend and updating the call sites in various models.

The overall approach is sound and correctly addresses the reported issue. However, I've identified a critical bug in the new implementation that could lead to crashes on platforms not supporting FlashAttention, like XPU. I've also pointed out a high-severity maintainability issue regarding the modification of function parameters, which could make the code harder to reason about. Addressing these points will improve the robustness and clarity of the code.

[gemini-code-assist]

The current logic unconditionally switches the backend to FLASH_ATTN if try_switch_to_fa is true. This can cause a crash on platforms that do not support FlashAttention, such as XPU, because the subsequent import of vllm.vllm_flash_attn will fail. The switch should be guarded to only occur on supported platforms (CUDA and ROCm).

Suggested change

	if try_switch_to_fa and not is_fa_backend(attn_backend):
	attn_backend = _Backend.FLASH_ATTN
	if try_switch_to_fa and not is_fa_backend(attn_backend) and (
	current_platform.is_cuda() or current_platform.is_rocm()):
	attn_backend = _Backend.FLASH_ATTN

[gemini-code-assist]

Modifying an input parameter try_use_upstream_fa directly is confusing and can lead to unexpected side effects. It's better to use a local variable to track the state within the function. For example, you could introduce use_upstream_fa = try_use_upstream_fa at the beginning of the function and then modify and use use_upstream_fa.

[chatgpt-codex-connector]

💡 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".

[bradleyhd]

Updated to try and mimic #26104 as closely as possible to make this an equivalent change. Not sure the behavior in the original PR's is good / should be preserved, though.

[zhewenl]

This PR also fix existing AMD failures(example):

(EngineCore_DP0 pid=50574)   File "/usr/local/lib/python3.12/dist-packages/vllm/attention/layers/cross_attention.py", line 168, in __init__
--
  | (EngineCore_DP0 pid=50574)     super().__init__(
  | (EngineCore_DP0 pid=50574)   File "/usr/local/lib/python3.12/dist-packages/vllm/attention/layer.py", line 236, in __init__
  | (EngineCore_DP0 pid=50574)     self.impl = impl_cls(
  | (EngineCore_DP0 pid=50574)                 ^^^^^^^^^
  | (EngineCore_DP0 pid=50574)   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/triton_attn.py", line 248, in __init__
  | (EngineCore_DP0 pid=50574)     raise NotImplementedError(
  | (EngineCore_DP0 pid=50574) NotImplementedError: Encoder self-attention and encoder/decoder cross-attention are not implemented for TritonAttentionImpl
  | [rank0]:[W1017 04:54:49.728888986 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

cc @Alexei-V-Ivanov-AMD

[DarkLight1337]

cc @tjtanaa

[LucasWilkinson]

This logic is very confusing now; would be good to get more context here and try to refactor this a bit more aggressively,

cc @wwl2755 @tjtanaa

Seems like the original intention of using upstream FA is: #24347 , i.e. use it for models with head dim that is not supported by vllm-FA but is supported by upstream-FA; anything thats a multiple of 8 but not a multiple of 32

Can we just make all this logic; if on cuda and if head dim is not a multiple of 32 use upstream FA? otherwise just use get_attn_backend (like pre: #24347)?

[tjtanaa]

This logic is very confusing now; would be good to get more context here and try to refactor this a bit more aggressively,

cc @wwl2755 @tjtanaa

Seems like the original intention of using upstream FA is: #24347 , i.e. use it for models with head dim that is not supported by vllm-FA but is supported by upstream-FA; anything thats a multiple of 8 but not a multiple of 32

Can we just make all this logic; if on cuda and if head dim is not a multiple of 32 use upstream FA? otherwise just use get_attn_backend (like pre: #24347)?

To add on to @LucasWilkinson 's feedback

The logic that we should update should be

vllm/vllm/platforms/rocm.py

Lines 204 to 211 in c3a2c6a

	def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> "_Backend":
	from vllm.attention.backends.registry import _Backend
	if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
	return _Backend.ROCM_AITER_FA
	if on_gfx9():
	return _Backend.FLASH_ATTN
	return _Backend.TORCH_SDPA

Right now, as long as on AMD instinct, we assume that ck-flash-attention library is installed. If we want to enable torch.sdpa , we can start from modifying this part first. If ck-flash-attention is not installed, or the head_dim is not supported by the specified backend, we fall back to torch.sdpa

Another thing that I notice that VLLM_ATTENTION_BACKEND semantics should be meant for Text Model backbone.

The set of ATTENTION BACKEND supported by ViT are TORCH_SDPA, FLASH_ATTN and ROCM_AITER_FA only.

However, the ATTENTION_BACKEND for LLM Backbone are TRITON_ATTN, ROCM_ATTN, ROCM_AITER_FA or ROCM_AITER_UNIFIED_ATTN.

So, I would suggest reserving the VLLM_ATTENTION_BACKEND environment variable for LLM Attention Backend Selection.

Moreover, on MI300 series, flash attention/ aiter flash attention is recommended to be used for ViT as it is the fastest. When torch.sdpa is selection, it is extremely slow as it does for loop to compute the attention output in majority of the vision models.

[DarkLight1337]

Heads up that we have decoupled the two backends in #27061

[bradleyhd]

Heads up that we have decoupled the two backends in #27061

@DarkLight1337 thanks. is maybe_get_vit_flash_attn_backend needed now in light of this PR?

[DarkLight1337]

cc @ywang96 @Isotr0py

[ywang96]

@LucasWilkinson @tjtanaa @bradleyhd FYI on parallel to this PR, I've also made #27061 which decouples ViT attn backend from LM attn backend (which should probably be something we should've done from the get go).

[bradleyhd]

@tjtanaa curious, is upstream FA in this case expected to be FAv3? (looking at #24347)

[bradleyhd]

@ywang96 #27061 only works if we override to ROCM_AITER_FA, because it is exempt from the logic in maybe_get_vit_flash_attn_backend. If we set to TORCH_SDPA, it just gets overwritten with FA because we have a module named flash_attn installed

[bradleyhd]

alright folks, I've updated this to make use of the new --mm_encoder_attn_backend. When supplied, it won't auto-upgrade to FA. We need this asap to unblock as it allows us to specify torch_sdpa usage. We can and should circle back for a more comprehensive refactor here

[ywang96]

LGTM - I think we do need to think about how to deal with override in a better way (whether we should honor it truly with the risk of failure or handle the fallback automatically)