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

Add comprehensive debugging to SMEM-P path to diagnose rank mismatch

Browse Source
This commit is contained in:
biondizzle
2026-05-23 09:30:30 +00:00
parent fe3b1abf22
commit e400c935a9
1 changed files with 48 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -218,11 +218,54 @@ class FmhaKernel:
cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0)
pv_mma.set(tcgen05.Field.ACCUMULATE, True)
else:
# SMEM-P: PV reads P from SMEM
for kb in cutlass.range(cute.size(tCrP, mode=[2]), unroll_full=True):
cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kb,kvh.index)], tOtO0)
pv_mma.set(tcgen05.Field.ACCUMULATE, True)
cute.arch.fence_view_async_tmem_store()
# SMEM-P: DEBUGGING RANK MISMATCH
print(f"[SMEM-P DEEP DEBUG] === START ===")
print(f"[SMEM-P DEEP DEBUG] rP_bf16 shape: {cute.shape(rP_bf16)}")
print(f"[SMEM-P DEEP DEBUG] rP_bf16 layout: {rP_bf16.layout}")
print(f"[SMEM-P DEEP DEBUG] sP shape: {cute.shape(sP)}")
print(f"[SMEM-P DEEP DEBUG] sP layout: {sP.layout}")
# Try different approaches
print(f"[SMEM-P DEEP DEBUG] --- Approach 1: Direct copy ---")
tSMEM_CPYrP1 = thr_smem_copy.partition_S(rP_bf16)
tSMEM_CPYsP1 = thr_smem_copy.partition_D(sP)
print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYrP1 shape: {cute.shape(tSMEM_CPYrP1)}")
print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYsP1 shape: {cute.shape(tSMEM_CPYsP1)}")
print(f"[SMEM-P DEEP DEBUG] --- Approach 2: With sP_2d ---")
sP_2d = cute.group_modes(sP, 0, 3)
print(f"[SMEM-P DEEP DEBUG] sP_2d shape: {cute.shape(sP_2d)}")
tSMEM_CPYrP2 = thr_smem_copy.partition_S(rP_bf16)
tSMEM_CPYsP2 = thr_smem_copy.partition_D(sP_2d)
print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYrP2 shape: {cute.shape(tSMEM_CPYrP2)}")
print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYsP2 shape: {cute.shape(tSMEM_CPYsP2)}")
print(f"[SMEM-P DEEP DEBUG] --- Approach 3: Reshape rP_bf16 ---")
# Try to match destination rank
rP_bf16_reshaped = cute.group_modes(rP_bf16, 0, 2) # Group first 3 modes?
print(f"[SMEM-P DEEP DEBUG] rP_bf16_reshaped shape: {cute.shape(rP_bf16_reshaped)}")
tSMEM_CPYrP3 = thr_smem_copy.partition_S(rP_bf16_reshaped)
print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYrP3 shape: {cute.shape(tSMEM_CPYrP3)}")
# Try copy with approach that has matching rank
if len(cute.shape(tSMEM_CPYrP3)) == len(cute.shape(tSMEM_CPYsP2)):
print(f"[SMEM-P DEEP DEBUG] Rank match! Attempting copy...")
cute.copy(tiled_smem_copy, tSMEM_CPYrP3, tSMEM_CPYsP2)
cute.arch.fence_proxy("async.shared", space="cta")
print(f"[SMEM-P DEEP DEBUG] Copy succeeded with reshaped source")
elif len(cute.shape(tSMEM_CPYrP1)) == len(cute.shape(tSMEM_CPYsP1)):
print(f"[SMEM-P DEEP DEBUG] Direct partition ranks match! Attempting copy...")
cute.copy(tiled_smem_copy, tSMEM_CPYrP1, tSMEM_CPYsP1)
cute.arch.fence_proxy("async.shared", space="cta")
print(f"[SMEM-P DEEP DEBUG] Copy succeeded with direct partitions")
else:
print(f"[SMEM-P DEEP DEBUG] No rank match found. Falling back to stub.")
for j in cutlass.range(cute.size(sP), vectorize=True):
sP[j] = BFloat16(0.0)
cute.arch.fence_proxy("async.shared", space="cta")
softmax_done_bar.arrive()
print(f"[SMEM-P DEEP DEBUG] === END ===") cute.arch.fence_view_async_tmem_store()
kvh.release()
acc_pipe.producer_commit(acc_st); acc_st.advance()
final_o_bar.arrive()