From c3065115128e5369ca51be5a9c40f68640d55a6e Mon Sep 17 00:00:00 2001
From: Lancer <maruixiang6688@gmail.com>
Date: Mon, 4 May 2026 22:22:06 +0800
Subject: [PATCH] [Bugfix] Fix CUBLAS_STATUS_EXECUTION_FAILED when native Flash
 Attention is available (Wan2.2) (#3327)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
---
 tests/diffusion/attention/test_flash_attn.py  | 33 +++++++++++++++++++
 .../attention/backends/flash_attn.py          | 11 +++++++
 .../diffusion/attention/backends/utils/fa.py  | 13 ++++----
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py
index 5ec1517d9..a6fcf8fcb 100644
--- a/tests/diffusion/attention/test_flash_attn.py
+++ b/tests/diffusion/attention/test_flash_attn.py
@@ -15,9 +15,12 @@ import torch
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
 from vllm_omni.platforms import current_omni_platform
 
 is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
+HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
+flash_attn_func = fa.flash_attn_func  # noqa: N813
 
 
 def create_attention_mask(batch_size: int, seq_len: int, valid_len: int, device: torch.device) -> torch.Tensor:
@@ -264,6 +267,36 @@ def test_fa_vs_sdpa():
     print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
 
 
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
+def test_flash_attn_func_preferred_over_varlen():
+    """Test flash_attn_func availability and basic forward call."""
+    if not HAS_FLASH_ATTN:
+        pytest.skip("No Flash Attention available")
+
+    if flash_attn_func is None:
+        pytest.skip("flash_attn_func not available, will use varlen fallback")
+
+    device = torch.device(current_omni_platform.device_type)
+    dtype = torch.bfloat16
+
+    num_heads, head_dim = 8, 64
+    fa_impl = FlashAttentionImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+
+    torch.manual_seed(42)
+    q = torch.randn(1, 32, num_heads, head_dim, device=device, dtype=dtype)
+    k = q.clone()
+    v = q.clone()
+
+    attn_metadata = AttentionMetadata(attn_mask=None)
+    output = fa_impl.forward(q, k, v, attn_metadata)
+
+    assert output.shape == q.shape
+    assert not torch.isnan(output).any()
+    print("✓ flash_attn_func forward works correctly!")
+
+
 if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)
diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py
index 5975313f3..f987fbbbe 100644
--- a/vllm_omni/diffusion/attention/backends/flash_attn.py
+++ b/vllm_omni/diffusion/attention/backends/flash_attn.py
@@ -140,6 +140,7 @@ class FlashAttentionImpl(AttentionImpl):
         """CUDA/ROCm/MUSA flash attention implementation."""
         from vllm_omni.diffusion.attention.backends.utils.fa import (
             HAS_FLASH_ATTN,
+            flash_attn_func,
         )
 
         if not HAS_FLASH_ATTN:
@@ -159,6 +160,16 @@ class FlashAttentionImpl(AttentionImpl):
                 attention_mask,
             )
 
+        if flash_attn_func is not None:
+            out = flash_attn_func(
+                query,
+                key,
+                value,
+                causal=self.causal,
+                softmax_scale=self.softmax_scale,
+            )
+            return self._unwrap_flash_output(out)
+
         return self._forward_varlen_dense(
             query,
             key,
diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py
index f6d4700f0..e30d45ff2 100644
--- a/vllm_omni/diffusion/attention/backends/utils/fa.py
+++ b/vllm_omni/diffusion/attention/backends/utils/fa.py
@@ -51,27 +51,28 @@ else:
     # CUDA: try FA3 -> FA2 fallback chain
     # Try FA3 from fa3-fwd PyPI package
     try:
-        from fa3_fwd_interface import flash_attn_varlen_func  # noqa: F401
+        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_varlen_func is None:
+    if flash_attn_func is None:
         try:
-            from flash_attn_interface import flash_attn_varlen_func  # noqa: F401
+            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_varlen_func is None:
+    if flash_attn_func is None:
         try:
-            from flash_attn import flash_attn_varlen_func  # noqa: F401
+            from flash_attn import flash_attn_func, flash_attn_varlen_func  # noqa: F401
         except (ImportError, ModuleNotFoundError):
             pass
 
-    if flash_attn_varlen_func is None:
+    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):