[ROCm] Unify ViT FlashAttention backend selection

docker/Dockerfile.rocm_base

@@ -8,7 +8,7 @@ ARG TORCHAUDIO_VERSION=2.11.0+rocm10.0.0
 ARG TRITON_BRANCH="669b31ac" # release/internal/3.8.x as of 09/29/2026
 ARG TRITON_VERSION=3.8.0+git669b31ac
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
-ARG FA_BRANCH="0e60e394"
+ARG FA_BRANCH="4a948e9c"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
 ARG AITER_BRANCH="v0.1.27"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
@@ -194,7 +194,7 @@ RUN git clone ${FA_REPO}
 RUN cd flash-attention \
     && git checkout ${FA_BRANCH} \
     && git submodule update --init \
-    && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist
+    && GPU_ARCHS="${PYTORCH_ROCM_ARCH}" python3 setup.py bdist_wheel --dist-dir=dist
 RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
 
 

vllm/platforms/rocm.py

@@ -499,40 +499,6 @@ def use_rocm_custom_paged_attention(
         )
 
 
-@cache
-def flash_attn_triton_available() -> bool:
-    if not on_gfx1x():
-        return False
-    try:
-        from importlib.util import find_spec
-
-        # Locate the Triton-AMD kernels. Older ROCm/flash-attention (pre
-        # 2026-03) shipped them as the flash_attn.flash_attn_triton_amd
-        # subpackage. The main_perf migration commit 3f94643 moved them
-        # into aiter at aiter.ops.triton._triton_kernels.flash_attn_triton_amd,
-        # so accept either location.
-        def _has_spec(name: str) -> bool:
-            try:
-                return find_spec(name) is not None
-            except (ImportError, ValueError):
-                return False
-
-        if not (
-            _has_spec("flash_attn.flash_attn_triton_amd")
-            or _has_spec("aiter.ops.triton._triton_kernels.flash_attn_triton_amd")
-        ):
-            return False
-        if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE") != "TRUE":
-            logger.info_once(
-                "Set FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE to enable "
-                "Flash Attention Triton backend on RDNA."
-            )
-            return False
-        return True
-    except ImportError:
-        return False
-
-
 def _get_backend_priorities(
     use_mla: bool,
     use_sparse: bool,
@@ -970,24 +936,13 @@ def get_vit_attn_backend(
             return AttentionBackendEnum.ROCM_AITER_FA
 
         if (
-            on_cdna()
+            (on_cdna() or on_gfx1x())
             and find_spec("flash_attn") is not None
             and (dtype == torch.float16 or dtype == torch.bfloat16)
         ):
             logger.info_once("Using Flash Attention backend for ViT model.")
             return AttentionBackendEnum.FLASH_ATTN
 
-        # RDNA3/RDNA4 (gfx11xx/gfx12xx): Use Flash Attention Triton backend
-        if (
-            on_gfx1x()
-            and flash_attn_triton_available()
-            and (dtype == torch.float16 or dtype == torch.bfloat16)
-        ):
-            logger.info_once(
-                "Using Flash Attention (Triton backend) for ViT model on RDNA."
-            )
-            return AttentionBackendEnum.FLASH_ATTN
-
         logger.info_once("Using Torch SDPA backend for ViT model.")
         return AttentionBackendEnum.TORCH_SDPA