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

D2: add shape debug print for mQ/mK

Browse Source
This commit is contained in:
biondizzle
2026-05-24 23:37:10 +00:00
parent df3146eb53
commit b5cd1b88c9
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
dsv4/kernels/attention/fmha.py
View File

@@ -189,6 +189,9 @@ class FmhaKernel:
# D2: Multi-CTA grid. Use block_idx_y to select Q and O for this CTA's head.
_bidx, head_cta_idx, _bidz = cute.arch.block_idx() # grid=(1, num_ctas, 1)
# Print mQ shape at trace time to understand mode structure
print(f"D2 DEBUG: mQ shape={cute.shape(mQ)}, mK shape={cute.shape(mK)}")
# Q: if num_ctas > 1, mQ has a head dimension. local_tile indexes into it.
# K/V: shared (MQA), always coordinate 0.
# For single-CTA (num_ctas=1), head_cta_idx=0 and coordinates are the same as before.