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

Try flattening sP and rP_bf16_qk with group_modes to fix rank mismatch

Browse Source
This commit is contained in:
biondizzle
2026-05-23 19:13:59 +00:00
parent 5ee0c20736
commit e765685951
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -273,9 +273,12 @@ class FmhaKernel:
thr_smem_copy = tiled_smem_copy.get_slice(sfw_idx)
# Debug sP shape
print(f"[SMEM-P PROPER] sP shape: {cute.shape(sP)} rank: {len(cute.shape(sP))}")
# Try without group_modes first
tSMEM_CPYsP = thr_smem_copy.partition_D(sP) # destination (SMEM)
# Try grouping ALL modes (0-4) to flatten
sP_flat = cute.group_modes(sP, 0, 4)
print(f"[SMEM-P PROPER] sP_flat shape after group_modes(0,4): {cute.shape(sP_flat)}")
tSMEM_CPYsP = thr_smem_copy.partition_D(sP_flat) # destination (SMEM)
print(f"[SMEM-P PROPER] After partition_D: tSMEM_CPYsP layout: {tSMEM_CPYsP.layout}")
print(f"[SMEM-P PROPER] After partition_D: tSMEM_CPYsP shape: {cute.shape(tSMEM_CPYsP)} rank: {len(cute.shape(tSMEM_CPYsP))}")
row_max = -Float32.inf
row_sum = Float32(0.0)
@@ -356,8 +359,12 @@ class FmhaKernel:
# Partition source with QK layout
print(f"[SMEM-P PROPER] Before partition_S: rP_bf16_qk shape: {cute.shape(rP_bf16_qk)} layout: {rP_bf16_qk.layout}")
tSMEM_CPYrP_qk = thr_smem_copy.partition_S(rP_bf16_qk)
# Try flattening source too (rP_bf16_qk has 3 modes: group_modes(0,3))
rP_bf16_qk_flat = cute.group_modes(rP_bf16_qk, 0, 3)
print(f"[SMEM-P PROPER] rP_bf16_qk_flat shape after group_modes(0,3): {cute.shape(rP_bf16_qk_flat)}")
tSMEM_CPYrP_qk = thr_smem_copy.partition_S(rP_bf16_qk_flat)
print(f"[SMEM-P PROPER] After partition_S: tSMEM_CPYrP_qk layout: {tSMEM_CPYrP_qk.layout}")
print(f"[SMEM-P PROPER] After partition_S: tSMEM_CPYrP_qk shape: {cute.shape(tSMEM_CPYrP_qk)} rank: {len(cute.shape(tSMEM_CPYrP_qk))}")
# Debug shapes
print(f"[SMEM-P PROPER] rP_bf16 shape: {cute.shape(rP_bf16)}, layout: TMEM")