From 4a9c850e9c25aaf5c0b56e5a386cd9f32ee810ca Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 03:14:06 +0000 Subject: [PATCH] feat: double-buffer TMA pipeline for K loads in single-tile kernel --- dsv4/kernels/attention/fmha_6warp_tma.cuh | 194 +++++++++++++--------- tests/unit/test_fmha_6warp_tma.cu | 2 + 2 files changed, 113 insertions(+), 83 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_tma.cuh b/dsv4/kernels/attention/fmha_6warp_tma.cuh index ab753032..5abb50bb 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma.cuh @@ -1,39 +1,28 @@ /** - * DSV4 FMHA — 6-warp specialized kernel with TMA async loads for Blackwell SM100. + * DSV4 FMHA — 6-warp specialized kernel with double-buffer TMA pipeline. * * ================================================================== - * WARP SPECIALIZATION (same as fmha_6warp.cuh) + * DOUBLE-BUFFER TMA PIPELINE * ================================================================== - * Warp 0-3 (tid 0-127): Softmax + correction + epilogue - * Warp 4 (tid 128-159): MMA (QK + PV) - * Warp 5 (tid 160-191): TMA loads (K, V) + Q direct load + P fill * - * ================================================================== - * TMA LOAD PIPELINE - * ================================================================== - * - Q: direct GMEM load (T=1 decode, only 1 row; small enough to not benefit from TMA) - * - K: TMA async load via cp.async.bulk.tensor.2d with mbarrier completion - * - V: TMA async load (future; currently direct GMEM load) - * - mbarrier: init once, arrive.expect_tx after TMA issue, phase parity tracking - * - ALL threads wait on mbarrier (works correctly at 192 threads, verified on B200) + * K is loaded via TMA with double-buffering: + * - Two canonical SMEM buffers: sK0 and sK1 + * - Preload kt=0 into sK0 + * - Loop: issue TMA for kt+1 → sTmaBuf (while MMA runs on current sK), + * then convert sTmaBuf → next sK, swap buffers + * - TMA DMA transfer overlaps with MMA compute * - * ================================================================== - * SMEM LAYOUT - * ================================================================== - * All SMEM regions are 128-byte aligned for TMA compatibility. - * sTmaBuf is used as a staging area: TMA writes row-major data here, - * then it's converted to canonical layout for MMA. + * The pipeline stages for kt=N: + * 1. Issue TMA load for kt=N+1 (if exists) → sTmaBuf + * 2. MMA QK using sK[current] (runs while TMA DMA is in flight) + * 3. Wait for TMA completion (usually already done after MMA) + * 4. Convert sTmaBuf → sK[next] (canonical layout) + * 5. Swap current/next buffer pointers * - * sTmemBase: 4 bytes - * sMbar: 16 bytes (128-byte aligned) - * sTmaBuf: TILE_SZ * 2 bytes (128-byte aligned) — TMA staging buffer - * sQ0: TILE_SZ * 2 bytes (128-byte aligned) — canonical for MMA - * sK0: TILE_SZ * 2 bytes (128-byte aligned) — canonical for MMA - * sPk: TILE_SZ * 2 bytes (128-byte aligned) — canonical for PV - * sV: V_SUB_SZ bytes (128-byte aligned) — canonical for PV - * sRowMax: MAX_ROWS * 4 bytes - * sRowSum: MAX_ROWS * 4 bytes - * s_p_vals: SK_TILE * 4 bytes + * V is also loaded via TMA but NOT double-buffered (V tiles are tiny: + * 16×16 = 512 bytes, TMA latency is negligible). + * + * Q is loaded directly from GMEM (T=1 decode, only 1 row). */ #pragma once @@ -52,17 +41,16 @@ fmha_6warp_tma_kernel( CUtensorMap* __restrict__ tma_v, bf16_t* __restrict__ o, float* __restrict__ lse, - int s_k, - float scale + int s_k, float scale ) { static constexpr int NKT_QK = HD / MMA_K_BF16; static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; static constexpr int N_NSUB = HD / 16; static constexpr int TILE_SZ = 128 * MMA_K_BF16; - static constexpr int V_SUB_SZ = 16 * MMA_K_BF16; // (16,16) canonical + static constexpr int V_SUB_SZ = 16 * MMA_K_BF16; static constexpr int TMEM_N = (HD <= 128) ? 128 : 256; static constexpr int MAX_ROWS = 128; - static constexpr int CORES_MN = 128 / 8; // 16 + static constexpr int CORES_MN = 128 / 8; static constexpr int NUM_READS = SK_TILE / 8; static constexpr int TMA_TILE_BYTES = TILE_SZ * sizeof(bf16_t); @@ -70,35 +58,31 @@ fmha_6warp_tma_kernel( const int wid = tid / 32; const int lane = tid % 32; - // Warp role predicates const bool is_softmax_warp = (wid < 4); const bool is_mma_warp = (wid == 4); const bool is_load_warp = (wid == 5); // ================================================================ - // SMEM allocation — all 128-byte aligned + // SMEM allocation — 128-byte aligned // ================================================================ extern __shared__ __align__(128) char sbuf[]; size_t off = 0; uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off += 4; - off = (off + 127) & ~(size_t)127; // 128-byte align - + off = (off + 127) & ~(size_t)127; uint64_t* sMbar = (uint64_t*)(sbuf + off); off += 16; off = (off + 127) & ~(size_t)127; bf16_t* sTmaBuf = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); off = (off + 127) & ~(size_t)127; - bf16_t* sQ0 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); off = (off + 127) & ~(size_t)127; - bf16_t* sK0 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); off = (off + 127) & ~(size_t)127; - + bf16_t* sK1 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); + off = (off + 127) & ~(size_t)127; bf16_t* sPk = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); off = (off + 127) & ~(size_t)127; - bf16_t* sV = (bf16_t*)(sbuf + off); off += V_SUB_SZ * sizeof(bf16_t); float* sRowMax = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); @@ -106,11 +90,9 @@ fmha_6warp_tma_kernel( float* s_p_vals = (float*)(sbuf + off); off += SK_TILE * sizeof(float); // ================================================================ - // Initialize TMEM + mbarrier + // Init // ================================================================ - if (is_mma_warp) { - tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N); - } + if (is_mma_warp) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N); if (tid == 0) { tma_mbarrier_init((uint32_t)__cvta_generic_to_shared(sMbar), 1); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); @@ -121,51 +103,98 @@ fmha_6warp_tma_kernel( int phase = 0; // ================================================================ - // QK GEMM loop: TMA for K, direct load for Q + // QK GEMM — double-buffer TMA pipeline // ================================================================ - for (int kt = 0; kt < NKT_QK; kt++) { - // ---- Warp 5: Load Q K-tile directly from GMEM ---- - if (is_load_warp) { - for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; - for (int d = lane; d < MMA_K_BF16; d += 32) { - int full_d = kt * MMA_K_BF16 + d; - if (full_d < HD) { - int ck = d / 8, lc = d % 8; - sQ0[ck * CORES_MN * 64 + lc] = q[full_d]; - } - } - } + { + uint32_t idesc = make_idesc(128, 128); + bf16_t* sK_bufs[2] = {sK0, sK1}; + int cur_buf = 0; // alternates 0,1,0,1,... - // ---- TMA load K sub-tile ---- - // Only 1 thread issues TMA (warp 5, lane 0) + // --- Preload kt=0: TMA → sTmaBuf → convert → sK0 --- if (is_load_warp && lane == 0) { tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_k, - mbar_addr, kt * MMA_K_BF16, 0); + mbar_addr, 0, 0); tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); } - // ALL threads wait for TMA completion tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; __syncthreads(); - // ---- Convert sTmaBuf (row-major) → sK0 (canonical) ---- - // All threads participate in the conversion + // Convert sTmaBuf → sK0 (canonical) for (int i = tid; i < TILE_SZ; i += 192) sK0[i] = 0; for (int i = tid; i < s_k * MMA_K_BF16; i += 192) { int r = i / MMA_K_BF16, c = i % MMA_K_BF16; int ck = c / 8, lc = c % 8, tmn = r / 8, lr = r % 8; sK0[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = sTmaBuf[i]; } - __syncthreads(); - // ---- Warp 4: QK MMA ---- - if (is_mma_warp) { - uint32_t idesc = make_idesc(128, 128); - uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ0), 128); - uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK0), 128); - if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + // Also load Q for kt=0 + if (is_load_warp) { + for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; + for (int d = lane; d < MMA_K_BF16; d += 32) { + if (d < HD) { + int ck = d / 8, lc = d % 8; + sQ0[ck * CORES_MN * 64 + lc] = q[d]; + } + } } __syncthreads(); + + // --- Pipeline loop --- + for (int kt = 0; kt < NKT_QK; kt++) { + bf16_t* cur_sK = sK_bufs[cur_buf]; + + // 1. Issue TMA for kt+1 (if exists) → sTmaBuf + bool has_next = (kt + 1 < NKT_QK); + if (has_next && is_load_warp && lane == 0) { + tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_k, + mbar_addr, (kt + 1) * MMA_K_BF16, 0); + tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); + } + + // 2. MMA QK using cur_sK (TMA DMA runs in parallel!) + if (is_mma_warp) { + uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ0), 128); + uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(cur_sK), 128); + if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + } + // Note: no __syncthreads() here — MMA and TMA can overlap across warps. + // But we need the MMA to complete before reading TMEM (softmax). + // The softmax warps will sync later. + + // 3. Wait for TMA kt+1 completion (usually already done after MMA) + if (has_next) { + tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; + } + __syncthreads(); + + // 4. Convert sTmaBuf → next sK buffer (canonical) + if (has_next) { + bf16_t* next_sK = sK_bufs[1 - cur_buf]; + for (int i = tid; i < TILE_SZ; i += 192) next_sK[i] = 0; + for (int i = tid; i < s_k * MMA_K_BF16; i += 192) { + int r = i / MMA_K_BF16, c = i % MMA_K_BF16; + int ck = c / 8, lc = c % 8, tmn = r / 8, lr = r % 8; + next_sK[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = sTmaBuf[i]; + } + } + + // 5. Load Q for next kt (if exists) + if (has_next && is_load_warp) { + for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; + for (int d = lane; d < MMA_K_BF16; d += 32) { + int full_d = (kt + 1) * MMA_K_BF16 + d; + if (full_d < HD) { + int ck = d / 8, lc = d % 8; + sQ0[ck * CORES_MN * 64 + lc] = q[full_d]; + } + } + } + + // 6. Swap buffers + cur_buf = 1 - cur_buf; + __syncthreads(); + } } // ================================================================ @@ -200,27 +229,29 @@ fmha_6warp_tma_kernel( __syncthreads(); // ================================================================ - // PV GEMM loop: N=16 sub-tiles × K-tiles + // PV GEMM loop: N=16 sub-tiles × K-tiles, V via TMA // ================================================================ for (int n = 0; n < N_NSUB; n++) { int d_base = n * 16; for (int kt = 0; kt < NKT_PV; kt++) { - // ---- Fill sPk ---- + const int col_start = kt * MMA_K_BF16; + + // Fill sPk if (is_load_warp) { for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0; if (lane < 16) { int c = lane; int ck = c / 8, lc = c % 8; - sPk[ck * CORES_MN * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[kt * MMA_K_BF16 + c]); + sPk[ck * CORES_MN * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[col_start + c]); } } __syncthreads(); - // ---- Load V via TMA ---- + // Load V via TMA if (is_load_warp && lane == 0) { tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_v, - mbar_addr, kt * MMA_K_BF16, d_base); + mbar_addr, col_start, d_base); tma_mbarrier_arrive_expect_tx(mbar_addr, V_SUB_SZ * sizeof(bf16_t)); } tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; @@ -235,7 +266,7 @@ fmha_6warp_tma_kernel( } __syncthreads(); - // ---- Warp 4: PV MMA ---- + // MMA if (is_mma_warp) { uint32_t idesc_pv16 = make_idesc(128, 16); uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), 128); @@ -248,7 +279,7 @@ fmha_6warp_tma_kernel( } // ================================================================ - // Epilogue: TMEM → regs → normalize → BF16 → GMEM (warp 0) + // Epilogue: TMEM → regs → BF16 → GMEM (warp 0) // ================================================================ if (wid == 0) { float o_vals[HD]; @@ -266,10 +297,7 @@ fmha_6warp_tma_kernel( } __syncthreads(); - // TMEM dealloc (warp 4) - if (is_mma_warp) { - tmem_dealloc(tb, TMEM_N); - } + if (is_mma_warp) tmem_dealloc(tb, TMEM_N); } } // namespace diff --git a/tests/unit/test_fmha_6warp_tma.cu b/tests/unit/test_fmha_6warp_tma.cu index d5a4f8e3..7a62fc7c 100644 --- a/tests/unit/test_fmha_6warp_tma.cu +++ b/tests/unit/test_fmha_6warp_tma.cu @@ -44,6 +44,8 @@ static size_t compute_smem() { off = (off + 127) & ~(size_t)127; off += TILE_SZ * sizeof(bf16_t); // sK0 off = (off + 127) & ~(size_t)127; + off += TILE_SZ * sizeof(bf16_t); // sK1 (double buffer) + off = (off + 127) & ~(size_t)127; off += TILE_SZ * sizeof(bf16_t); // sPk off = (off + 127) & ~(size_t)127; off += V_SUB_SZ * sizeof(bf16_t); // sV