From 46ded465da154a72a05a8634a4372b9ee2add406 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 22 May 2026 17:28:45 +0000 Subject: [PATCH] DEBUG: print tBgK/tVgV shapes before/after slice --- tests/unit/test_fmha_v3_stage_c_full.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/test_fmha_v3_stage_c_full.py b/tests/unit/test_fmha_v3_stage_c_full.py index 13f7b69a..7976bc5b 100644 --- a/tests/unit/test_fmha_v3_stage_c_full.py +++ b/tests/unit/test_fmha_v3_stage_c_full.py @@ -171,9 +171,13 @@ class FmhaV3StageC: # CUTLASS reference: tKgK = tKgK_kdl[None, None, 0, 0] (keeps TMA_atom + GMEM_iter) # tVgV = tVgV_dkl[None, 0, None, 0] (keeps TMA_atom + GMEM_iter at mode 2) # SMEM tensors from tma_partition are already 2D — don't re-slice them. + print(f"DEBUG tBgK shape before slice: {tBgK.shape}") + print(f"DEBUG tVgV shape before slice: {tVgV.shape}") tAgQ = tAgQ[(None,0,None,0)] # Q: 1 tile, hardcode is fine tBgK = tBgK[(None,None,0,0)] # K: keep mode 1 (GMEM iter) free tVgV = tVgV[(None,0,None,0)] # V: keep mode 2 (GMEM iter) free + print(f"DEBUG tBgK shape after slice: {tBgK.shape}") + print(f"DEBUG tVgV shape after slice: {tVgV.shape}") tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) tCrV = pv_mma.make_fragment_B(sV)