[diffusion] use fa3 by default when device supports it (#783)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import torch
+from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
@@ -13,18 +14,6 @@ from vllm_omni.diffusion.attention.backends.utils.fa import _pad_input, _unpad_i
 
 logger = init_logger(__name__)
 
-try:
-    # only tested with flash_attn v3
-    # from flash_attn_interface import flash_attn_func as flash_attn_3_func  # not available in flash-attn 2.8.1
-    from flash_attn import flash_attn_func, flash_attn_varlen_func  # can be FA2 or FA3
-except ImportError:
-    logger.warning(
-        "FlashAttentionBackend is not available. You may install flash-attn "
-        "by running `uv pip install flash-attn==2.8.1 --no-build-isolation`"
-        " or install pre-built flash-attn from https://github.com/Dao-AILab/flash-attention/releases"
-    )
-    raise ImportError
-
 
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True

vllm_omni/diffusion/attention/selector.py

@@ -5,12 +5,14 @@ import importlib
 import os
 from functools import cache
 
+import torch
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
 )
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.utils.platform_utils import detect_device_type, is_rocm
 
 logger = init_logger(__name__)
 
@@ -63,7 +65,22 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
         The selected attention backend class
     """
     # Check environment variable
-    backend_name: str | None = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+
+    backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND", None)
+
+    if detect_device_type() == "cuda" and not is_rocm():
+        compute_capability = torch.cuda.get_device_capability()
+        major, minor = compute_capability
+        if 80 <= major * 10 + minor < 100:
+            if backend_name is None:
+                backend_name = "FLASH_ATTN"
+        else:
+            if backend_name == "FLASH_ATTN":
+                logger.warning(
+                    """Flash Attention requires GPU with compute capability >= 8.0 or < 10.0. "
+                               "Falling back to TORCH_SDPA backend."""
+                )
+                backend_name = "TORCH_SDPA"
 
     if backend_name is not None:
         backend_name_upper = backend_name.upper()