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

Add debug prints to SMEM-P path to understand rank mismatch

Browse Source
This commit is contained in:
biondizzle
2026-05-23 09:25:48 +00:00
parent f6b43227e5
commit 69f458ac30
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -346,6 +346,13 @@ class FmhaKernel:
# rP_bf16 contains P values in QK C-fragment layout (BF16)
# Use rP_bf16 directly (already in correct layout for QK C-fragment)
tSMEM_CPYrP = thr_smem_copy.partition_S(rP_bf16)
# DEBUG: Print shapes before copy
print(f"[SMEM-P DEBUG] rP_bf16 shape: {cute.shape(rP_bf16)}")
print(f"[SMEM-P DEBUG] tSMEM_CPYrP shape: {cute.shape(tSMEM_CPYrP)}")
print(f"[SMEM-P DEBUG] tSMEM_CPYsP shape: {cute.shape(tSMEM_CPYsP)}")
print(f"[SMEM-P DEBUG] rP_bf16 rank: {len(cute.shape(rP_bf16))}")
print(f"[SMEM-P DEBUG] tSMEM_CPYrP rank: {len(cute.shape(tSMEM_CPYrP))}")
print(f"[SMEM-P DEBUG] tSMEM_CPYsP rank: {len(cute.shape(tSMEM_CPYsP))}")
cute.copy(tiled_smem_copy, tSMEM_CPYrP, tSMEM_CPYsP)
cute.arch.fence_proxy("async.shared", space="cta")
softmax_done_bar.arrive()