biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

DEBUG: print flat_divide shapes

Browse Source
This commit is contained in:
biondizzle
2026-05-22 22:31:24 +00:00
parent f64af7b0e0
commit 88b70e56fb
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -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])