Support latest FA3 API changes

Author: ChrisLiu6

src/diffusers/models/attention_dispatch.py

@@ -621,24 +621,32 @@ def _wrapped_flash_attn_3(
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     # Hardcoded for now because pytorch does not support tuple/int type hints
     window_size = (-1, -1)
-    out, lse, *_ = flash_attn_3_func(
-        q=q,
-        k=k,
-        v=v,
-        softmax_scale=softmax_scale,
-        causal=causal,
-        qv=qv,
-        q_descale=q_descale,
-        k_descale=k_descale,
-        v_descale=v_descale,
-        window_size=window_size,
-        attention_chunk=attention_chunk,
-        softcap=softcap,
-        num_splits=num_splits,
-        pack_gqa=pack_gqa,
-        deterministic=deterministic,
-        sm_margin=sm_margin,
-    )
+
+    kwargs = {
+        "q": q,
+        "k": k,
+        "v": v,
+        "softmax_scale": softmax_scale,
+        "causal": causal,
+        "qv": qv,
+        "q_descale": q_descale,
+        "k_descale": k_descale,
+        "v_descale": v_descale,
+        "window_size": window_size,
+        "attention_chunk": attention_chunk,
+        "softcap": softcap,
+        "num_splits": num_splits,
+        "pack_gqa": pack_gqa,
+        "deterministic": deterministic,
+        "sm_margin": sm_margin,
+    }
+
+    # For backward compatibility with early flash-attn-3 APIs.
+    if "return_attn_probs" in inspect.signature(flash_attn_3_func).parameters:
+        kwargs["return_attn_probs"] = True
+
+    out, lse, *_ = flash_attn_3_func(**kwargs)
+
     lse = lse.permute(0, 2, 1)
     return out, lse
 
@@ -1504,17 +1512,29 @@ def _flash_varlen_attention_3(
     key_packed = torch.cat(key_valid, dim=0)
     value_packed = torch.cat(value_valid, dim=0)
 
-    out, lse, *_ = flash_attn_3_varlen_func(
-        q=query_packed,
-        k=key_packed,
-        v=value_packed,
-        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
-        softmax_scale=scale,
-        causal=is_causal,
-    )
+    kwargs = {
+        "q": query_packed,
+        "k": key_packed,
+        "v": value_packed,
+        "cu_seqlens_q": cu_seqlens_q,
+        "cu_seqlens_k": cu_seqlens_k,
+        "max_seqlen_q": max_seqlen_q,
+        "max_seqlen_k": max_seqlen_k,
+        "softmax_scale": scale,
+        "causal": is_causal,
+    }
+
+    if "return_attn_probs" in inspect.signature(flash_attn_3_varlen_func).parameters:
+        kwargs["return_attn_probs"] = return_lse
+        out = flash_attn_3_varlen_func(**kwargs)
+        if return_lse:
+            out, lse = out[0], out[1]
+        else:
+            lse = None
+    else:
+        # For backward compatibility with early flash-attn-3 APIs.
+        out, lse, *_ = flash_attn_3_varlen_func(**kwargs)
+
     out = out.unflatten(0, (batch_size, -1))
 
     return (out, lse) if return_lse else out

src/diffusers/models/transformers/transformer_z_image.py

@@ -638,7 +638,9 @@ def forward(
 
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for layer in self.layers:
-                unified = self._gradient_checkpointing_func(layer, unified, unified_attn_mask, unified_freqs_cis, adaln_input)
+                unified = self._gradient_checkpointing_func(
+                    layer, unified, unified_attn_mask, unified_freqs_cis, adaln_input
+                )
         else:
             for layer in self.layers:
                 unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input)

src/diffusers/pipelines/z_image/pipeline_z_image.py

@@ -45,7 +45,7 @@
         >>> # pipe.transformer.set_attention_backend("flash")
         >>> # (2) Use flash attention 3
         >>> # pipe.transformer.set_attention_backend("_flash_3")
-        
+
         >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。"
         >>> image = pipe(
         ...     prompt,