[diffusion][Wan]: fix sparse attention backends being applied to cross-attention (#17596)

2026-02-16 21:57:58 +08:00
parent b86c6491fa
commit c1d1337afc
3 changed files with 28 additions and 6 deletions
@@ -301,7 +301,16 @@ class CausalWanTransformerBlock(nn.Module):

        # 2. Cross-attention
        # Only T2V for now
-        self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
+        cross_attn_backends = {
+            b for b in supported_attention_backends if not b.is_sparse
+        }
+        self.attn2 = WanT2VCrossAttention(
+            dim,
+            num_heads,
+            qk_norm=qk_norm,
+            eps=eps,
+            supported_attention_backends=cross_attn_backends,
+        )
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
        )