From c403b57617055f67c3edeac8a6dc62c51a5ff9a5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 02:44:47 +0000 Subject: [PATCH] diag: print bSG shapes for TMA store indexing --- tests/unit/test_fmha_v3_stage_c.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 8cf4003a..6b60e81b 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -140,6 +140,9 @@ class FmhaV3StageCMulti: cute.group_modes(sC, 0, 2), cute.group_modes(gC_epi, 0, 2), ) + # Print TMA partition shapes for debugging + print(f'DEBUG bSG_sC shape: {bSG_sC.shape} rank: {cute.rank(bSG_sC)}', flush=True) + print(f'DEBUG bSG_gC shape: {bSG_gC.shape} rank: {cute.rank(bSG_gC)}', flush=True) n_kv_tiles = cute.size(gK, mode=[3])