update sage_attn2

lightx2v/attentions/common/sage_attn2.py

@@ -13,34 +13,28 @@ else:
 
 
 def sage_attn2(q, k, v, cu_seqlens_q=None, cu_seqlens_kv=None, max_seqlen_q=None, max_seqlen_kv=None, model_cls="hunyuan"):
-    q, k, v = (
-        q.transpose(1, 0).contiguous(),
-        k.transpose(1, 0).contiguous(),
-        v.transpose(1, 0).contiguous(),
-    )
-
+    q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
     if model_cls == "hunyuan":
         x1 = sageattn(
-            q[:, : cu_seqlens_q[1], :].unsqueeze(0),
-            k[:, : cu_seqlens_kv[1], :].unsqueeze(0),
-            v[:, : cu_seqlens_kv[1], :].unsqueeze(0),
+            q[: cu_seqlens_q[1]].unsqueeze(0),
+            k[: cu_seqlens_kv[1]].unsqueeze(0),
+            v[: cu_seqlens_kv[1]].unsqueeze(0),
+            tensor_layout="NHD",
         )
         x2 = sageattn(
-            q[:, cu_seqlens_q[1] :, :].unsqueeze(0),
-            k[:, cu_seqlens_kv[1] :, :].unsqueeze(0),
-            v[:, cu_seqlens_kv[1] :, :].unsqueeze(0),
+            q[cu_seqlens_q[1] :].unsqueeze(0),
+            k[cu_seqlens_kv[1] :].unsqueeze(0),
+            v[cu_seqlens_kv[1] :].unsqueeze(0),
+            tensor_layout="NHD",
         )
-        x = torch.cat((x1, x2), dim=-2).transpose(2, 1).contiguous()
+        x = torch.cat((x1, x2), dim=1)
         x = x.view(max_seqlen_q, -1)
     elif model_cls == "wan2.1":
-        x = (
-            sageattn(
-                q[:, : cu_seqlens_q[1], :].unsqueeze(0),
-                k[:, : cu_seqlens_kv[1], :].unsqueeze(0),
-                v[:, : cu_seqlens_kv[1], :].unsqueeze(0),
-            )
-            .transpose(2, 1)
-            .contiguous()
+        x = sageattn(
+            q.unsqueeze(0),
+            k.unsqueeze(0),
+            v.unsqueeze(0),
+            tensor_layout="NHD",
         )
         x = x.view(max_seqlen_q, -1)
     return x