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

Fix diag: remove .rank

Browse Source
This commit is contained in:
biondizzle
2026-05-22 21:56:56 +00:00
parent 4c1dbfd0f3
commit 254f7be884
1 changed files with 6 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
tests/fmha_v3_stage_c_example9.py
View File

@@ -180,16 +180,12 @@ class FmhaV3StageCMulti:
tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3))
tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3))
# DIAG: print shapes before pre-slice
cute.printf("tAgQ shape: {} rank: {}\n", cute.shape(tAgQ), tAgQ.rank)
cute.printf("tBgK shape: {} rank: {}\n", cute.shape(tBgK), tBgK.rank)
cute.printf("tVgV shape: {} rank: {}\n", cute.shape(tVgV), tVgV.rank)
cute.printf("tAsQ shape: {} rank: {}\n", cute.shape(tAsQ), tAsQ.rank)
cute.printf("tBsK shape: {} rank: {}\n", cute.shape(tBsK), tBsK.rank)
cute.printf("tVsV shape: {} rank: {}\n", cute.shape(tVsV), tVsV.rank)
for i in range(tBgK.rank):
cute.printf(" tBgK mode {} size: {}\n", i, cute.size(tBgK, mode=[i]))
for i in range(tVgV.rank):
cute.printf(" tVgV mode {} size: {}\n", i, cute.size(tVgV, mode=[i]))
cute.printf("tAgQ shape: {}\n", cute.shape(tAgQ))
cute.printf("tBgK shape: {}\n", cute.shape(tBgK))
cute.printf("tVgV shape: {}\n", cute.shape(tVgV))
cute.printf("tAsQ shape: {}\n", cute.shape(tAsQ))
cute.printf("tBsK shape: {}\n", cute.shape(tBsK))
cute.printf("tVsV shape: {}\n", cute.shape(tVsV))
tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,0,0)]; tVgV = tVgV[(None,0,None,0)]
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)