Update qwenimage/qwen_fa3_processor.py

qwenimage/qwen_fa3_processor.py

@@ -1,138 +1,116 @@
 """
-QwenDoubleStreamAttnProcessorFA3 - Optimized lazy loading + Blackwell support
+Paired with a good language model. Thanks!
 """
 
 import torch
-import torch.nn.functional as F
-import warnings
 from typing import Optional, Tuple
 from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
 
-# Global lazy cache
-_fa3_available: bool = False
-_flash_attn_func = None
-_fa3_unavailable_reason: str = ""
-
-
-def _ensure_fa3_loaded():
-    """Lazy-load FA3 kernel only the first time it's needed."""
-    global _fa3_available, _flash_attn_func, _fa3_unavailable_reason
-
-    if _fa3_available or _fa3_unavailable_reason:
-        return
-
-    # Blackwell (sm_100+) not supported yet by vllm-flash-attn3
-    if torch.cuda.is_available():
-        cap = torch.cuda.get_device_capability()
-        if cap[0] >= 10:
-            _fa3_unavailable_reason = (
-                "FlashAttention-3 is not yet supported on Blackwell (sm_100) GPUs. "
-                "Falling back to scaled-dot-product attention (SDPA)."
-            )
-            return
+import torch.nn.functional as F
 
-    try:
-        from kernels import get_kernel
+_flash_attn_func = None
+_kernels_err = None
+
+try:
+    from kernels import get_kernel
+    # 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
-        _fa3_available = True
-    except Exception as e:
-        _fa3_unavailable_reason = (
-            f"FlashAttention-3 via Hugging Face kernels unavailable.\n"
-            f"get_kernel('kernels-community/vllm-flash-attn3') failed with:\n{e}\n"
-            "Falling back to SDPA."
-        )
-
-
-# ---------------------------------------------------------------------------
-# Custom op (only registered once FA3 is confirmed available)
-# ---------------------------------------------------------------------------
-def _register_fa3_op():
-    if not _fa3_available:
-        return
+except Exception as e:
+    _flash_attn_func = None
+    _kernels_err = e
 
-    @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
-    def flash_attn_func(
-        q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
-    ) -> torch.Tensor:
-        output, _ = _flash_attn_func(q, k, v, causal=causal)
-        return output
 
-    @flash_attn_func.register_fake
-    def _fake(q, k, v, causal=False):
-        return torch.empty_like(q).contiguous()
+def _fa3_available() -> bool:
+    return _flash_attn_func is not None
 
 
-# Only register if FA3 is actually available
-if _fa3_available:
-    _register_fa3_op()
-
-
-# ---------------------------------------------------------------------------
-# SDPA fallback
-# ---------------------------------------------------------------------------
-def _sdpa_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    # (B, S, H, D_h) → SDPA format
+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)
+    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(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
+) -> torch.Tensor:
+    outputs, lse = _flash_attn_func(q, k, v, causal=causal)
+    return outputs
+
+@flash_attn_func.register_fake
+def _(q, k, v, **kwargs):
+    return torch.empty_like(q).contiguous()
 
 
-# ---------------------------------------------------------------------------
-# Main processor
-# ---------------------------------------------------------------------------
 class QwenDoubleStreamAttnProcessorFA3:
-    def __init__(self):
-        _ensure_fa3_loaded()
-        self._backend = "fa3" if _fa3_available else "sdpa"
+    """
+    FA3-based attention processor for Qwen double-stream architecture.
+    Computes joint attention over concatenated [text, image] streams using vLLM FlashAttention-3
+    accessed via Hugging Face `kernels`.
 
-        if not _fa3_available:
-            warnings.warn(f"QwenDoubleStreamAttnProcessorFA3: {_fa3_unavailable_reason}", stacklevel=2)
+    Notes / limitations:
+    - General attention masks are not supported here (FA3 path). `is_causal=False` and no arbitrary mask.
+    - Optional windowed attention / sink tokens / softcap can be plumbed through if you use those features.
+    - Expects an available `apply_rotary_emb_qwen` in scope (same as your non-FA3 processor).
+    """
 
-    def _attend(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-        if self._backend == "fa3":
-            return torch.ops.flash.flash_attn_func(q, k, v, False)  # custom op
-        return _sdpa_attention(q, k, v)
+    _attention_backend = "fa3"  # for parity with your other processors, not used internally
 
-    @torch.inference_mode()
+    def __init__(self):
+        if not _fa3_available():
+            print(f"[QwenDoubleStreamAttnProcessorFA3] FA3 unavailable, using SDPA fallback. Reason: {_kernels_err}")
+
+    @torch.no_grad()
     def __call__(
         self,
-        attn,
-        hidden_states: torch.FloatTensor,
-        encoder_hidden_states: torch.FloatTensor = None,
-        encoder_hidden_states_mask: torch.FloatTensor = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        attn,  # Attention module with to_q/to_k/to_v/add_*_proj, norms, to_out, to_add_out, and .heads
+        hidden_states: torch.FloatTensor,                 # (B, S_img, D_model)  image stream
+        encoder_hidden_states: torch.FloatTensor = None,  # (B, S_txt, D_model)  text stream
+        encoder_hidden_states_mask: torch.FloatTensor = None,  # unused in FA3 path
+        attention_mask: Optional[torch.FloatTensor] = None,    # unused in FA3 path
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # (img_freqs, txt_freqs)
     ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
         if encoder_hidden_states is None:
-            raise ValueError("QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states.")
-
-        if attention_mask is not None and self._backend == "fa3":
-            raise NotImplementedError("attention_mask not supported on FA3 path.")
+            raise ValueError("QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states (text stream).")
+        if attention_mask is not None:
+            # 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.")
 
         B, S_img, _ = hidden_states.shape
         S_txt = encoder_hidden_states.shape[1]
 
-        # QKV projections
-        img_q = attn.to_q(hidden_states)
+        # ---- QKV projections (image/sample stream) ----
+        img_q = attn.to_q(hidden_states)   # (B, S_img, D)
         img_k = attn.to_k(hidden_states)
         img_v = attn.to_v(hidden_states)
-        txt_q = attn.add_q_proj(encoder_hidden_states)
+
+        # ---- QKV projections (text/context stream) ----
+        txt_q = attn.add_q_proj(encoder_hidden_states)  # (B, S_txt, D)
         txt_k = attn.add_k_proj(encoder_hidden_states)
         txt_v = attn.add_v_proj(encoder_hidden_states)
 
-        # Reshape
+        # ---- Reshape to (B, S, H, D_h) ----
         H = attn.heads
         img_q = img_q.unflatten(-1, (H, -1))
         img_k = img_k.unflatten(-1, (H, -1))
         img_v = img_v.unflatten(-1, (H, -1))
+
         txt_q = txt_q.unflatten(-1, (H, -1))
         txt_k = txt_k.unflatten(-1, (H, -1))
         txt_v = txt_v.unflatten(-1, (H, -1))
 
-        # Norms
+        # ---- Q/K normalization (per your module contract) ----
         if getattr(attn, "norm_q", None) is not None:
             img_q = attn.norm_q(img_q)
         if getattr(attn, "norm_k", None) is not None:
@@ -142,32 +120,38 @@ class QwenDoubleStreamAttnProcessorFA3:
         if getattr(attn, "norm_added_k", None) is not None:
             txt_k = attn.norm_added_k(txt_k)
 
-        # RoPE
+        # ---- RoPE (Qwen variant) ----
         if image_rotary_emb is not None:
             img_freqs, txt_freqs = image_rotary_emb
+            # expects tensors shaped (B, S, H, D_h)
             img_q = apply_rotary_emb_qwen(img_q, img_freqs, use_real=False)
             img_k = apply_rotary_emb_qwen(img_k, img_freqs, use_real=False)
             txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs, use_real=False)
             txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs, use_real=False)
 
-        # Joint attention
+        # ---- Joint attention over [text, image] along sequence axis ----
+        # Shapes: (B, S_total, H, D_h)
         q = torch.cat([txt_q, img_q], dim=1)
         k = torch.cat([txt_k, img_k], dim=1)
         v = torch.cat([txt_v, img_v], dim=1)
 
-        out = self._attend(q, k, v)
+        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)
+        # ---- Back to (B, S, D_model) ----
         out = out.flatten(2, 3).to(q.dtype)
 
-        # Split & output projections
-        txt_out = out[:, :S_txt, :]
-        img_out = out[:, S_txt:, :]
+        # Split back to text / image segments
+        txt_attn_out = out[:, :S_txt, :]
+        img_attn_out = out[:, S_txt:, :]
 
-        img_out = attn.to_out[0](img_out)
+        # ---- Output projections ----
+        img_attn_out = attn.to_out[0](img_attn_out)
         if len(attn.to_out) > 1:
-            img_out = attn.to_out[1](img_out)
+            img_attn_out = attn.to_out[1](img_attn_out)  # dropout if present
 
-        txt_out = attn.to_add_out(txt_out)
+        txt_attn_out = attn.to_add_out(txt_attn_out)
 
-        return img_out, txt_out
+        return img_attn_out, txt_attn_out