[feature] add optional attention mask to native flash attention and update backend handling in SkyReelsV2AttnProcessor

Author: tolgacangoz

src/diffusers/models/attention_dispatch.py

@@ -873,6 +873,7 @@ def _native_flash_attention(
     query: torch.Tensor,
     key: torch.Tensor,
     value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
     dropout_p: float = 0.0,
     is_causal: bool = False,
     scale: Optional[float] = None,
@@ -884,7 +885,7 @@ def _native_flash_attention(
             query=query,
             key=key,
             value=value,
-            attn_mask=None,  # not supported
+            attn_mask=attn_mask,
             dropout_p=dropout_p,
             is_causal=is_causal,
             scale=scale,

src/diffusers/models/transformers/transformer_skyreels_v2.py

@@ -123,6 +123,12 @@ def apply_rotary_emb(
             query = apply_rotary_emb(query, *rotary_emb)
             key = apply_rotary_emb(key, *rotary_emb)
 
+        if self._attention_backend == "_native_flash-flash_varlen":
+            if not self.is_cross_attention:
+                self._attention_backend = "_native_flash"
+            else:
+                self._attention_backend = "flash_varlen"
+
         # I2V task
         hidden_states_img = None
         if encoder_hidden_states_img is not None:
@@ -153,7 +159,6 @@ def apply_rotary_emb(
             is_causal=False,
             backend=self._attention_backend,
         )
-
         hidden_states = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.type_as(query)