diag: print bSG shapes for TMA store indexing

2026-05-23 02:44:47 +00:00
parent 1d3c169805
commit c403b57617
1 changed files with 3 additions and 0 deletions
--- 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])