From 4c1dbfd0f3afb51146a0f96b28b07f963e273d78 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 22 May 2026 21:56:22 +0000 Subject: [PATCH] Add cute.printf shape diagnostics to example9 --- tests/fmha_v3_stage_c_example9.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/fmha_v3_stage_c_example9.py b/tests/fmha_v3_stage_c_example9.py index f1cc666f..0381c543 100644 --- a/tests/fmha_v3_stage_c_example9.py +++ b/tests/fmha_v3_stage_c_example9.py @@ -179,6 +179,17 @@ class FmhaV3StageCMulti: b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) 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)) + # DIAG: print shapes before pre-slice + cute.printf("tAgQ shape: {} rank: {}\n", cute.shape(tAgQ), tAgQ.rank) + cute.printf("tBgK shape: {} rank: {}\n", cute.shape(tBgK), tBgK.rank) + cute.printf("tVgV shape: {} rank: {}\n", cute.shape(tVgV), tVgV.rank) + cute.printf("tAsQ shape: {} rank: {}\n", cute.shape(tAsQ), tAsQ.rank) + cute.printf("tBsK shape: {} rank: {}\n", cute.shape(tBsK), tBsK.rank) + cute.printf("tVsV shape: {} rank: {}\n", cute.shape(tVsV), tVsV.rank) + for i in range(tBgK.rank): + cute.printf(" tBgK mode {} size: {}\n", i, cute.size(tBgK, mode=[i])) + for i in range(tVgV.rank): + cute.printf(" tVgV mode {} size: {}\n", i, cute.size(tVgV, mode=[i])) 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)