Merge pull request #245 from aigc-apps/fix_sage

Set tensor_layout in sageattn
2025-07-04 15:18:51 +08:00
parent dcd1b32ee6 9c16b3a73f
commit 21ec5bf326
1 changed files with 1 additions and 6 deletions
@@ -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,