From ca23a2fc59bfc1b6062793cf1994db3b1ae8a7e9 Mon Sep 17 00:00:00 2001
From: Emanuele Bugliarello <emanuele.bugliarello@gmail.com>
Date: Wed, 27 Aug 2025 11:43:35 +0200
Subject: [PATCH] Fix flash attention

fa3 latest version changed the return shape of the varlen func to be consistent w fa2. this pr fixes the fa3 attention call as done in https://github.com/Wan-Video/Wan2.2/pull/64
---
 wan/modules/attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wan/modules/attention.py b/wan/modules/attention.py
index 4dbbe03..c8f69c2 100644
--- a/wan/modules/attention.py
+++ b/wan/modules/attention.py
@@ -107,7 +107,7 @@ def flash_attention(
             max_seqlen_k=lk,
             softmax_scale=softmax_scale,
             causal=causal,
-            deterministic=deterministic)[0].unflatten(0, (b, lq))
+            deterministic=deterministic).unflatten(0, (b, lq))
     else:
         assert FLASH_ATTN_2_AVAILABLE
         x = flash_attn.flash_attn_varlen_func(