fix: fa3 broken on Blackwell+

qwenimage/qwen_fa3_processor.py

@@ -6,22 +6,40 @@ import torch
 from typing import Optional, Tuple
 from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
 
+import torch.nn.functional as F
+
+_flash_attn_func = None
+_kernels_err = None
+
 try:
     from kernels import get_kernel
-    _k = get_kernel("kernels-community/vllm-flash-attn3")
-    _flash_attn_func = _k.flash_attn_func
+    # Blackwell (sm_120+) is not yet supported by the vllm-flash-attn3 kernel binary
+    _cap = torch.cuda.get_device_capability() if torch.cuda.is_available() else (0, 0)
+    if _cap >= (12, 0):
+        _kernels_err = RuntimeError(
+            f"GPU compute capability sm_{_cap[0]}{_cap[1]} (Blackwell+) is not supported "
+            "by the kernels-community/vllm-flash-attn3 binary; using SDPA fallback."
+        )
+    else:
+        _k = get_kernel("kernels-community/vllm-flash-attn3")
+        _flash_attn_func = _k.flash_attn_func
 except Exception as e:
     _flash_attn_func = None
     _kernels_err = e
 
 
-def _ensure_fa3_available():
-    if _flash_attn_func is None:
-        raise ImportError(
-            "FlashAttention-3 via Hugging Face `kernels` is required. "
-            "Tried `get_kernel('kernels-community/vllm-flash-attn3')` and failed with:\n"
-            f"{_kernels_err}"
-        )
+def _fa3_available() -> bool:
+    return _flash_attn_func is not None
+
+
+def _sdpa_fallback(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    # q/k/v: (B, S, H, D_h) → SDPA expects (B, H, S, D_h)
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v = v.transpose(1, 2)
+    out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+    return out.transpose(1, 2)  # back to (B, S, H, D_h)
+
 
 @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
 def flash_attn_func(
@@ -32,11 +50,7 @@ def flash_attn_func(
 
 @flash_attn_func.register_fake
 def _(q, k, v, **kwargs):
-    # two outputs:
-    # 1. output: (batch, seq_len, num_heads, head_dim)
-    # 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32
-    meta_q = torch.empty_like(q).contiguous()
-    return meta_q #, q.new_empty((q.size(0), q.size(2), q.size(1)), dtype=torch.float32)
+    return torch.empty_like(q).contiguous()
 
 
 class QwenDoubleStreamAttnProcessorFA3:
@@ -54,7 +68,8 @@ class QwenDoubleStreamAttnProcessorFA3:
     _attention_backend = "fa3"  # for parity with your other processors, not used internally
 
     def __init__(self):
-        _ensure_fa3_available()
+        if not _fa3_available():
+            print(f"[QwenDoubleStreamAttnProcessorFA3] FA3 unavailable, using SDPA fallback. Reason: {_kernels_err}")
 
     @torch.no_grad()
     def __call__(
@@ -72,8 +87,6 @@ class QwenDoubleStreamAttnProcessorFA3:
             # FA3 kernel path here does not consume arbitrary masks; fail fast to avoid silent correctness issues.
             raise NotImplementedError("attention_mask is not supported in this FA3 implementation.")
 
-        _ensure_fa3_available()
-
         B, S_img, _ = hidden_states.shape
         S_txt = encoder_hidden_states.shape[1]
 
@@ -122,8 +135,10 @@ class QwenDoubleStreamAttnProcessorFA3:
         k = torch.cat([txt_k, img_k], dim=1)
         v = torch.cat([txt_v, img_v], dim=1)
 
-        # FlashAttention-3 path expects (B, S, H, D_h) and returns (out, softmax_lse)
-        out = flash_attn_func(q, k, v, causal=False)  # out: (B, S_total, H, D_h)
+        if _fa3_available():
+            out = flash_attn_func(q, k, v, causal=False)  # out: (B, S_total, H, D_h)
+        else:
+            out = _sdpa_fallback(q, k, v)  # out: (B, S_total, H, D_h)
 
         # ---- Back to (B, S, D_model) ----
         out = out.flatten(2, 3).to(q.dtype)