From 88b70e56fb029fc91e3e7a62ec6d86a568c5941e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 22 May 2026 22:31:24 +0000 Subject: [PATCH] DEBUG: print flat_divide shapes --- tests/unit/test_fmha_v3_stage_c.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 07727012..08909cce 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -173,6 +173,10 @@ class FmhaV3StageCMulti: gK = cute.flat_divide(mK, cute.select(self.qk_mma_tiler, mode=[1, 2])) gV = cute.flat_divide(mV, cute.select(self.pv_mma_tiler, mode=[1, 2])) gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) + # DEBUG + print(f"gK flat_divide shape: {cute.shape(gK)}") + print(f"gV flat_divide shape: {cute.shape(gV)}") + print(f"gQ local_tile shape: {cute.shape(gQ)}") # flat_divide layout: (tile_K, tile_L, n_kv_tiles, ...) # n_kv_tiles is at mode 2 for K (0-indexed) n_kv_tiles = cute.size(gK, mode=[2])