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

fix: don't re-init mbarrier in loop — use phase parity tracking

Browse Source 
The mbarrier is initialized once before the loop with count=1.
Inside the loop: issue TMA → arrive.expect_tx → wait(phase) → flip phase.
Re-initializing the mbarrier inside the loop resets the phase, which
breaks the parity tracking and causes the wait to hang.

This matches the CUTLASS/gau-nernst pattern exactly.
This commit is contained in:
biondizzle
2026-05-29 18:24:47 +00:00
parent bd169ccb0f
commit 3e14a25bb0
1 changed files with 5 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
dsv4/kernels/attention/fmha_6warp_tma.cuh
View File

@@ -145,17 +145,14 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
// ==================================================================
// QK GEMM → S in TMEM (loop over K sub-tiles)
// ==================================================================
int phase = 0; // track mbarrier phase parity across the kernel
// mbarrier phase parity — tracks which phase we're waiting on.
// The mbarrier was initialized with count=1 before the loop.
// After each TMA + arrive.expect_tx + wait, we flip the phase.
// We do NOT re-init the mbarrier inside the loop.
int phase = 0;
for (int kt = 0; kt < NKT_QK; kt++) {
// --- TMA load Q sub-tile ---
// Re-init mbarrier for each TMA transaction
if (tid == 0) {
tma_mbarrier_init(mbar_addr, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
__syncthreads();
if (is_load_warp && lane == 0) {
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sQ_tma);
tma_load_2d(smem_dst, (uint64_t)tma_q, mbar_addr, kt * MMA_K_BF16, 0);
@@ -169,12 +166,6 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
__syncthreads();
// --- TMA load K sub-tile ---
if (tid == 0) {
tma_mbarrier_init(mbar_addr, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
__syncthreads();
if (is_load_warp && lane == 0) {
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sK_tma);
tma_load_2d(smem_dst, (uint64_t)tma_k, mbar_addr, kt * MMA_K_BF16, 0);
@@ -272,12 +263,6 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
__syncthreads();
// --- TMA load V sub-tile ---
if (tid == 0) {
tma_mbarrier_init(mbar_addr, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
__syncthreads();
if (is_load_warp && lane == 0) {
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sV_tma);
// V is (HD, s_k). TMA 2D: coord {col_start, d_base}