SDPA fallback for flash_attn_func in the action module

wan/modules/action_module.py

@@ -1,6 +1,5 @@
 from typing import Any, List, Tuple, Optional, Union, Dict
 from einops import rearrange
-from flash_attn import flash_attn_func
 import torch
 import torch.nn as nn
 from .posemb_layers import apply_rotary_emb, get_nd_rotary_pos_embed
@@ -9,11 +8,21 @@ from torch.nn.attention.flex_attention import flex_attention
 
 try:
     import flash_attn_interface
-    FLASH_ATTN_3_AVAILABLE = True
-except:
-    from flash_attn import flash_attn_func
+    FLASH_ATTN_3_AVAILABLE = torch.cuda.is_available() and (
+        "h100" in torch.cuda.get_device_name(0).lower()
+        or "hopper" in torch.cuda.get_device_name(0).lower())
+except ModuleNotFoundError:
     FLASH_ATTN_3_AVAILABLE = False
 
+try:
+    from flash_attn import flash_attn_func
+except ModuleNotFoundError:
+    # No FlashAttention wheel on this platform.  Every call site below runs
+    # without a causal flag or a sliding window (causality is enforced by the
+    # KV-cache layout, not the kernel), so PyTorch SDPA is a numerically
+    # equivalent substitute.
+    from .attention import sdpa_flash_attn_func as flash_attn_func
+
 
 DISABLE_COMPILE = False  # get os env
 flex_attention = torch.compile(

wan/modules/attention.py

@@ -183,3 +183,46 @@ def attention(
 
         out = out.transpose(1, 2).contiguous()
         return out
+
+
+def sdpa_flash_attn_func(
+    q,
+    k,
+    v,
+    dropout_p=0.,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    **kwargs,
+):
+    """Drop-in replacement for ``flash_attn.flash_attn_func`` built on SDPA.
+
+    Same ``[B, L, H, D]`` in/out layout.  Only the argument combinations the
+    action module actually uses are supported (no varlen, no sliding window).
+    Causal masking follows FlashAttention's bottom-right alignment, which
+    differs from ``F.scaled_dot_product_attention(is_causal=True)`` whenever the
+    query and key lengths differ.
+    """
+    if window_size != (-1, -1):
+        raise NotImplementedError(
+            "sliding-window attention requires the real flash_attn kernel")
+
+    q_t, k_t, v_t = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+    attn_mask, is_causal = None, False
+    if causal:
+        lq, lk = q_t.shape[-2], k_t.shape[-2]
+        if lq == lk:
+            is_causal = True
+        else:
+            rows = torch.arange(lq, device=q.device).unsqueeze(1) + (lk - lq)
+            cols = torch.arange(lk, device=q.device).unsqueeze(0)
+            attn_mask = cols <= rows
+
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q_t, k_t, v_t,
+        attn_mask=attn_mask,
+        dropout_p=dropout_p,
+        is_causal=is_causal,
+        scale=softmax_scale,
+    )
+    return out.transpose(1, 2)