[Diffusion][Bugfix] Fix the flash_attn backends selection logic (vllm-project#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>

Author: mxuax

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -10,8 +10,26 @@
     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 @@ def forward_cuda(
         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 @@ def forward_cuda(
             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,57 +3,80 @@
 # 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
 
+# 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)
+
+# Try flash_attn_interface first (from flash-attention source build)
 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
+    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
 
-    HAS_FLASH_ATTN_HOPPER = True
-except ImportError:
-    HAS_FLASH_ATTN_HOPPER = False
+# 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:
     import aiter  # noqa: F401
     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 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 @@
     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.
+
+    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"
 
-    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,
+    # 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 @@
 import torch
 
 from .ring_globals import (
-    HAS_NPU,
     HAS_SAGE_ATTENTION,
     HAS_SPARSE_SAGE_ATTENTION,
 )
@@ -28,9 +27,6 @@
 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 fn(q, k, v, causal=False, softmax_scale=None, *args, **kwargs):
 
         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 @@
 
 from vllm_omni.platforms import current_omni_platform
 
+# Flash Attention function detection with fallback chain
+flash_attn_func = None
+flash_attn_varlen_func = None
+
 if current_omni_platform.is_rocm():
-    from vllm._aiter_ops import is_aiter_found_and_supported
+    # ROCm: try Aiter first
+    try:
+        from vllm._aiter_ops import is_aiter_found_and_supported
 
-    # 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 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):