Merge 5f7e7ed289 into ec902046f6

wan/modules/attention.py

@@ -79,6 +79,7 @@ def flash_attention(
         k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
         v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
 
+    try:
         q = q.to(v.dtype)
         k = k.to(v.dtype)
     
@@ -126,6 +127,53 @@ def flash_attention(
                 window_size=window_size,
                 deterministic=deterministic).unflatten(0, (b, lq))
     
+    except RuntimeError as e:
+        if "FlashAttention only supports Ampere GPUs or newer" in str(e):
+            #for cards like 2080ti that aren't Ampere structure
+            from torch import nn
+            import torch.nn.functional as F
+        
+            q = q.to(half(k).dtype)
+
+            # 转置维度，保证形状为 [B, N, L, C]
+            q = q.view(b, lq, q.size(1), q.size(2)).transpose(1, 2)
+            k = k.view(b, lk, k.size(1), k.size(2)).transpose(1, 2)
+            v = v.view(b, lk, v.size(1), v.size(2)).transpose(1, 2)
+
+            # 计算注意力
+            # 注意：确保 Q、K、V 的形状为 [B, N, L, C]
+            # 设置默认缩放因子
+            if softmax_scale is None:
+                softmax_scale = 1.0 / q.size(-1) ** 0.5
+
+            # 如果 q_scale 存在，则应用缩放
+            if q_scale is not None:
+                q = q * q_scale
+
+            # 创建掩码
+            if causal:
+                attn_mask = torch.triu(torch.full((q.size(2), k.size(2)), -torch.inf), diagonal=1).to(q.device)
+            else:
+                attn_mask = None
+
+            # 计算注意力
+            # 使用 scaled_dot_product_attention
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attn_mask,
+                dropout_p=dropout_p,
+                is_causal=causal,
+            )
+
+            # 转换回原形状 [B, L, N, C]
+            x = x.transpose(1, 2).contiguous()
+
+            # 对输出应用 Dropout
+            dropout = nn.Dropout(dropout_p)
+            x = dropout(x)            
+        else:
+            raise
+            
     # output
     return x.type(out_dtype)