From fe58d07b8cb395b11a052292ce41c93a8f0f64b5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 22 May 2026 18:51:22 +0000 Subject: [PATCH] FIX: K slice (None,None,0,0) like working diag --- tests/fmha_v3_identity_diag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/fmha_v3_identity_diag.py b/tests/fmha_v3_identity_diag.py index 76241ddf..c6ed5c81 100644 --- a/tests/fmha_v3_identity_diag.py +++ b/tests/fmha_v3_identity_diag.py @@ -138,7 +138,7 @@ class FmhaV3IdentityDiag: tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) # Use the SAME slice as the working diag test - tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] + tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,0,0)]; tVgV = tVgV[(None,0,None,0)] tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) tCrV = pv_mma.make_fragment_B(sV)