From 9c16b3a73f7739e643a4d4b89fbd9c88e8477fb4 Mon Sep 17 00:00:00 2001 From: "huangkunzhe.hkz" Date: Fri, 27 Jun 2025 17:58:41 +0800 Subject: [PATCH] Set tensor_layout in sageattn --- videox_fun/models/wan_transformer3d.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/videox_fun/models/wan_transformer3d.py b/videox_fun/models/wan_transformer3d.py index ef451b5..c03155a 100755 --- a/videox_fun/models/wan_transformer3d.py +++ b/videox_fun/models/wan_transformer3d.py @@ -196,14 +196,9 @@ def attention( ) attn_mask = None - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - out = sageattn( - q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + q, k, v, attn_mask=attn_mask, tensor_layout="NHD", is_causal=causal, dropout_p=dropout_p) - out = out.transpose(1, 2).contiguous() elif attention_type == "FLASH_ATTENTION" and (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): return flash_attention( q=q,