From 1b35e0f9672e3c29fa50fa06d981627f576162da Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 22 May 2026 22:25:36 +0000 Subject: [PATCH] DEBUG: print TMA partition tensor shapes --- tests/unit/test_fmha_v3_stage_c.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 78f983ca..8f29afb0 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -179,7 +179,16 @@ 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)) + # DEBUG: print TMA partition tensor shapes BEFORE pre-slice + print(f"tAgQ shape before: {cute.shape(tAgQ)}") + print(f"tBgK shape before: {cute.shape(tBgK)}") + print(f"tVgV shape before: {cute.shape(tVgV)}") + print(f"tBsK shape: {cute.shape(tBsK)}") + print(f"tVsV shape: {cute.shape(tVsV)}") tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,0,0)]; tVgV = tVgV[(None,0,None,0)] + # DEBUG: after pre-slice + print(f"tBgK shape after pre-slice: {cute.shape(tBgK)}") + print(f"tVgV shape after pre-slice: {cute.shape(tVgV)}") tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) tCrV = pv_mma.make_fragment_B(sV)