make flash_attn ViT upgrade opt-in

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

Author: bradleyhd

vllm/attention/layer.py

@@ -65,7 +65,7 @@ def check_xformers_availability():
     return USE_XFORMERS_OPS
 
 
-def check_upstream_fa_availability(dtype: torch.dtype):
+def check_upstream_fa_availability(dtype: torch.dtype) -> bool:
     if (
         dtype in (torch.float16, torch.bfloat16)
         and current_platform.is_cuda()
@@ -80,26 +80,31 @@ def check_upstream_fa_availability(dtype: torch.dtype):
         return find_spec("flash_attn") is not None
     return False
 
+def is_fa_backend(backend: _Backend) -> bool:
+    return backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}
 
 def maybe_get_vit_flash_attn_backend(
-    attn_backend: _Backend, use_upstream_fa: bool
-) -> tuple[_Backend, Callable]:
-    if (
-        attn_backend != _Backend.FLASH_ATTN
-        and attn_backend != _Backend.ROCM_AITER_FA
-        and check_upstream_fa_availability(torch.get_default_dtype())
-    ):
-        attn_backend = _Backend.FLASH_ATTN
-        use_upstream_fa = True
+        attn_backend: _Backend,
+        try_switch_to_fa: bool = False,
+        try_use_upstream_fa: bool = False) -> tuple[_Backend, Callable]:
 
-    if current_platform.is_rocm() and attn_backend == _Backend.FLASH_ATTN:
-        use_upstream_fa = True
-
-    if attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
+    if try_switch_to_fa and not is_fa_backend(attn_backend):
+        attn_backend = _Backend.FLASH_ATTN
+    
+    if current_platform.is_rocm() and \
+        attn_backend == _Backend.FLASH_ATTN:
+        # Always try upstream on ROCM.
+        logger.info_once("maybe_get_vit_flash_attn_backend: forcing upstream FlashAttn on ROCM.")
+        try_use_upstream_fa = True
+    
+    if is_fa_backend(attn_backend):
         if attn_backend == _Backend.ROCM_AITER_FA:
             from aiter import flash_attn_varlen_func
         else:
-            if use_upstream_fa:
+            if try_use_upstream_fa:
+                assert check_upstream_fa_availability( \
+                    torch.get_default_dtype()), \
+                    "Upstream FlashAttn is not available."
                 from flash_attn import flash_attn_varlen_func
             else:
                 from vllm.vllm_flash_attn import flash_attn_varlen_func
@@ -108,7 +113,6 @@ def maybe_get_vit_flash_attn_backend(
 
     return attn_backend, flash_attn_varlen_func
 
-
 class Attention(nn.Module, AttentionLayerBase):
     """Attention layer.
 
@@ -428,11 +432,6 @@ def __init__(
         # Determine the attention backend
         backend = get_vit_attn_backend(head_size=head_size, dtype=dtype)
 
-        # Some auto-selected backends can be upgraded
-        # to upstream flash attention if available.
-        # If vllm native fa is selected, we use it directly.
-        use_upstream_fa = False
-
         if current_platform.is_xpu():
             # currently, only torch_sdpa is supported on xpu
             self.attn_backend = _Backend.TORCH_SDPA
@@ -450,30 +449,20 @@ def __init__(
                 else _Backend.TORCH_SDPA
             )
 
-        self.attn_backend, self._flash_attn_varlen_func = (
-            maybe_get_vit_flash_attn_backend(
+        self.attn_backend, self._flash_attn_varlen_func \
+            = maybe_get_vit_flash_attn_backend(
                 self.attn_backend,
-                use_upstream_fa,
+                try_switch_to_fa=False,
+                try_use_upstream_fa=False,
             )
-        )
 
         if self.attn_backend == _Backend.XFORMERS and not check_xformers_availability():
             self.attn_backend = _Backend.TORCH_SDPA
 
-        self.is_flash_attn_backend = self.attn_backend in {
-            _Backend.FLASH_ATTN,
-            _Backend.ROCM_AITER_FA,
-        }
-
-        # this condition is just to make sure that the
-        # use_upstream_fa in the log is correct
-        if current_platform.is_rocm() and self.attn_backend == _Backend.FLASH_ATTN:
-            use_upstream_fa = True
+        self.is_flash_attn_backend = is_fa_backend(self.attn_backend)
 
         logger.info_once(
-            f"MultiHeadAttention attn_backend: {self.attn_backend}, "
-            f"use_upstream_fa: {use_upstream_fa}"
-        )
+            f"MultiHeadAttention attn_backend: {self.attn_backend}")
 
     def forward(
         self,

vllm/model_executor/models/dots_ocr.py

@@ -295,7 +295,8 @@ def __init__(
         self.attn_backend, self.flash_attn_varlen_func = (
             maybe_get_vit_flash_attn_backend(
                 self.attn_backend,
-                self.use_upstream_fa,
+                try_switch_to_fa=True,
+                try_use_upstream_fa=self.use_upstream_fa,
             )
         )
         if self.attn_backend not in {

vllm/model_executor/models/ernie45_vl.py

@@ -203,7 +203,8 @@ def __init__(
         self.attn_backend, self.flash_attn_varlen_func = (
             maybe_get_vit_flash_attn_backend(
                 self.attn_backend,
-                self.use_upstream_fa,
+                try_switch_to_fa=True,
+                try_use_upstream_fa=self.use_upstream_fa,
             )
         )
 

vllm/model_executor/models/glm4_1v.py

@@ -293,7 +293,8 @@ def __init__(
         self.attn_backend, self.flash_attn_varlen_func = (
             maybe_get_vit_flash_attn_backend(
                 self.attn_backend,
-                self.use_upstream_fa,
+                try_switch_to_fa=True,
+                try_use_upstream_fa=self.use_upstream_fa,
             )
         )
 

vllm/model_executor/models/qwen2_5_vl.py

@@ -345,7 +345,8 @@ def __init__(
         self.attn_backend, self.flash_attn_varlen_func = (
             maybe_get_vit_flash_attn_backend(
                 self.attn_backend,
-                self.use_upstream_fa,
+                try_switch_to_fa=True,
+                try_use_upstream_fa=self.use_upstream_fa,
             )
         )
         self.is_flash_attn_backend = self.attn_backend in {

vllm/model_executor/models/qwen2_vl.py

@@ -361,7 +361,8 @@ def __init__(
         self.attn_backend, self.flash_attn_varlen_func = (
             maybe_get_vit_flash_attn_backend(
                 self.attn_backend,
-                self.use_upstream_fa,
+                try_switch_to_fa=True,
+                try_use_upstream_fa=self.use_upstream_fa,
             )
         )
 

vllm/model_executor/models/siglip2navit.py

@@ -255,7 +255,8 @@ def __init__(
         self.attn_backend, self.flash_attn_varlen_func = (
             maybe_get_vit_flash_attn_backend(
                 self.attn_backend,
-                self.use_upstream_fa,
+                try_switch_to_fa=True,
+                try_use_upstream_fa=self.use_upstream_fa,
             )
         )