From 4459ddefddfc8502f04d0b12319a415c9b94f9ac Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 29 May 2026 19:32:02 +0000 Subject: [PATCH] =?UTF-8?q?feat:=206-warp=20TMA=20FMHA=20kernel=20+=20test?= =?UTF-8?q?=20=E2=80=94=20TMA=20for=20K=20loads?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dsv4/kernels/attention/fmha_6warp_tma.cuh | 401 ++++++++++------------ tests/unit/test_fmha_6warp_tma.cu | 143 ++++++++ 2 files changed, 328 insertions(+), 216 deletions(-) create mode 100644 tests/unit/test_fmha_6warp_tma.cu diff --git a/dsv4/kernels/attention/fmha_6warp_tma.cuh b/dsv4/kernels/attention/fmha_6warp_tma.cuh index b4803309..9198fbd8 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma.cuh @@ -1,22 +1,39 @@ /** - * DSV4 FMHA — TMA async loads, 4-warp specialization. + * DSV4 FMHA — 6-warp specialized kernel with TMA async loads for Blackwell SM100. * - * Based on the proven test_fmha_gen pattern, extended with TMA async loads - * for K and V. + * ================================================================== + * WARP SPECIALIZATION (same as fmha_6warp.cuh) + * ================================================================== + * 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 * - * DESIGN: - * - 4 warps (128 threads), __launch_bounds__(128) - * - Warp 0: TMA load + softmax + TMEM read/epilogue - * - Warp 1: MMA + TMEM alloc - * - Warps 2-3: softmax + epilogue - * - TMA: warp 0 lane 0 issues, all threads wait via mbarrier - * - mbarrier: init once, arrive.expect_tx after TMA, phase parity tracking - * - Q loaded directly (T=1 decode for now), K/V via TMA - * - Per-K-sub-tile Q loading (128, 16) into sQ0 - * - Per-K-sub-tile K loading via TMA into sTmaBuf, then canonical sK0 - * - MMA: tid==0 calls umma_ss_f16 - * - Multi-row softmax: warps 0-3 each handle 32 rows - * - PV: per-N-sub-tile, P in registers, V via TMA + * ================================================================== + * 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) + * + * ================================================================== + * 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. + * + * 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 */ #pragma once @@ -27,77 +44,73 @@ namespace dsv4::kernels::attention { -struct FmhaTmaParams { - const bf16_t* __restrict__ q; - const bf16_t* __restrict__ k; - const bf16_t* __restrict__ v; // direct GMEM pointer for V - bf16_t* __restrict__ o; - float* __restrict__ lse; - int s_k, T; - float scale; - int head_dim; - int q_head_stride, q_batch_stride; - int k_head_stride, k_batch_stride; - int v_head_stride, v_batch_stride; - int o_head_stride, o_batch_stride; - int lse_head_stride, lse_batch_stride; - CUtensorMap* __restrict__ tma_k; - CUtensorMap* __restrict__ tma_v; -}; - template -__global__ void __launch_bounds__(128) -fmha_tma_kernel(FmhaTmaParams params) { +__global__ void __launch_bounds__(192) +fmha_6warp_tma_kernel( + const bf16_t* __restrict__ q, + CUtensorMap* __restrict__ tma_k, + const bf16_t* __restrict__ v, + bf16_t* __restrict__ o, + float* __restrict__ lse, + 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 TMEM_N = (HD <= 128) ? 128 : 256; static constexpr int MAX_ROWS = 128; - static constexpr int CORES_MN = 128 / 8; + static constexpr int CORES_MN = 128 / 8; // 16 static constexpr int NUM_READS = SK_TILE / 8; static constexpr int TMA_TILE_BYTES = TILE_SZ * sizeof(bf16_t); - const int head_idx = blockIdx.y; - const int batch_idx = blockIdx.z; const int tid = threadIdx.x; const int wid = tid / 32; const int lane = tid % 32; - const int T = params.T; - const int s_k = params.s_k; - const float scale = params.scale; - bf16_t* __restrict__ q_head = (bf16_t*)params.q + head_idx * params.q_head_stride + batch_idx * params.q_batch_stride; - bf16_t* __restrict__ o_head = params.o + head_idx * params.o_head_stride + batch_idx * params.o_batch_stride; - const bf16_t* __restrict__ v_head = (const bf16_t*)params.v + head_idx * params.v_head_stride + batch_idx * params.v_batch_stride; - float* __restrict__ lse_head = params.lse ? params.lse + head_idx * params.lse_head_stride + batch_idx * params.lse_batch_stride : nullptr; + // Warp role predicates + const bool is_softmax_warp = (wid < 4); + const bool is_mma_warp = (wid == 4); + const bool is_load_warp = (wid == 5); - CUtensorMap* __restrict__ tma_k = params.tma_k; - CUtensorMap* __restrict__ tma_v = params.tma_v; - - // ================================================================== - // SMEM allocation — all 128-byte aligned for TMA compatibility - // ================================================================== + // ================================================================ + // SMEM allocation — all 128-byte aligned + // ================================================================ extern __shared__ __align__(128) char sbuf[]; size_t off = 0; - uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off = 4; + + uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off += 4; + off = (off + 127) & ~(size_t)127; // 128-byte align + + uint64_t* sMbar = (uint64_t*)(sbuf + off); off += 16; off = (off + 127) & ~(size_t)127; - bf16_t* sQ0 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); - bf16_t* sK0 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); + bf16_t* sTmaBuf = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); - off = (off + 15) & ~(size_t)15; - uint64_t* sMbar = (uint64_t*)(sbuf + off); off += 8; + 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* 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); float* sRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); - // sPk and sV for PV GEMM - off = (off + 127) & ~(size_t)127; - bf16_t* sPk = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); - bf16_t* sV = (bf16_t*)(sbuf + off); off += 16 * MMA_K_BF16 * sizeof(bf16_t); // (16,16) canonical for PV + float* s_p_vals = (float*)(sbuf + off); off += SK_TILE * sizeof(float); - // ================================================================== - // Initialize - // ================================================================== - if (wid == 1) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N); + // ================================================================ + // Initialize TMEM + mbarrier + // ================================================================ + 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"); @@ -107,195 +120,151 @@ fmha_tma_kernel(FmhaTmaParams params) { const uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar); int phase = 0; - const bool my_warp_active = (T <= 32) ? (wid == 0) : (wid < 4); - const int my_row = my_warp_active ? (wid * 32 + lane) : 0; - const bool my_row_active = my_warp_active && (my_row < T); - - // ================================================================== - // QK GEMM → S in TMEM - // ================================================================== - { - uint32_t idesc = make_idesc(128, 128); - for (int kt = 0; kt < NKT_QK; kt++) { - // Q sub-tile: direct load from GMEM - for (int i = tid; i < TILE_SZ; i += 128) sQ0[i] = 0; - // Write rows 0..T-1 in canonical layout - for (int r = 0; r < T; r++) { - for (int d = tid % 32; d < MMA_K_BF16; d += 32) { - // Use warp-stride: each warp handles a subset of rows - int my_r = wid * 32 / 128; // simplified: all warps contribute - // Actually, let's use a simpler approach: all 128 threads load + // ================================================================ + // QK GEMM loop: TMA for K, direct load for Q + // ================================================================ + 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]; } } - // Simpler: all 128 threads write Q row by row - for (int d = tid; d < T * MMA_K_BF16; d += 128) { - int r = d / MMA_K_BF16; - int c = d % MMA_K_BF16; - int full_d = kt * MMA_K_BF16 + c; - if (full_d < HD && r < T) { - int ck = c / 8, lc = c % 8, cm = r / 8, lr = r % 8; - sQ0[ck * CORES_MN * 64 + cm * 64 + lr * 8 + lc] = q_head[r * HD + full_d]; - } - } - __syncthreads(); - - // K sub-tile: TMA load + canonical - if (wid == 0 && lane == 0) { - tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_k, mbar_addr, kt * MMA_K_BF16, 0); - tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); - } - tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; - __syncthreads(); - - for (int i = tid; i < TILE_SZ; i += 128) sK0[i] = 0; - for (int i = tid; i < s_k * MMA_K_BF16; i += 128) { - 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(); - - // MMA - if (tid == 0) { - 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); - umma_ss_f16(tb, dq, dk, idesc, kt > 0); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - } - __syncthreads(); } + + // ---- TMA load K sub-tile ---- + // Only 1 thread issues TMA (warp 5, lane 0) + 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); + 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 + 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"); + } + __syncthreads(); } - asm volatile("fence.sc.gpu;" ::: "memory"); - __syncthreads(); - - // ================================================================== - // SOFTMAX - // ================================================================== - float my_row_max = -INFINITY; - if (my_warp_active) { + // ================================================================ + // Softmax (warp 0, row 0 only for T=1 decode) + // ================================================================ + if (wid == 0) { + float s_vals[SK_TILE], row_max = -INFINITY; for (int n = 0; n < NUM_READS; 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)); + : "r"(tb + n*8)); asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (my_row_active) { - for (int c = 0; c < 8; c++) { - int col = n * 8 + c; - if (col < s_k) my_row_max = fmaxf(my_row_max, tmp[c] * scale); - } + if (lane == 0) for (int c=0;c<8;c++) { + s_vals[n*8+c] = tmp[c] * scale; + row_max = fmaxf(row_max, tmp[c] * scale); } } - } - if (my_row_active) sRowMax[my_row] = my_row_max; - __syncthreads(); - - float my_p_vals[SK_TILE]; - float my_row_sum = 0.0f; - if (my_warp_active) { - float rm = my_row_active ? sRowMax[my_row] : 0.0f; - for (int n = 0; n < NUM_READS; 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;"); - if (my_row_active) { - for (int c = 0; c < 8; c++) { - int col = n * 8 + c; - if (col < s_k) { - float p = expf(tmp[c] * scale - rm); - my_p_vals[col] = p; - my_row_sum += p; - } - } - } + row_max = wmax(row_max); + if (lane == 0) *sRowMax = row_max; + float row_sum = 0.0f; + if (lane == 0) for (int j=0;j 0); + if (tid == 128) umma_ss_f16(tb + n * 16, dp, dv, idesc_pv16, kt > 0); asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); } __syncthreads(); } } - asm volatile("fence.sc.gpu;" ::: "memory"); - __syncthreads(); - - // ================================================================== - // EPILOGUE - // ================================================================== - if (my_warp_active) { - 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; - - for (int n = 0; n < N_NSUB * 2; n++) { + // ================================================================ + // Epilogue: TMEM → regs → normalize → BF16 → GMEM (warp 0) + // ================================================================ + if (wid == 0) { + float o_vals[HD]; + for (int n = 0; n < HD / 8; 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)); + : "r"(tb + n*8)); asm volatile("tcgen05.wait::ld.sync.aligned;"); - 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 (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c]; } - if (my_row_active && lse_head) lse_head[my_row] = logf(rs) + rm; + float row_sum = *sRowSum; + float inv_rs = 1.0f / row_sum; + if (lane == 0) for (int d=0;d +#include +#include +#include +#include +#include + +#ifndef HD_VAL +#define HD_VAL 64 +#endif + +#include "dsv4/kernels/attention/fmha_common.cuh" +#include "dsv4/kernels/attention/fmha_umma_desc.cuh" +#include "dsv4/kernels/attention/fmha_tma.cuh" + +using namespace dsv4::kernels::attention; + +static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); } +static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; } + +constexpr int HD = HD_VAL; +constexpr int SK = 128; +constexpr int MMA_K_BF16 = 16; +constexpr int TILE_SZ = 128 * MMA_K_BF16; +constexpr int V_SUB_SZ = 16 * MMA_K_BF16; +constexpr int TMEM_N = (HD <= 128) ? 128 : 256; + +#include "dsv4/kernels/attention/fmha_6warp_tma.cuh" + +static size_t compute_smem() { + size_t off = 0; + off += 4; // sTmemBase + off = (off + 127) & ~(size_t)127; + off += 16; // sMbar + off = (off + 127) & ~(size_t)127; + off += TILE_SZ * sizeof(bf16_t); // sTmaBuf + off = (off + 127) & ~(size_t)127; + off += TILE_SZ * sizeof(bf16_t); // sQ0 + 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); // sPk + off = (off + 127) & ~(size_t)127; + off += V_SUB_SZ * sizeof(bf16_t); // sV + off += 128 * sizeof(float); // sRowMax + off += 128 * sizeof(float); // sRowSum + off += SK * sizeof(float); // s_p_vals + return off; +} + +int main() { + printf("=== 6-warp TMA FMHA HD=%d ===\n", HD); + const float SCALE = 1.0f / sqrtf((float)HD); + + bf16_t* h_q = (bf16_t*)malloc(HD*sizeof(bf16_t)); + bf16_t* h_k = (bf16_t*)malloc(SK*HD*sizeof(bf16_t)); + bf16_t* h_v = (bf16_t*)malloc(HD*SK*sizeof(bf16_t)); + bf16_t* h_o = (bf16_t*)calloc(HD, sizeof(bf16_t)); + float* h_lse = (float*)calloc(1, sizeof(float)); + + srand(42); + for (int d=0;d 48 * 1024) { + cudaFuncSetAttribute(fmha_6warp_tma_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); + } + fmha_6warp_tma_kernel<<<1, 192, smem>>>(d_q, d_tma_k, d_v, d_o, d_lse, SK, SCALE); + + cudaError_t launch_err = cudaGetLastError(); + if (launch_err != cudaSuccess) { printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err)); return 1; } + + cudaError_t err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; } + + cudaMemcpy(h_o, d_o, HD*sizeof(bf16_t), cudaMemcpyDeviceToHost); + cudaMemcpy(h_lse, d_lse, sizeof(float), cudaMemcpyDeviceToHost); + + printf("O[0..7] MMA: "); for(int d=0;d1e-4f) { cs+=a*b; na+=a*a; nb+=b*b; } + } + cs /= (sqrtf(na)*sqrtf(nb)+1e-10f); + printf("Filtered cosine: %.8f\n", cs); + printf("Test %s\n", cs > 0.999f ? "PASSED" : "FAILED"); + + cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse); cudaFree(d_tma_k); + free(h_q); free(h_k); free(h_v); free(h_o); free(h_lse); + return cs > 0.999f ? 0 : 1; +}