[Diffusion][Bugfix] Fix the flash_attn backends selection logic (#983)

Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -10,8 +10,26 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
 )
 
+# Import flash attention functions with fallback chain from utils/fa.py
+# FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
+from vllm_omni.diffusion.attention.backends.utils.fa import (
+    HAS_FLASH_ATTN,
+    _pad_input,
+    _unpad_input,
+    _upad_input,
+    flash_attn_func,
+    flash_attn_varlen_func,
+)
+
 logger = init_logger(__name__)
 
+if not HAS_FLASH_ATTN:
+    raise ImportError(
+        "FlashAttentionBackend requires Flash Attention. "
+        "Please install one of: fa3-fwd, flash-attention, or flash-attn. "
+        "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+    )
+
 
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
@@ -56,14 +74,6 @@ class FlashAttentionImpl(AttentionImpl):
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
         """CUDA/ROCm flash attention implementation."""
-        from vllm_omni.diffusion.attention.backends.utils.fa import (
-            _pad_input,
-            _unpad_input,
-            _upad_input,
-            flash_attn_func,
-            flash_attn_varlen_func,
-        )
-
         query_length = query.size(1)
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         #  Contains at least one padding token in the sequence
@@ -92,13 +102,16 @@ class FlashAttentionImpl(AttentionImpl):
             out = _pad_input(out_unpad, indices_q, query.size(0), query_length)
 
         else:
-            out: torch.Tensor = flash_attn_func(
+            out = flash_attn_func(
                 query,
                 key,
                 value,
                 causal=self.causal,
                 softmax_scale=self.softmax_scale,
             )
+            # FA3 may return (out, lse) tuple, FA2 returns just out
+            if isinstance(out, tuple):
+                out = out[0]
         return out
 
     def forward_npu(

vllm_omni/diffusion/attention/backends/ring/ring_globals.py

@@ -3,30 +3,53 @@
 # Copyright (c) 2024, Jiarui Fang.
 # Adapted from https://github.com/feifeibear/long-context-attention
 
-
-# test if flash_attn is available
+# test if flash_attn (FA2) is available
 try:
     import flash_attn  # noqa: F401
-    from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward  # noqa: F401
+    from flash_attn.flash_attn_interface import _flash_attn_forward  # noqa: F401
 
     HAS_FLASH_ATTN = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_FLASH_ATTN = False
 
-try:
-    from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward  # noqa: F401
-    from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper  # noqa: F401
-    from flash_attn_interface import flash_attn_func as flash3_attn_func  # noqa: F401
+# FA3 detection: try multiple sources (forward only, no backward needed for inference)
+# Source 1: flash_attn_interface (from flash-attention source build)
+# Source 2: fa3_fwd_interface (from fa3-fwd PyPI package, supports Ampere/Ada/Hopper)
+# Note: FA3 high-level API may or may not return softmax_lse depending on version.
+#       For Ring Attention which requires LSE, we fall back to low-level API if needed.
+HAS_FA3 = False
+fa3_fwd_func = None  # Low-level forward function (_flash_attn_forward)
+fa3_attn_func = None  # High-level attention function (flash_attn_func)
 
-    HAS_FLASH_ATTN_HOPPER = True
-except ImportError:
-    HAS_FLASH_ATTN_HOPPER = False
+# Try flash_attn_interface first (from flash-attention source build)
+try:
+    from flash_attn_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401
+    from flash_attn_interface import flash_attn_func as fa3_attn_func  # noqa: F401
+
+    HAS_FA3 = True
+except (ImportError, ModuleNotFoundError):
+    pass
+
+# Fallback: try fa3_fwd_interface (PyPI package, supports Ampere/Ada/Hopper)
+if not HAS_FA3:
+    try:
+        from fa3_fwd_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401
+        from fa3_fwd_interface import flash_attn_func as fa3_attn_func  # noqa: F401
+
+        HAS_FA3 = True
+    except (ImportError, ModuleNotFoundError):
+        pass
+
+# Legacy aliases for backward compatibility
+HAS_FLASH_ATTN_HOPPER = HAS_FA3
+flash_attn_forward_hopper = fa3_fwd_func
+flash3_attn_func = fa3_attn_func
 
 try:
     from flashinfer.prefill import single_prefill_with_kv_cache  # noqa: F401
 
     HAS_FLASHINFER = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_FLASHINFER = False
 
 try:
@@ -34,26 +57,26 @@ try:
     from aiter import flash_attn_func as flash_attn_func_aiter  # noqa: F401
 
     HAS_AITER = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_AITER = False
 
 try:
     import sageattention  # noqa: F401
 
     HAS_SAGE_ATTENTION = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_SAGE_ATTENTION = False
 
 try:
     import spas_sage_attn  # noqa: F401
 
     HAS_SPARSE_SAGE_ATTENTION = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_SPARSE_SAGE_ATTENTION = False
 
 try:
     import torch_npu  # noqa: F401
 
     HAS_NPU = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_NPU = False

vllm_omni/diffusion/attention/backends/ring/ring_kernels.py

@@ -6,7 +6,13 @@ import math
 
 import torch
 
-from .ring_globals import HAS_AITER, HAS_FLASH_ATTN, HAS_FLASH_ATTN_HOPPER, HAS_FLASHINFER, HAS_NPU
+from .ring_globals import (
+    HAS_AITER,
+    HAS_FA3,
+    HAS_FLASH_ATTN,
+    HAS_FLASHINFER,
+    fa3_fwd_func,
+)
 
 _scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention
 _scaled_dot_product_efficient_attention = torch.ops.aten._scaled_dot_product_efficient_attention
@@ -26,23 +32,11 @@ if HAS_FLASH_ATTN:
     import flash_attn
     from flash_attn.flash_attn_interface import _flash_attn_forward
 
-if HAS_FLASH_ATTN_HOPPER:
-    from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward
-    from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper
-    from flash_attn_interface import flash_attn_func as flash3_attn_func
-else:
-    flash_attn_forward_hopper = None
-    flash_attn_func_hopper_backward = None
-    flash3_attn_func = None
-
 if HAS_FLASHINFER:
     from flashinfer.prefill import single_prefill_with_kv_cache
 
     _LOG2_E = math.log2(math.e)
 
-if HAS_NPU:
-    import torch_npu
-
 
 def pytorch_attn_forward(
     q: torch.Tensor,
@@ -146,52 +140,35 @@ def flash_attn_forward(
     return block_out, block_lse
 
 
-def flash_attn3_func_forward(
-    q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax
-):
-    assert HAS_FLASH_ATTN_HOPPER
-    # current signature of flash_attn_forward_hopper:
-    # (q, k, v, softmax_scale, causal, window_size, descale_q=None, descale_k=None, descale_v=None, gqa_parallel=False)
+def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax):
+    """FA3 forward pass for inference.
 
-    out, softmax_lse, *unused = flash_attn_forward_hopper(
-        q=q,
-        k=k,
-        v=v,
-        k_new=None,
-        v_new=None,
-        qv=None,
-        out=None,
-        cu_seqlens_q=None,
-        cu_seqlens_k=None,
-        cu_seqlens_k_new=None,
-        seqused_q=None,
-        seqused_k=None,
-        max_seqlen_q=None,
-        max_seqlen_k=None,
-        page_table=None,
-        kv_batch_idx=None,
-        leftpad_k=None,
-        rotary_cos=None,
-        rotary_sin=None,
-        seqlens_rotary=None,
-        q_descale=None,
-        k_descale=None,
-        v_descale=None,
+    FA3 supports Ampere, Ada, and Hopper GPUs. Dropout is ignored since FA3 is inference-only.
+    Uses low-level API (_flash_attn_forward) which always returns softmax_lse,
+    required for Ring Attention's correct accumulation.
+    """
+    assert HAS_FA3, "FA3 is not available"
+    assert fa3_fwd_func is not None, "FA3 low-level API (fa3_fwd_func) not available"
+
+    # Low-level API always returns (out, softmax_lse, S_dmask, rng_state)
+    out, softmax_lse, *_ = fa3_fwd_func(
+        q,
+        k,
+        v,
         softmax_scale=softmax_scale,
-        causal=False,
-        window_size=(-1, -1),
-        attention_chunk=0,
-        softcap=0.0,
-        rotary_interleaved=True,
-        scheduler_metadata=None,
-        num_splits=0,
-        pack_gqa=None,
-        sm_margin=0,
+        causal=causal,
+        window_size_left=window_size[0] if window_size else -1,
+        window_size_right=window_size[1] if window_size else -1,
+        softcap=softcap if softcap else 0.0,
     )
 
     return out, softmax_lse
 
 
+# Legacy alias for backward compatibility
+flash_attn3_func_forward = fa3_forward
+
+
 def flash_attn_forward_aiter(
     q,
     k,
@@ -264,20 +241,3 @@ def flashinfer_attn_forward(
         raise ValueError(f"Invalid input shape: {q.shape}")
     lse = lse / _LOG2_E
     return out, lse
-
-
-def npu_attn_forward(q, k, v, softmax_scale=None, layout="BSND"):
-    assert HAS_NPU, "torch_npu is not available"
-    softmax_scale = q.shape[-1] ** (-0.5)
-    block_out, block_lse = torch_npu.npu_fused_infer_attention_score(
-        q,
-        k,
-        v,
-        num_heads=q.shape[-2],
-        input_layout=layout,
-        scale=softmax_scale,
-        softmax_lse_flag=True,
-        pre_tokens=65535,
-        next_tokens=65535,
-    )
-    return block_out, block_lse.squeeze(dim=-1)

vllm_omni/diffusion/attention/backends/ring/ring_selector.py

@@ -10,7 +10,6 @@ from functools import partial
 import torch
 
 from .ring_globals import (
-    HAS_NPU,
     HAS_SAGE_ATTENTION,
     HAS_SPARSE_SAGE_ATTENTION,
 )
@@ -28,9 +27,6 @@ if HAS_SAGE_ATTENTION:
 if HAS_SPARSE_SAGE_ATTENTION:
     from spas_sage_attn.autotune import SparseAttentionMeansim
 
-if HAS_NPU:
-    from torch_npu import npu_fused_infer_attention_score
-
 
 class AttnType(Enum):
     AITER = "aiter"
@@ -44,7 +40,6 @@ class AttnType(Enum):
     SAGE_FP8 = "sage_fp8"
     SAGE_FP8_SM90 = "sage_fp8_sm90"
     SPARSE_SAGE = "sparse_sage"
-    NPU = "npu"
 
     @classmethod
     def from_string(cls, s: str):
@@ -157,11 +152,6 @@ def select_flash_attn_impl(
 
         return fn
 
-    elif impl_type == AttnType.NPU:
-        if not HAS_NPU:
-            raise ImportError("torch_npu is not available!")
-        return npu_fused_infer_attention_score
-
     elif attn_processor is not None:
         return attn_processor
 

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -17,17 +17,53 @@ import torch.nn.functional as F
 
 from vllm_omni.platforms import current_omni_platform
 
-if current_omni_platform.is_rocm():
-    from vllm._aiter_ops import is_aiter_found_and_supported
+# Flash Attention function detection with fallback chain
+flash_attn_func = None
+flash_attn_varlen_func = None
 
-    # Choose to enable this by default on ROCm
-    # Whenever possible as it is the fastest backend
-    if is_aiter_found_and_supported():
-        from aiter import flash_attn_func, flash_attn_varlen_func  # noqa: F401
-    else:
-        raise ImportError("Aiter is not found and supported on currentROCm device.")
+if current_omni_platform.is_rocm():
+    # ROCm: try Aiter first
+    try:
+        from vllm._aiter_ops import is_aiter_found_and_supported
+
+        if is_aiter_found_and_supported():
+            from aiter import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+    except (ImportError, ModuleNotFoundError):
+        pass
 else:
-    from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+    # CUDA: try FA3 -> FA2 fallback chain
+    # Try FA3 from fa3-fwd PyPI package
+    try:
+        from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+    except (ImportError, ModuleNotFoundError):
+        pass
+
+    # Fallback: Try FA3 from flash-attention source build
+    if flash_attn_func is None:
+        try:
+            from flash_attn_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+    # Fallback: Try FA2 from flash-attn package (try multiple import paths)
+    if flash_attn_func is None:
+        try:
+            from flash_attn import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+    if flash_attn_func is None:
+        try:
+            from flash_attn.flash_attn_interface import (  # noqa: F401
+                flash_attn_func,
+                flash_attn_varlen_func,
+            )
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+# If no FA backend available, SDPA backend will be selected at the platform level
+# flash_attn_func and flash_attn_varlen_func will be None
+HAS_FLASH_ATTN = flash_attn_func is not None
 
 
 def _index_first_axis(tensor, indices):

vllm_omni/diffusion/attention/parallel/ring.py

@@ -10,7 +10,7 @@ import torch
 from vllm.logger import init_logger
 
 # import torch.distributed as dist # Not used directly here, but good practice if needed
-from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FLASH_ATTN
+from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FA3, HAS_FLASH_ATTN
 from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
 from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
@@ -116,11 +116,14 @@ class RingParallelAttention:
             except Exception:
                 backend_pref = None
 
-        # Fallback for FP32 or if Flash Attention is not available
-        if query.dtype == torch.float32 or not HAS_FLASH_ATTN:
-            if not HAS_FLASH_ATTN and backend_pref != "sdpa":
+        # Determine attention type with fallback chain: FA3 -> FA2 -> SDPA
+        # FP32 is not supported by Flash Attention, force SDPA
+        if query.dtype == torch.float32:
+            backend_pref = "sdpa"
+        elif not HAS_FA3 and not HAS_FLASH_ATTN:
+            if backend_pref != "sdpa":
                 logger = init_logger(__name__)
-                logger.warning_once("Flash Attention is not available! Force enabling SDPA.")
+                logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.")
             backend_pref = "sdpa"
 
         # Extract joint tensors
@@ -150,6 +153,9 @@ class RingParallelAttention:
 
         from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
 
+        # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper)
+        attn_type = AttnType.FA3 if HAS_FA3 else AttnType.FA
+
         return ring_flash_attn_func(
             query,
             key,
@@ -161,9 +167,8 @@ class RingParallelAttention:
             softcap=0.0,
             alibi_slopes=None,
             deterministic=False,
-            # return_attn_probs=False, # Removed as it might not be supported in signature
             group=self._sp_group.ring_group,
-            attn_type=AttnType.FA,
+            attn_type=attn_type,
             joint_tensor_key=joint_key,
             joint_tensor_value=joint_value,
             joint_strategy=joint_strategy,

vllm_omni/diffusion/envs.py

@@ -66,14 +66,32 @@ class PackagesEnvChecker:
             if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
                 return False
 
+            # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn)
+            # Try FA3 from fa3-fwd PyPI package
+            try:
+                import fa3_fwd_interface  # noqa: F401
+
+                return True
+            except (ImportError, ModuleNotFoundError):
+                pass
+
+            # Try FA3 from flash-attention source build
+            try:
+                import flash_attn_interface  # noqa: F401
+
+                return True
+            except (ImportError, ModuleNotFoundError):
+                pass
+
+            # Try FA2 from flash-attn package
             from flash_attn import __version__
 
             if __version__ < "2.6.0":
                 raise ImportError("install flash_attn >= 2.6.0")
             return True
-        except ImportError:
+        except (ImportError, ModuleNotFoundError):
             if not packages_info.get("has_aiter", False):
-                logger.warning('Flash Attention library "flash_attn" not found, using pytorch attention implementation')
+                logger.warning("No Flash Attention backend found, using pytorch SDPA implementation")
             return False
 
     def get_packages_info(self) -> dict:

vllm_omni/platforms/cuda/platform.py

@@ -39,22 +39,34 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         selected_backend: str | None,
         head_size: int,
     ) -> str:
+        from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+
         # Check compute capability for Flash Attention support
         # Flash Attention requires compute capability >= 8.0 and < 10.0
         compute_capability = cls.get_device_capability()
-        flash_attn_supported = False
+        compute_supported = False
         if compute_capability is not None:
             major, minor = compute_capability
             capability = major * 10 + minor
-            flash_attn_supported = 80 <= capability < 100
+            compute_supported = 80 <= capability < 100
+
+        # Check if FA packages are available
+        packages_info = PACKAGES_CHECKER.get_packages_info()
+        packages_available = packages_info.get("has_flash_attn", False)
+
+        # Both compute capability and packages must be available for FA
+        flash_attn_supported = compute_supported and packages_available
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
-                logger.warning(
-                    "Flash Attention requires GPU with compute capability >= 8.0 "
-                    "and < 10.0. Falling back to TORCH_SDPA backend."
-                )
+                if not compute_supported:
+                    logger.warning(
+                        "Flash Attention requires GPU with compute capability >= 8.0 "
+                        "and < 10.0. Falling back to TORCH_SDPA backend."
+                    )
+                elif not packages_available:
+                    logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.")
                 logger.info("Defaulting to diffusion attention backend SDPA")
                 return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             backend = DiffusionAttentionBackendEnum[backend_upper]