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

fix: move epilogue TMEM loads outside my_row_active guard (warp-collective hang)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 23:46:46 +00:00
parent 55c0604a71
commit ac8fa779e2
1 changed files with 21 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
dsv4/kernels/attention/fmha_6warp_multirow.cuh
View File

@@ -87,7 +87,7 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
// SMEM allocation
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
float* sRowMax = (float*)(sbuf + 8);
float* sRowMax = (float*)(sbuf + 4); // sTmemBase is 1 uint32_t = 4 bytes
float* sRowSum = sRowMax + MAX_ROWS;
bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 127) & ~(uintptr_t)127);
bf16_t* sK0 = sQ0 + TILE_SZ;
@@ -265,30 +265,36 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
// ================================================================
// EPILOGUE: TMEM → regs → normalize → BF16 → GMEM
//
// CRITICAL: TMEM loads (32x32b.x8) are WARP-COLLECTIVE.
// ALL 32 lanes must execute them. The load MUST be outside
// the my_row_active guard. Only the GMEM store is conditional.
// ================================================================
if (my_warp_active) {
if (my_row_active) {
float rm = sRowMax[my_row];
float rs = sRowSum[my_row];
float inv_rs = 1.0f / rs;
float rm = my_row_active ? sRowMax[my_row] : 0.0f;
float rs = my_row_active ? sRowSum[my_row] : 0.0f;
float inv_rs = my_row_active ? (1.0f / rs) : 0.0f;
// Read O from TMEM: N_NSUB*2 groups of 8 columns
for (int n = 0; n < N_NSUB * 2; n++) {
float tmp[8];
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
: "r"(tb + n * 8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
// Read O from TMEM: N_NSUB*2 groups of 8 columns
// ALL lanes in the warp must execute the TMEM load (warp-collective)
for (int n = 0; n < N_NSUB * 2; n++) {
float tmp[8];
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
: "r"(tb + n * 8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
// Only store to GMEM for active rows
if (my_row_active) {
for (int c = 0; c < 8; c++) {
int d = n * 8 + c;
if (d < HD) o_head[my_row * HD + d] = f32_to_bf16(tmp[c] * inv_rs);
}
}
if (lse_head) lse_head[my_row] = logf(rs) + rm;
}
if (my_row_active && lse_head) lse_head[my_row] = logf(rs) + rm;
}
__syncthreads();