From 212fc8562757d6e6bec1c0aca94ebacf2872d814 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 16:56:07 +0000 Subject: [PATCH] =?UTF-8?q?P6:=20One-way=20TMEM=E2=86=92regs=E2=86=92SMEM?= =?UTF-8?q?=E2=86=92TMA=20store=20epilogue?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fmha_6warp_multihead.cuh: Rewritten epilogue with proper Blackwell pipeline 1. TMEM → regs (tcgen05.ld, warp-collective) 2. epilogue_op in regs (normalize, FP4 hook via ENABLE_FP4_EPILOGUE) 3. regs → SMEM row-major (sO_epi, for TMA tile format) 4. TMA store SMEM → GMEM (async, enables multi-CTA) Fallback to direct GMEM write when tma_o is nullptr. Added FmhaParams.tma_o field and ENABLE_FP4_EPILOGUE template param. - fmha_6warp_tma_multirow_multitile.cuh: Same epilogue pattern for multi-tile. Writes normalized output to sO_epi_rowmajor + TMA store (or direct GMEM). Added tma_o to FmhaTmaMultiRowMultiTileParams. - fmha_tma.cuh: Added tma_store_2d and tma_store_wait for async GMEM writes. - fmha_multihead_capi.cu: Added fmha_multihead_decode_tma_launch with per-(head,batch) TMA descriptors. Updated SMEM size calculation for sO_epi + sMbarStore. - fmha_multitile_capi.cu: Added tma_o=nullptr (backward compatible), updated SMEM size. --- .../attention/fmha_6warp_multihead.cuh | 186 ++++++++++++------ .../fmha_6warp_tma_multirow_multitile.cuh | 52 ++++- dsv4/kernels/attention/fmha_multihead_capi.cu | 132 +++++++++++-- dsv4/kernels/attention/fmha_multitile_capi.cu | 5 + dsv4/kernels/attention/fmha_tma.cuh | 54 +++++ 5 files changed, 354 insertions(+), 75 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_multihead.cuh b/dsv4/kernels/attention/fmha_6warp_multihead.cuh index a8a8c87a..11543235 100644 --- a/dsv4/kernels/attention/fmha_6warp_multihead.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multihead.cuh @@ -12,42 +12,27 @@ * No cross-CTA synchronization required. * * ================================================================== - * MQA / GQA SUPPORT + * EPILOGUE (P6 — One-way TMEM → regs → SMEM → TMA store → GMEM) * ================================================================== - * - MQA: all Q heads share one KV head. Pass k_head_stride=0, v_head_stride=0 - * so all CTAs read the same K/V. - * - GQA: groups of Q heads share a KV head. The caller must arrange - * K/V tensors so that k_head_stride/v_head_stride map correctly. - * - MHA: k_head_stride = k_row_stride * N, same for V. + * The proper Blackwell output pipeline: + * 1. TMEM → registers (tcgen05.ld, warp-collective) + * 2. epilogue_op in registers (normalize + optional FP4 pack) + * 3. Registers → SMEM (row-major, matching TMA tile format) + * 4. TMA store SMEM → GMEM (async, enables multi-CTA) + * + * This replaces the old direct GMEM write and unblocks: + * - D2 multi-CTA grid (TMA store with flat_divide coords) + * - NVFP4-1.2 FP4 output fusion (register slot for amax + pack) + * - Proper async pipeline overlap + * + * When tma_o is nullptr, falls back to direct GMEM write from registers. + * When tma_o is set, uses the proper TMA store pipeline. * * ================================================================== - * TENSOR LAYOUTS (GMEM) + * OUTPUT * ================================================================== - * Q: [batch, n_h, T, hd] — head stride = T * hd, batch stride = n_h * T * hd - * K: [batch, n_kv, N, hd] — head stride = N * hd (or 0 for MQA) - * V: [batch, n_kv, hd, N] — head stride = hd * N (or 0 for MQA) - * O: [batch, n_h, T, hd] — same strides as Q - * - * For decode (T=1): q_head_offset = blockIdx.y * hd, q_batch_offset = blockIdx.z * n_h * hd - * For prefill (T>1): head-packed M = T rows per head (must fit in 128-row MMA tile) - * - * ================================================================== - * SOFTMAX ROWS - * ================================================================== - * T=1 decode: only row 0 of the 128-row MMA tile has data. Only warp 0 - * computes softmax for row 0. - * T>1 prefill: rows 0..T-1 have data. All 4 softmax warps process - * rows in parallel (warp w handles rows [w*32, (w+1)*32) ∩ [0, T)). - * This is Milestone 4 territory — current implementation handles T=1 only. - * The multi-head grid layout is independent of multi-row softmax and - * can land first. - * - * ================================================================== - * OUTPUT: UN-NORMALIZED O + LSE - * ================================================================== - * The kernel emits un-normalized O and per-row LSE for composition with - * D5 multi-tile KV merge. External code normalizes: O_norm = O / row_sum. - * For single-segment decode, normalization is done in the epilogue. + * For single-segment decode: normalized O written to GMEM + LSE. + * For multi-segment: un-normalized O + LSE for external merge. * LSE layout: [batch, n_h, T] — one float per head per row. */ @@ -55,6 +40,7 @@ #include "fmha_common.cuh" #include "fmha_umma_desc.cuh" +#include "fmha_tma.cuh" namespace dsv4::kernels::attention { @@ -79,15 +65,19 @@ struct FmhaParams { int q_batch_stride; // stride between Q batch items = n_h * T * hd int k_head_stride; // stride between K heads = N * hd (0 for MQA) int k_batch_stride; // stride between K batch items = n_kv * N * hd - int v_head_stride; // stride between V heads = hd * N (0 for MQA) + int v_head_stride; // stride between V heads = hd * N (or 0 for MQA) int v_batch_stride; // stride between V batch items = n_kv * hd * N int o_head_stride; // stride between O heads = T * hd int o_batch_stride; // stride between O batch items = n_h * T * hd int lse_head_stride; // stride between LSE heads = T int lse_batch_stride; // stride between LSE batch items = n_h * T + + // TMA descriptor for O output (device pointer). When nullptr, uses + // direct GMEM write. When set, uses TMA store for async output. + CUtensorMap* __restrict__ tma_o; }; -template +template __global__ void __launch_bounds__(192) fmha_6warp_multihead_kernel(FmhaParams params) { static constexpr int NKT_QK = HD / MMA_K_BF16; @@ -96,6 +86,7 @@ fmha_6warp_multihead_kernel(FmhaParams params) { static constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16 static constexpr int V_SUB_SZ = 256; // (16,16) canonical BF16 static constexpr int TMEM_N = (HD <= 128) ? 128 : 256; + static constexpr int CORES_MN = 128 / 8; // 16 const int head_idx = blockIdx.y; const int batch_idx = blockIdx.z; @@ -134,6 +125,18 @@ fmha_6warp_multihead_kernel(FmhaParams params) { // ================================================================ // SMEM allocation (shared across all warps) // ================================================================ + // Layout: + // [0..3] sTmemBase (4 bytes, written by tcgen05.alloc) + // [4..7] sRowMax (4 bytes, float) + // [8..11] sRowSum (4 bytes, float) + // [12..15] alignment padding + // [16..16+TILE_SZ*2) sQ0 (4KB, 128×16 canonical BF16) + // [sQ0+TILE_SZ*2..) sK0 (4KB, 128×16 canonical BF16) + // [sK0+TILE_SZ*2..) sPk (4KB, 128B aligned) + // [sPk+TILE_SZ*2..) sV (512B, 16×16 canonical BF16, 128B aligned) + // [sV+V_SUB_SZ*2..) s_p_vals (SK_TILE*4 = 512B) + // [s_p_vals+SK_TILE*4..) sO_epi (HD*2 bytes, row-major BF16, 128B aligned) + // [sO_epi+HD*2..) sMbarStore (16 bytes, 128B aligned) extern __shared__ char sbuf[]; uint32_t* sTmemBase = (uint32_t*)sbuf; float* sRowMax = (float*)(sbuf + 4); @@ -143,6 +146,10 @@ fmha_6warp_multihead_kernel(FmhaParams params) { bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + TILE_SZ) + 127) & ~(uintptr_t)127); bf16_t* sV = (bf16_t*)(((uintptr_t)(sPk + TILE_SZ) + 127) & ~(uintptr_t)127); float* s_p_vals = (float*)(sV + V_SUB_SZ); + // Epilogue SMEM: row-major O tile for TMA store + bf16_t* sO_epi = (bf16_t*)(((uintptr_t)(s_p_vals + SK_TILE) + 127) & ~(uintptr_t)127); + // TMA store mbarrier (16 bytes, 128B aligned) + uint64_t* sMbarStore = (uint64_t*)(((uintptr_t)(sO_epi + HD) + 127) & ~(uintptr_t)127); // ================================================================ // TMEM allocation (warp 4) @@ -160,25 +167,22 @@ fmha_6warp_multihead_kernel(FmhaParams params) { for (int kt = 0; kt < NKT_QK; kt++) { // ---- Warp 5: Load Q and K for this K-tile ---- if (is_load_warp) { - // Load Q K-tile: Q is (1, hd) for decode, row 0 only for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; for (int d = lane; d < MMA_K_BF16; d += 32) { int ck = d / 8, lc = d % 8; - sQ0[ck * 16 * 64 + lc] = q_head[kt * MMA_K_BF16 + d]; + sQ0[ck * CORES_MN * 64 + lc] = q_head[kt * MMA_K_BF16 + d]; } - // Load K K-tile: K is (s_k, hd) for (int i = lane; i < TILE_SZ; i += 32) sK0[i] = 0; for (int r = 0; r < s_k; r++) { for (int d = lane; d < MMA_K_BF16; d += 32) { int ck = d / 8, lc = d % 8; int tmn = r / 8, lr = r % 8; - sK0[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k_head[r * HD + kt * MMA_K_BF16 + d]; + sK0[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = k_head[r * HD + kt * MMA_K_BF16 + d]; } } } __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); @@ -227,16 +231,13 @@ fmha_6warp_multihead_kernel(FmhaParams params) { int d_base = n * 16; for (int kt = 0; kt < NKT_PV; kt++) { - // ---- Warp 5: Fill sPk and load V sub-tile ---- if (is_load_warp) { - // Fill sPk from s_p_vals 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 * 16 * 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[kt * MMA_K_BF16 + c]); } - // Load V sub-tile: V is (hd, s_k) in GMEM for (int i = lane; i < V_SUB_SZ; i += 32) sV[i] = 0; for (int dd = lane; dd < 16; dd += 32) { for (int lr = 0; lr < MMA_K_BF16; lr++) { @@ -249,7 +250,6 @@ fmha_6warp_multihead_kernel(FmhaParams params) { } __syncthreads(); - // ---- Warp 4: PV 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); @@ -262,13 +262,15 @@ fmha_6warp_multihead_kernel(FmhaParams params) { } // ================================================================ - // Epilogue: TMEM → regs → normalize → BF16 → GMEM - // For single-segment decode: normalize in-kernel. - // For multi-segment: emit un-normalized O + LSE. + // EPILOGUE: One-way TMEM → regs → epilogue_op → SMEM → TMA store → GMEM + // ================================================================ + // Step 1: TMEM → registers (warp 0, warp-collective tcgen05.ld) + // Step 2: epilogue_op in registers (normalize + optional FP4 pack) + // Step 3: Registers → SMEM (row-major, matching TMA tile format) + // Step 4: TMA store SMEM → GMEM (or direct GMEM write if no TMA desc) // ================================================================ if (wid == 0) { - float row_max = *sRowMax; - float row_sum = *sRowSum; + // Step 1: TMEM → registers float o_vals[HD]; for (int n = 0; n < HD / 8; n++) { float tmp[8]; @@ -279,20 +281,84 @@ fmha_6warp_multihead_kernel(FmhaParams params) { asm volatile("tcgen05.wait::ld.sync.aligned;"); if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c]; } - // P was NORMALIZED in softmax step. PV = P @ V is already the normalized - // attention output. No further division by row_sum needed. - // For single-segment decode, write O directly. - // LSE is written for multi-segment merge (P5). + + // Step 2: epilogue_op + // P was NORMALIZED in softmax step. PV = P @ V is already normalized. + // Default epilogue_op = identity (o_vals unchanged). + // + // FP4 quantization hook (ENABLE_FP4_EPILOGUE): + // 1. Find amax across o_vals[0..HD-1] + // 2. scale = amax / 6.0 (NVFP4 E2M1 max representable = 6.0) + // 3. For each d: e2m1_val = clamp(o_vals[d] / scale, 0, 6.0) → quantize + // 4. Pack 2 E2M1 values per byte, write scale as FP8 E4M3 + // Output would be (HD/2 bytes FP4 + 1 FP8 scale per 16-element block) + // instead of (HD * 2 bytes BF16). + // For now, BF16 output path — just cast to BF16. + + // Step 3: Registers → SMEM (row-major layout for TMA store) + // All 32 lanes participate. Lane i writes o_vals at offset i*4..i*4+3 + // (matching the TMEM lane mapping for row 0). + // But only lane 0 has valid o_vals. For T=1 decode, this is correct + // because only row 0 of the MMA tile has data. if (lane == 0) { for (int d = 0; d < HD; d++) { - o_head[d] = f32_to_bf16(o_vals[d]); + sO_epi[d] = f32_to_bf16(o_vals[d]); } - // Write LSE if pointer is valid - if (lse_head) { - // LSE = ln(row_sum) + row_max (natural log) - // This is the log of the softmax denominator, useful for - // multi-segment merge: O = sum(exp(lse_i) * O_i) / sum(exp(lse_i)) - lse_head[0] = logf(row_sum) + row_max; + } + // SMEM fence: make writes visible to TMA store + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + + // Write LSE + if (lane == 0 && lse_head) { + float row_max = *sRowMax; + float row_sum = *sRowSum; + lse_head[0] = logf(row_sum) + row_max; + } + } + __syncthreads(); + + // Step 4: TMA store SMEM → GMEM (or direct GMEM write) + if (params.tma_o) { + // Proper TMA store path — async, enables multi-CTA grids. + // One thread initializes the mbarrier and issues the TMA store. + if (tid == 0) { + uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore); + tma_mbarrier_init(mbar_addr, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + __syncthreads(); + + // The O tensor is [batch, n_h, T, HD]. Each head's output starts + // at a different GMEM offset. The TMA descriptor covers the full + // O tensor. We index by (head, batch) to find the right descriptor. + // + // TMA coords for this head's row 0: (x=0, y=0) + // The descriptor was created with the head's GMEM base pointer, + // so coords are relative to that head's start. + if (tid == 0) { + uint32_t smem_addr = (uint32_t)__cvta_generic_to_shared(sO_epi); + uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore); + // Get this (head, batch) TMA descriptor + CUtensorMap* my_tma = params.tma_o + batch_idx * gridDim.y + head_idx; + uint64_t tma_desc = (uint64_t)my_tma; + tma_store_2d(smem_addr, tma_desc, mbar_addr, 0, 0); + // TMA tile: (1, HD) BF16 = HD*2 bytes + tma_mbarrier_arrive_expect_tx(mbar_addr, HD * sizeof(bf16_t)); + } + __syncthreads(); + + // Wait for TMA store completion + if (tid == 0) { + uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore); + tma_store_wait(mbar_addr, 0); + } + __syncthreads(); + } else { + // Fallback: direct GMEM write from SMEM (backward compatible) + // This path is used when no TMA descriptor is provided. + if (wid == 0 && lane == 0) { + for (int d = 0; d < HD; d++) { + o_head[d] = sO_epi[d]; } } } diff --git a/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh b/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh index 26bcfd17..b4afe40f 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh @@ -39,6 +39,8 @@ struct FmhaTmaMultiRowMultiTileParams { int q_head_stride, q_batch_stride; int o_head_stride, o_batch_stride; int lse_head_stride, lse_batch_stride; + // P6: TMA descriptor for O output. When nullptr, direct GMEM write. + CUtensorMap* __restrict__ tma_o; }; template @@ -103,6 +105,10 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params) float* sRunningSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); float* sTileRowMax = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); float* sTileRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); + // P6: Row-major O epilogue buffer + TMA store mbarrier + off = (off + 127) & ~(size_t)127; + bf16_t* sO_epi_rowmajor = (bf16_t*)(sbuf + off); off += MAX_ROWS * HD_CHUNK * sizeof(bf16_t); + uint64_t* sMbarStore = (uint64_t*)(sbuf + off); off += 16; // Init TMEM + mbarrier (once, shared across hd_chunks) if (is_mma_warp) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N); @@ -329,11 +335,18 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params) __syncthreads(); } // kv_tile loop - // ---- Write chunk to GMEM ---- + // ---- Write chunk to SMEM row-major, then TMA store to GMEM ---- + // P6: One-way epilogue pattern — normalize in registers, + // write to SMEM row-major, then TMA store to GMEM. + // This enables multi-CTA grids and FP4 output fusion. + // + // Each active row writes its normalized output to sO_epi_rowmajor + // in row-major layout. Then a single TMA store issues the async + // write for this hd_chunk's output tile. if (my_row_active) { float inv_rs = 1.0f / sRunningSum[my_row]; for (int d = 0; d < HD_CHUNK; d++) { - o_head[my_row * HD + hd_chunk_start + d] = f32_to_bf16(sOacc[my_row * HD_CHUNK + d] * inv_rs); + sO_epi_rowmajor[my_row * HD_CHUNK + d] = f32_to_bf16(sOacc[my_row * HD_CHUNK + d] * inv_rs); } // LSE: same for all chunks, write once if (lse_head && !lse_written) { @@ -341,6 +354,41 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params) } } if (hd_chunk == 0) lse_written = true; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncthreads(); + + // TMA store sO_epi_rowmajor → GMEM + if (params.tma_o) { + if (tid == 0) { + uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore); + tma_mbarrier_init(mbar_addr, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + __syncthreads(); + + if (tid == 0) { + uint32_t smem_addr = (uint32_t)__cvta_generic_to_shared(sO_epi_rowmajor); + uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore); + CUtensorMap* my_tma = params.tma_o + batch_idx * params.n_h + head_idx; + uint64_t tma_desc = (uint64_t)my_tma; + // TMA coords: x = hd_chunk_start, y = 0 + tma_store_2d(smem_addr, tma_desc, mbar_addr, hd_chunk_start, 0); + tma_mbarrier_arrive_expect_tx(mbar_addr, T * HD_CHUNK * sizeof(bf16_t)); + } + __syncthreads(); + if (tid == 0) { + uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore); + tma_store_wait(mbar_addr, 0); + } + __syncthreads(); + } else { + // Fallback: direct GMEM write from SMEM + if (my_row_active) { + for (int d = 0; d < HD_CHUNK; d++) { + o_head[my_row * HD + hd_chunk_start + d] = sO_epi_rowmajor[my_row * HD_CHUNK + d]; + } + } + } __syncthreads(); } // hd_chunk loop diff --git a/dsv4/kernels/attention/fmha_multihead_capi.cu b/dsv4/kernels/attention/fmha_multihead_capi.cu index 92979192..47bcb442 100644 --- a/dsv4/kernels/attention/fmha_multihead_capi.cu +++ b/dsv4/kernels/attention/fmha_multihead_capi.cu @@ -2,32 +2,48 @@ * DSV4 FMHA Multi-Head — C API for ctypes loading. * * Supports single and multi-KV-tile with FlashAttention-2 online softmax. + * P6: One-way TMEM→regs→SMEM→TMA store epilogue with FP4 hook. */ +#include #include #include +#include #include "fmha_common.cuh" #include "fmha_umma_desc.cuh" +#include "fmha_tma.cuh" #include "fmha_6warp_multihead.cuh" +using namespace dsv4::kernels::attention; + extern "C" { int fmha_compute_smem(int hd) { using namespace dsv4::kernels::attention; constexpr int SK = 128; - constexpr int TILE_SZ = 128 * MMA_K_BF16; + constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16 constexpr int V_SUB_SZ = 256; - // sTmemBase(4) + sRowMax(4) + sRowSum(4) + align(4) + sQ0 + sK0 + sPk + sV + s_p_vals + sOacc + slack - int base = 4 + 4 + 4 + 4; // tmem_base, row_max, row_sum, alignment - int sQ0 = TILE_SZ * 2; - int sK0 = TILE_SZ * 2; - int sPk = TILE_SZ * 2; // 128*16*2 = 4096 (P sub-tile) - int sV = V_SUB_SZ * 2; - int sp_vals = SK * 4; - int sOacc = hd * 4; // float accumulator for 1 row - int total = base + sQ0 + sK0 + sPk + sV + sp_vals + sOacc + 256 + 127; - return total & ~127; + // sTmemBase(4) + sRowMax(4) + sRowSum(4) + align(4) = 16 + // sQ0(TILE_SZ*2) = 4096 + // sK0(TILE_SZ*2) = 4096 + // sPk(TILE_SZ*2 + 127) = 4223 (128B aligned) + // sV(V_SUB_SZ*2 + 127) = 639 (128B aligned) + // s_p_vals(SK*4) = 512 + // sO_epi(hd*2 + 127) (128B aligned, row-major BF16) + // sMbarStore(16 + 127) (128B aligned) + int base = 16; + int sQ0 = TILE_SZ * 2; // 4096 + int sK0 = TILE_SZ * 2; // 4096 + int sPk = TILE_SZ * 2; // 4096 (before 128B alignment) + int sV = V_SUB_SZ * 2; // 512 + int sp = SK * 4; // 512 + int sO = hd * 2; // row-major O (HD BF16) + int sMbar = 16; // TMA store mbarrier + // With 128B alignment between sections + int total = base + sQ0 + sK0 + (sPk + 127) + (sV + 127) + sp + (sO + 127) + (sMbar + 127) + 256; + // Round up to 128B + return (total + 127) & ~127; } int fmha_multihead_decode_launch( @@ -44,8 +60,6 @@ int fmha_multihead_decode_launch( int lse_head_stride, int lse_batch_stride, float scale ) { - using namespace dsv4::kernels::attention; - FmhaParams params; params.q = reinterpret_cast(q_ptr); params.k = reinterpret_cast(k_ptr); @@ -65,6 +79,10 @@ int fmha_multihead_decode_launch( params.o_batch_stride = o_batch_stride; params.lse_head_stride = lse_head_stride; params.lse_batch_stride = lse_batch_stride; + // P6: TMA descriptor for O output. + // When nullptr, epilogue uses direct GMEM write (backward compatible). + // To enable TMA store, create per-(head, batch) descriptors and pass them. + params.tma_o = nullptr; int smem = fmha_compute_smem(hd); dim3 grid(1, n_h, batch); @@ -90,4 +108,92 @@ int fmha_multihead_decode_launch( return 0; } +/** + * Launch with TMA store epilogue enabled. + * Creates per-(head, batch) TMA descriptors for the O output tensor. + */ +int fmha_multihead_decode_tma_launch( + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + void* o_ptr, + void* lse_ptr, + int batch, int n_h, int n_kv, int N, int hd, + int q_head_stride, int q_batch_stride, + int k_head_stride, int k_batch_stride, + int v_head_stride, int v_batch_stride, + int o_head_stride, int o_batch_stride, + int lse_head_stride, int lse_batch_stride, + float scale +) { + // Create TMA descriptors for O output: one per (batch, head) + // Each descriptor covers a (1, HD) BF16 tile at the head's GMEM offset. + size_t desc_count = n_h * batch; + CUtensorMap* d_tma_o; + cudaError_t err = cudaMalloc(&d_tma_o, desc_count * sizeof(CUtensorMap)); + if (err != cudaSuccess) return (int)err; + + for (int b = 0; b < batch; b++) { + for (int h = 0; h < n_h; h++) { + int idx = b * n_h + h; + bf16_t* o_head = (bf16_t*)o_ptr + h * o_head_stride + b * o_batch_stride; + // O tile: (1, hd) — one row per head for T=1 decode + CUtensorMap tma_desc; + bool ok = create_tma_desc_2d_bf16(&tma_desc, o_head, 1, hd, 1, hd); + if (!ok) { + cudaFree(d_tma_o); + return -2; + } + cudaMemcpy(&d_tma_o[idx], &tma_desc, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + } + } + + FmhaParams params; + params.q = reinterpret_cast(q_ptr); + params.k = reinterpret_cast(k_ptr); + params.v = reinterpret_cast(v_ptr); + params.o = reinterpret_cast(o_ptr); + params.lse = reinterpret_cast(lse_ptr); + params.s_k = N; + params.scale = scale; + params.head_dim = hd; + params.q_head_stride = q_head_stride; + params.q_batch_stride = q_batch_stride; + params.k_head_stride = k_head_stride; + params.k_batch_stride = k_batch_stride; + params.v_head_stride = v_head_stride; + params.v_batch_stride = v_batch_stride; + params.o_head_stride = o_head_stride; + params.o_batch_stride = o_batch_stride; + params.lse_head_stride = lse_head_stride; + params.lse_batch_stride = lse_batch_stride; + params.tma_o = d_tma_o; + + int smem = fmha_compute_smem(hd); + dim3 grid(1, n_h, batch); + dim3 block(NTHREADS); + + if (smem > 48 * 1024) { + if (hd == 64) cudaFuncSetAttribute(fmha_6warp_multihead_kernel<64, 128>, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + else if (hd == 128) cudaFuncSetAttribute(fmha_6warp_multihead_kernel<128, 128>, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + else if (hd == 256) cudaFuncSetAttribute(fmha_6warp_multihead_kernel<256, 128>, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + + if (hd == 64) fmha_6warp_multihead_kernel<64, 128><<>>(params); + else if (hd == 128) fmha_6warp_multihead_kernel<128, 128><<>>(params); + else if (hd == 256) fmha_6warp_multihead_kernel<256, 128><<>>(params); + else { + cudaFree(d_tma_o); + return -1; + } + + err = cudaGetLastError(); + cudaFree(d_tma_o); + if (err != cudaSuccess) return (int)err; + return 0; +} + } // extern "C" diff --git a/dsv4/kernels/attention/fmha_multitile_capi.cu b/dsv4/kernels/attention/fmha_multitile_capi.cu index 24b57010..676f7655 100644 --- a/dsv4/kernels/attention/fmha_multitile_capi.cu +++ b/dsv4/kernels/attention/fmha_multitile_capi.cu @@ -80,6 +80,7 @@ int fmha_multitile_decode_launch( params.o_batch_stride = o_batch_stride; params.lse_head_stride = lse_head_stride; params.lse_batch_stride = lse_batch_stride; + params.tma_o = nullptr; // P6: TMA O descriptor (future: pass from caller) // SMEM size (match kernel layout) constexpr int HD_CHUNK = 256; @@ -105,6 +106,10 @@ int fmha_multitile_decode_launch( off += 128 * 4; // sRunningSum off += 128 * 4; // sTileRowMax off += 128 * 4; // sTileRowSum + // P6: sO_epi_rowmajor + sMbarStore + off = (off+127)&~(size_t)127; + off += 128 * hc * 2; // sO_epi_rowmajor (MAX_ROWS * HD_CHUNK BF16) + off += 16; // sMbarStore off += 256; // slack int smem = (int)((off + 127) & ~(size_t)127); diff --git a/dsv4/kernels/attention/fmha_tma.cuh b/dsv4/kernels/attention/fmha_tma.cuh index 3f139e78..d83aa89d 100644 --- a/dsv4/kernels/attention/fmha_tma.cuh +++ b/dsv4/kernels/attention/fmha_tma.cuh @@ -186,4 +186,58 @@ struct FmhaTmaDescriptors { CUtensorMap* __restrict__ tma_v; // V descriptor: (HD, s_k) }; +// ================================================================== +// TMA store (GMEM write) operations +// ================================================================== +// The one-way epilogue writes O from SMEM to GMEM via TMA store. +// This replaces direct GMEM writes and enables multi-CTA grids. +// ================================================================== + +/** + * Issue a 2D TMA async copy from SMEM to GMEM (store). + * + * @param smem_src SMEM source address (via __cvta_generic_to_shared) + * @param tma_desc Pointer to CUtensorMap in device memory (uint64_t cast) + * @param smem_mbar SMEM mbarrier address (via __cvta_generic_to_shared) + * @param coord_x Column coordinate (innermost dimension) + * @param coord_y Row coordinate (outermost dimension) + */ +__device__ __forceinline__ void tma_store_2d( + uint32_t smem_src, + uint64_t tma_desc, + uint32_t smem_mbar, + int coord_x, + int coord_y +) { + asm volatile( + "cp.async.bulk.tensor.2d.global.shared::cluster.mbarrier::complete_tx::bytes " + "[%0, {%3, %4}], [%1], [%2];" + :: "l"(tma_desc), + "r"(smem_src), + "r"(smem_mbar), + "r"(coord_x), + "r"(coord_y) + : "memory" + ); +} + +/** + * Wait for TMA store completion using mbarrier try_wait. + * Same pattern as tma_mbarrier_wait but for store mbarriers. + */ +__device__ __forceinline__ void tma_store_wait(uint32_t smem_mbar, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}" + :: "r"(smem_mbar), "r"(phase), "r"(0x989680) + : "memory" + ); +} + } // namespace dsv4::kernels::attention