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

DEBUG: print tBgK/tVgV layout to check strides

Browse Source
This commit is contained in:
biondizzle
2026-05-22 22:28:32 +00:00
parent be103deb6d
commit daa4017505
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

@@ -181,6 +181,10 @@ class FmhaV3StageCMulti:
tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3))
tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)]
# DEBUG: print shapes and strides to verify kv_tiles dimension has nonzero stride
print(f"tBgK pre-sliced shape: {cute.shape(tBgK)} layout: {tBgK.layout}")
print(f"tVgV pre-sliced shape: {cute.shape(tVgV)} layout: {tVgV.layout}")
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
tCrV = pv_mma.make_fragment_B(sV)