[Bugfix] Fix ViT with FlashAttention on ROCm (#30703)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

[Bugfix] Fix ViT with FlashAttention on ROCm (#30703)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
51e5b3e3 · Matthew Bonanni · GitHub · ec154c36 · 51e5b3e3
Unverified Commit 51e5b3e3 authored Dec 15, 2025 by Matthew Bonanni Committed by GitHub Dec 15, 2025
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

vllm/attention/layer.py vllm/attention/layer.py +4 -1

No files found.
--- a/vllm/attention/layer.py
+++ b/vllm/attention/layer.py
@@ -464,7 +464,10 @@ class MultiHeadAttention(nn.Module):
        }

        self.fa_version = None
-        if self.attn_backend == AttentionBackendEnum.FLASH_ATTN:
+        if (
+            self.attn_backend == AttentionBackendEnum.FLASH_ATTN
+            and current_platform.is_cuda()
+        ):
            self.fa_version = get_flash_attn_version()
            assert self._flash_attn_varlen_func is not None
            self._flash_attn_varlen_func = functools.partial(