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

Clean debug SMEM-P path to understand rank mismatch

Browse Source
This commit is contained in:
biondizzle
2026-05-23 09:31:24 +00:00
parent ef0ac2c187
commit 303c2c5275
1 changed files with 37 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -218,54 +218,10 @@ 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: 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 ===")
# 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()
kvh.release()
acc_pipe.producer_commit(acc_st); acc_st.advance()
@@ -388,15 +344,39 @@ class FmhaKernel:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
else:
# SMEM-P: write P to SMEM via tiled_smem_copy
# 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))}")
else:
# SMEM-P: DEBUGGING RANK MISMATCH
print(f"[SMEM-P] rP_bf16 shape: {cute.shape(rP_bf16)} rank: {len(cute.shape(rP_bf16))}")
print(f"[SMEM-P] sP shape: {cute.shape(sP)} rank: {len(cute.shape(sP))}")
# Approach 1: Try sP directly (no group_modes)
tSMEM_CPYrP1 = thr_smem_copy.partition_S(rP_bf16)
tSMEM_CPYsP1 = thr_smem_copy.partition_D(sP)
print(f"[SMEM-P] Approach1 - src shape: {cute.shape(tSMEM_CPYrP1)} rank: {len(cute.shape(tSMEM_CPYrP1))}")
print(f"[SMEM-P] Approach1 - dst shape: {cute.shape(tSMEM_CPYsP1)} rank: {len(cute.shape(tSMEM_CPYsP1))}")
# Approach 2: Try with sP_2d
sP_2d = cute.group_modes(sP, 0, 3)
print(f"[SMEM-P] sP_2d shape: {cute.shape(sP_2d)} rank: {len(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] Approach2 - src shape: {cute.shape(tSMEM_CPYrP2)} rank: {len(cute.shape(tSMEM_CPYrP2))}")
print(f"[SMEM-P] Approach2 - dst shape: {cute.shape(tSMEM_CPYsP2)} rank: {len(cute.shape(tSMEM_CPYsP2))}")
# Try copy with matching ranks
if len(cute.shape(tSMEM_CPYrP1)) == len(cute.shape(tSMEM_CPYsP1)):
print(f"[SMEM-P] Attempting copy with direct partitions")
cute.copy(tiled_smem_copy, tSMEM_CPYrP1, tSMEM_CPYsP1)
elif len(cute.shape(tSMEM_CPYrP2)) == len(cute.shape(tSMEM_CPYsP2)):
print(f"[SMEM-P] Attempting copy with sP_2d partitions")
cute.copy(tiled_smem_copy, tSMEM_CPYrP2, tSMEM_CPYsP2)
else:
print(f"[SMEM-P] No rank match. Using 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 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)