From 3e14a25bb0a93a5b36f4b3ed5942a23f5b9adefe Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 29 May 2026 18:24:47 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20don't=20re-init=20mbarrier=20in=20loop?= =?UTF-8?q?=20=E2=80=94=20use=20phase=20parity=20tracking?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- dsv4/kernels/attention/fmha_6warp_tma.cuh | 25 +++++------------------ 1 file changed, 5 insertions(+), 20 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_tma.cuh b/dsv4/kernels/attention/fmha_6warp_tma.cuh index 6f0bba88..0e3f7a08 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma.cuh @@ -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}