From 95725f1df00f1fc7929883ce1e1679339aa0af49 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 17:21:15 +0000 Subject: [PATCH] P8: Delete 6 redundant .cuh variants + multihead CAPI/op Kept: fmha_6warp_tma_multirow_multitile.cuh (production kernel) Deleted: fmha_6warp.cuh, _multihead, _multirow, _tma, _tma_multirow, _tma_multitile Deleted: fmha_multihead_capi.cu, fmha_multihead_op.py production.py: Removed _dsv4_attention_fast_decode, unified dispatch to _dsv4_attention_multitile for all fast-path cases. --- dsv4/kernels/attention/fmha_6warp.cuh | 234 ------------ .../attention/fmha_6warp_multihead.cuh | 332 ----------------- .../kernels/attention/fmha_6warp_multirow.cuh | 309 ---------------- dsv4/kernels/attention/fmha_6warp_tma.cuh | 303 ---------------- .../attention/fmha_6warp_tma_multirow.cuh | 336 ------------------ .../attention/fmha_6warp_tma_multitile.cuh | 315 ---------------- dsv4/kernels/attention/fmha_multihead_capi.cu | 105 ------ dsv4/kernels/attention/fmha_multihead_op.py | 267 -------------- dsv4/kernels/attention/production.py | 51 --- 9 files changed, 2252 deletions(-) delete mode 100644 dsv4/kernels/attention/fmha_6warp.cuh delete mode 100644 dsv4/kernels/attention/fmha_6warp_multihead.cuh delete mode 100644 dsv4/kernels/attention/fmha_6warp_multirow.cuh delete mode 100644 dsv4/kernels/attention/fmha_6warp_tma.cuh delete mode 100644 dsv4/kernels/attention/fmha_6warp_tma_multirow.cuh delete mode 100644 dsv4/kernels/attention/fmha_6warp_tma_multitile.cuh delete mode 100644 dsv4/kernels/attention/fmha_multihead_capi.cu delete mode 100644 dsv4/kernels/attention/fmha_multihead_op.py diff --git a/dsv4/kernels/attention/fmha_6warp.cuh b/dsv4/kernels/attention/fmha_6warp.cuh deleted file mode 100644 index 7740d84b..00000000 --- a/dsv4/kernels/attention/fmha_6warp.cuh +++ /dev/null @@ -1,234 +0,0 @@ -/** - * DSV4 FMHA — 6-warp specialized kernel for Blackwell SM100. - * - * ================================================================== - * WARP SPECIALIZATION - * ================================================================== - * Warp 0-3 (tid 0-127): Softmax + correction + epilogue - * - Read S from TMEM, compute softmax, write P to SMEM - * - After PV: read O from TMEM, normalize, write to GMEM - * - For T=1 decode: only warp 0 processes row 0 - * - * Warp 4 (tid 128-159): MMA (QK + PV) - * - Call tcgen05.mma for QK and PV - * - TMEM alloc/dealloc - * - Only 1 thread calls MMA, but TMEM ops are warp-collective - * - * Warp 5 (tid 160-191): Data staging (Q/K/V loads) - * - Load Q, K, V from GMEM to SMEM in canonical layout - * - Future: TMA loads with mbarrier - * - Fill sPk from s_p_vals - * - * ================================================================== - * SYNCHRONIZATION - * ================================================================== - * CTA-wide __syncthreads() barriers between phases: - * 1. After Q/K/V loads → QK MMA - * 2. After QK MMA → softmax - * 3. After softmax + P fill + V load → PV MMA - * 4. After PV MMA → epilogue - * - * Future: mbarrier-based producer-consumer sync between warp 5 (producer) - * and warp 4 (consumer) for pipeline overlap. - * - * ================================================================== - * SMEM LAYOUT (shared across all warps) - * ================================================================== - * sQ: (128, 16) canonical = 8 KB (1 K-tile, reused) - * sK: (128, 16) canonical = 8 KB (1 K-tile, reused) - * sPk: (128, 16) canonical = 8 KB (1 sub-tile, reused) - * sV: (16, 16) canonical = 512 bytes (1 N-sub-tile) - * s_p_vals: 128 floats = 512 bytes (softmax output) - * sRowMax: 1 float (row 0 max, for T=1) - * sRowSum: 1 float (row 0 sum, for T=1) - * sTmemBase: 4 bytes (TMEM allocation) - * Total: ~26 KB - */ - -#pragma once - -#include "fmha_common.cuh" -#include "fmha_umma_desc.cuh" - -namespace dsv4::kernels::attention { - -template -__global__ void __launch_bounds__(192) -fmha_6warp_kernel( - const bf16_t* __restrict__ q, - const bf16_t* __restrict__ k, - const bf16_t* __restrict__ v, - bf16_t* __restrict__ o, - int s_k, float scale -) { - static constexpr int NKT_QK = HD / MMA_K_BF16; - static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; // 8 - static constexpr int N_NSUB = HD / 16; - 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; - - const int tid = threadIdx.x; - const int wid = tid / 32; - const int lane = tid % 32; - - // Warp role predicates - const bool is_softmax_warp = (wid < 4); // Warps 0-3 - const bool is_mma_warp = (wid == 4); // Warp 4 - const bool is_load_warp = (wid == 5); // Warp 5 - - // ================================================================ - // SMEM allocation (shared across all warps) - // ================================================================ - extern __shared__ char sbuf[]; - uint32_t* sTmemBase = (uint32_t*)sbuf; - float* sRowMax = (float*)(sbuf + 4); - float* sRowSum = sRowMax + 1; - bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + 1) + 15) & ~(uintptr_t)15); - bf16_t* sK0 = sQ0 + TILE_SZ; - 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); - - // ================================================================ - // TMEM allocation (warp 4) - // ================================================================ - if (is_mma_warp) { - uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); - tmem_alloc(smem_ptr, TMEM_N); - } - __syncthreads(); - uint32_t tb = *sTmemBase; - - // ================================================================ - // QK GEMM loop: for each K-tile, load Q+K, then MMA - // ================================================================ - 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 - 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[kt * MMA_K_BF16 + d]; - } - // Load K K-tile - 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[r * HD + kt * MMA_K_BF16 + d]; - } - } - } - __syncthreads(); // Wait for loads - - // ---- 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(); // Wait for MMA - } - - // ================================================================ - // 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 < SK_TILE / 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)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - 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); - } - } - 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); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - } - __syncthreads(); // Wait for MMA - } - } - - // ================================================================ - // 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)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c]; - } - if (lane == 0) for (int d=0;d -__global__ void __launch_bounds__(192) -fmha_6warp_multihead_kernel(FmhaParams params) { - static constexpr int NKT_QK = HD / MMA_K_BF16; - static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; // 8 - static constexpr int N_NSUB = HD / 16; - 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; - const int tid = threadIdx.x; - const int wid = tid / 32; - const int lane = tid % 32; - - // Warp role predicates - const bool is_softmax_warp = (wid < 4); // Warps 0-3 - const bool is_mma_warp = (wid == 4); // Warp 4 - const bool is_load_warp = (wid == 5); // Warp 5 - - // ================================================================== - // Compute per-head GMEM pointers - // ================================================================== - const bf16_t* __restrict__ q_head = params.q - + head_idx * params.q_head_stride - + batch_idx * params.q_batch_stride; - const bf16_t* __restrict__ k_head = params.k - + head_idx * params.k_head_stride - + batch_idx * params.k_batch_stride; - const bf16_t* __restrict__ v_head = params.v - + head_idx * params.v_head_stride - + batch_idx * params.v_batch_stride; - bf16_t* __restrict__ o_head = params.o - + head_idx * params.o_head_stride - + batch_idx * params.o_batch_stride; - float* __restrict__ lse_head = params.lse - ? params.lse + head_idx * params.lse_head_stride - + batch_idx * params.lse_batch_stride - : nullptr; - - const int s_k = params.s_k; - const float scale = params.scale; - - // ================================================================ - // 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); - float* sRowSum = sRowMax + 1; - bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + 1) + 15) & ~(uintptr_t)15); - bf16_t* sK0 = sQ0 + TILE_SZ; - 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 for GMEM write - bf16_t* sO_epi = (bf16_t*)(((uintptr_t)(s_p_vals + SK_TILE) + 127) & ~(uintptr_t)127); - - // ================================================================ - // TMEM allocation (warp 4) - // ================================================================ - if (is_mma_warp) { - uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); - tmem_alloc(smem_ptr, TMEM_N); - } - __syncthreads(); - uint32_t tb = *sTmemBase; - - // ================================================================ - // QK GEMM loop: for each K-tile, load Q+K, then MMA - // ================================================================ - for (int kt = 0; kt < NKT_QK; kt++) { - // ---- Warp 5: Load Q and K for this K-tile ---- - 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 ck = d / 8, lc = d % 8; - sQ0[ck * CORES_MN * 64 + lc] = q_head[kt * MMA_K_BF16 + d]; - } - 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 * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = k_head[r * HD + kt * MMA_K_BF16 + d]; - } - } - } - __syncthreads(); - - 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(); - } - - // ================================================================ - // 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 < SK_TILE / 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)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - 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); - } - } - 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); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - } - __syncthreads(); - } - } - - // ================================================================ - // EPILOGUE: One-way TMEM → regs → epilogue_op → SMEM → 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) - // Step 4: SMEM → GMEM (direct write) - // - // NOTE: cp.async.bulk.tensor store (SMEM → GMEM) is NOT available - // on SM100 for this use case. The CUTLASS SM100 epilogue uses - // st.global directly. For multi-CTA (D2), we'll use st.global - // with flat_divide coordinates. - // ================================================================ - if (wid == 0) { - // Step 1: TMEM → registers (warp-collective) - 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)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c]; - } - - // Step 2: epilogue_op in registers - // P was NORMALIZED in softmax step. PV = P @ V is already normalized. - // Default epilogue_op = identity. - // - // FP4 quantization hook (ENABLE_FP4_EPILOGUE): - // 1. Find amax across o_vals[0..HD-1] - // 2. scale = amax / 6.0 (NVFP4 E2M1 range) - // 3. Quantize each o_vals[d] to E2M1, pack 2 per byte - // 4. Write scale as FP8 E4M3 per 16-element block - // For now, BF16 output: just cast. - - // Step 3: Registers → SMEM (row-major) - // This is the SMEM staging step of the one-way epilogue. - // In the MoE kernel (CuTeDSL), this is done by - // epilogue_smem_copy_and_partition. Here we write directly. - if (lane == 0) { - for (int d = 0; d < HD; d++) { - sO_epi[d] = f32_to_bf16(o_vals[d]); - } - } - 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: SMEM → GMEM (direct write) - // The SMEM buffer is the single source of truth for the output. - // For single-CTA decode, this is correct and efficient. - // For multi-CTA (D2), will use st.global with flat_divide coords. - if (wid == 0 && lane == 0) { - for (int d = 0; d < HD; d++) { - o_head[d] = sO_epi[d]; - } - } - __syncthreads(); - - // TMEM dealloc (warp 4) - if (is_mma_warp) { - tmem_dealloc(tb, TMEM_N); - } -} - -} // namespace dsv4::kernels::attention diff --git a/dsv4/kernels/attention/fmha_6warp_multirow.cuh b/dsv4/kernels/attention/fmha_6warp_multirow.cuh deleted file mode 100644 index 5625f0f0..00000000 --- a/dsv4/kernels/attention/fmha_6warp_multirow.cuh +++ /dev/null @@ -1,309 +0,0 @@ -/** - * DSV4 FMHA — 6-warp specialized kernel, multi-row softmax (prefill T>1). - * - * ================================================================== - * DESIGN - * ================================================================== - * - * 6-warp CTA: warps 0-3 = softmax, warp 4 = MMA, warp 5 = TMA load. - * Grid: (1, n_h, batch) — each CTA processes one head of one batch item. - * - * KEY: TMEM is shared between QK output (S) and PV output (O). - * We CANNOT interleave softmax reads with PV writes. - * Instead: softmax reads ALL of S first (into registers), THEN PV writes O. - * - * After UMMA, 4 warps reading TMEM with 32x32b.x8 each see a different - * 32-row partition (verified on B200): - * Warp 0 → rows 0-31, Warp 1 → rows 32-63, etc. - * Lane l in warp w reads row w*32 + l across 8 columns. - * - * FLOW: - * 1. QK GEMM → S in TMEM - * 2. Softmax: read all S from TMEM, compute P in registers - * - Pass 1: row_max - * - Pass 2: exp(S*scale - row_max), accumulate row_sum, store P in regs - * 3. PV GEMM: write P to sPk per K-tile, PV accumulate into TMEM - * 4. Epilogue: read O from TMEM, normalize, write to GMEM - * - * P in registers: each lane holds float p_vals[SK_TILE] = 512 bytes. - * For T=128, 128 lanes × 512 = 64KB register space (within B200 budget). - * ================================================================== - */ - -#pragma once - -#include "fmha_common.cuh" -#include "fmha_umma_desc.cuh" - -namespace dsv4::kernels::attention { - -struct FmhaMultiRowParams { - const bf16_t* __restrict__ q; - const bf16_t* __restrict__ k; - const bf16_t* __restrict__ v; - bf16_t* __restrict__ o; - float* __restrict__ lse; // [batch, n_h, T] — per-row LSE for multi-tile KV merge - 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; -}; - -template -__global__ void __launch_bounds__(192) -fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { - 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; - 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 NUM_READS = SK_TILE / 8; // 16 groups of 8 columns - - 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 bool is_softmax_warp = (wid < 4); - const bool is_mma_warp = (wid == 4); - const bool is_load_warp = (wid == 5); - const int T = params.T; - const int s_k = params.s_k; - const float scale = params.scale; - - const bf16_t* __restrict__ q_head = params.q + head_idx * params.q_head_stride + batch_idx * params.q_batch_stride; - const bf16_t* __restrict__ k_head = params.k + head_idx * params.k_head_stride + batch_idx * params.k_batch_stride; - const bf16_t* __restrict__ v_head = params.v + head_idx * params.v_head_stride + batch_idx * params.v_batch_stride; - bf16_t* __restrict__ o_head = params.o + head_idx * params.o_head_stride + batch_idx * params.o_batch_stride; - float* __restrict__ lse_head = params.lse ? params.lse + head_idx * params.lse_head_stride + batch_idx * params.lse_batch_stride : nullptr; - - // SMEM allocation - extern __shared__ char sbuf[]; - uint32_t* sTmemBase = (uint32_t*)sbuf; - float* sRowMax = (float*)(sbuf + 4); // sTmemBase is 1 uint32_t = 4 bytes - float* sRowSum = sRowMax + MAX_ROWS; - bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 127) & ~(uintptr_t)127); - bf16_t* sK0 = sQ0 + TILE_SZ; - 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); - - // TMEM alloc - if (is_mma_warp) { - uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); - tmem_alloc(smem_ptr, TMEM_N); - } - __syncthreads(); - uint32_t tb = *sTmemBase; - - // Row assignment - const bool my_warp_active = (T <= 32) ? (wid == 0) : is_softmax_warp; - 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 - // ================================================================ - for (int kt = 0; kt < NKT_QK; kt++) { - if (is_load_warp) { - for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; - for (int r = 0; r < T; r++) { - 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, cm = r/8, lr = r%8; - sQ0[ck*CORES_MN*64 + cm*64 + lr*8 + lc] = q_head[r * HD + full_d]; - } - } - } - 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 full_d = kt * MMA_K_BF16 + d; - if (full_d < HD) { - int ck = d/8, lc = d%8, tmn = r/8, lr = r%8; - sK0[ck*CORES_MN*64 + tmn*64 + lr*8 + lc] = k_head[r * HD + full_d]; - } - } - } - } - __syncthreads(); - 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(); - } - - // TMEM visibility fence - asm volatile("fence.sc.gpu;" ::: "memory"); - __syncthreads(); - - // ================================================================ - // SOFTMAX — compute P in registers (TWO passes over TMEM) - // ================================================================ - - // Pass 1: row_max - float my_row_max = -INFINITY; - if (my_warp_active) { - 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) my_row_max = fmaxf(my_row_max, tmp[c] * scale); - } - } - } - } - - // Store row_max to SMEM for cross-warp visibility - if (my_row_active) sRowMax[my_row] = my_row_max; - __syncthreads(); - - // Pass 2: compute P values (exp, sum, normalize) - float my_p_vals[SK_TILE]; // P for my row, all columns - 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;"); // NOTE: ld not ld.sync — some B200 builds need the dot - 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; - } - } - } - } - } - - // Store row_sum - if (my_row_active) sRowSum[my_row] = my_row_sum; - __syncthreads(); - - // ================================================================ - // PV GEMM — write P to sPk per K-tile, accumulate O in TMEM - // ================================================================ - for (int n_sub = 0; n_sub < N_NSUB; n_sub++) { - int d_base = n_sub * 16; - for (int pv_kt = 0; pv_kt < NKT_PV; pv_kt++) { - const int col_start = pv_kt * MMA_K_BF16; - - // Load warp: zero sPk - if (is_load_warp) { - for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0; - } - __syncthreads(); - - // Softmax warps: write P to sPk - if (my_row_active) { - for (int c = 0; c < MMA_K_BF16; c++) { - int gc = col_start + c; - int ck = c/8, lc = c%8; - int core_mn = my_row/8, local_r = my_row%8; - sPk[ck*CORES_MN*64 + core_mn*64 + local_r*8 + lc] = f32_to_bf16(my_p_vals[gc]); - } - } - __syncthreads(); - - // Load warp: fill sV - if (is_load_warp) { - 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++) { - int r = col_start + lr; - if (r < s_k && (d_base + dd) < HD) { - int g_mn = dd/8, g_k = lr/8, llr = dd%8, lc = lr%8; - sV[g_k*2*64 + g_mn*64 + llr*8 + lc] = v_head[(d_base+dd)*s_k + r]; - } - } - } - } - __syncthreads(); - - // MMA: sPk × sV → TMEM - if (is_mma_warp) { - uint32_t idesc_pv = make_idesc(128, 16); - uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), 128); - uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16); - // accumulate: first K-tile for this n_sub should NOT accumulate - // pv_kt > 0 means we're accumulating across K-tiles for the same n_sub - if (tid == 128) umma_ss_f16(tb + n_sub*16, dp, dv, idesc_pv, pv_kt > 0); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - } - __syncthreads(); - } - } - - // Ensure PV output is visible to all warps before epilogue - asm volatile("fence.sc.gpu;" ::: "memory"); - __syncthreads(); - - // ================================================================ - // EPILOGUE: TMEM → regs → normalize → BF16 → GMEM + LSE output - // - // CRITICAL: TMEM loads (32x32b.x8) are WARP-COLLECTIVE. - // ALL 32 lanes must execute them. The load MUST be outside - // the my_row_active guard. Only the GMEM store is conditional. - // - // Output: normalized O (O_unnorm / row_sum) + per-row LSE. - // LSE = ln(row_sum) + row_max, for multi-tile KV merge: - // O = Σ exp(lse_i - L) * O_i / Σ exp(lse_i - L) - // where L = max(lse_i) for numerical stability. - // ================================================================ - 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; - - // Read O from TMEM: N_NSUB*2 groups of 8 columns - // ALL lanes in the warp must execute the TMEM load (warp-collective) - for (int n = 0; n < N_NSUB * 2; 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;"); - - // Only store to GMEM for active rows - 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 (my_row_active && lse_head) lse_head[my_row] = logf(rs) + rm; - } - __syncthreads(); - - if (is_mma_warp) tmem_dealloc(tb, TMEM_N); -} - -} // namespace dsv4::kernels::attention diff --git a/dsv4/kernels/attention/fmha_6warp_tma.cuh b/dsv4/kernels/attention/fmha_6warp_tma.cuh deleted file mode 100644 index cafd7753..00000000 --- a/dsv4/kernels/attention/fmha_6warp_tma.cuh +++ /dev/null @@ -1,303 +0,0 @@ -/** - * DSV4 FMHA — 6-warp specialized kernel with double-buffer TMA pipeline. - * - * ================================================================== - * DOUBLE-BUFFER TMA PIPELINE - * ================================================================== - * - * 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 - * - * 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 - * - * 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 - -#include "fmha_common.cuh" -#include "fmha_umma_desc.cuh" -#include "fmha_tma.cuh" - -namespace dsv4::kernels::attention { - -template -__global__ void __launch_bounds__(192) -fmha_6warp_tma_kernel( - const bf16_t* __restrict__ q, - CUtensorMap* __restrict__ tma_k, - CUtensorMap* __restrict__ tma_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; - static constexpr int TMEM_N = (HD <= 128) ? 128 : (HD <= 256) ? 256 : 512; - static constexpr int MAX_ROWS = 128; - 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); - - const int tid = threadIdx.x; - const int wid = tid / 32; - const int lane = tid % 32; - - const bool is_softmax_warp = (wid < 4); - const bool is_mma_warp = (wid == 4); - const bool is_load_warp = (wid == 5); - - // ================================================================ - // 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; - 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); - float* sRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); - float* s_p_vals = (float*)(sbuf + off); off += SK_TILE * sizeof(float); - - // ================================================================ - // Init - // ================================================================ - 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"); - } - __syncthreads(); - uint32_t tb = *sTmemBase; - const uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar); - int phase = 0; - - // ================================================================ - // QK GEMM — double-buffer TMA pipeline - // ================================================================ - { - uint32_t idesc = make_idesc(128, 128); - bf16_t* sK_bufs[2] = {sK0, sK1}; - int cur_buf = 0; // alternates 0,1,0,1,... - - // --- 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, 0, 0); - tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); - } - tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; - __syncthreads(); - - // 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]; - } - - // 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(); - } - } - - // ================================================================ - // 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)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - 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); - } - } - 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); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - } - __syncthreads(); - } - } - - // ================================================================ - // Epilogue: TMEM → regs → 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)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c]; - } - if (lane == 0) for (int d=0;d -__global__ void __launch_bounds__(192) -fmha_6warp_tma_multirow_kernel(FmhaTmaMultiRowParams params) { - 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; - static constexpr int TMEM_N = (HD <= 128) ? 128 : (HD <= 256) ? 256 : 512; - static constexpr int MAX_ROWS = 128; - 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); - - 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 bool is_softmax_warp = (wid < 4); - const bool is_mma_warp = (wid == 4); - const bool is_load_warp = (wid == 5); - const int T = params.T; - const int s_k = params.s_k; - const float scale = params.scale; - - const bf16_t* __restrict__ q_head = params.q + head_idx * params.q_head_stride + batch_idx * params.q_batch_stride; - // v_head not used — V loaded via TMA - bf16_t* __restrict__ o_head = params.o + head_idx * params.o_head_stride + batch_idx * params.o_batch_stride; - float* __restrict__ lse_head = params.lse ? params.lse + head_idx * params.lse_head_stride + batch_idx * params.lse_batch_stride : nullptr; - - CUtensorMap* __restrict__ my_tma_k = params.tma_k + batch_idx * params.n_h + head_idx; - CUtensorMap* __restrict__ my_tma_v = params.tma_v + batch_idx * params.n_h + head_idx; - - // ================================================================ - // 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; - 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); - float* sRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); - - // Init - 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"); - } - __syncthreads(); - uint32_t tb = *sTmemBase; - const uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar); - int phase = 0; - - const bool my_warp_active = (T <= 32) ? (wid == 0) : is_softmax_warp; - const int my_row = my_warp_active ? (wid * 32 + lane) : 0; - const bool my_row_active = my_warp_active && (my_row < T); - - // ================================================================ - // QK GEMM — double-buffer TMA pipeline - // ================================================================ - { - bf16_t* sK_bufs[2] = {sK0, sK1}; - int cur_buf = 0; - - // Preload kt=0 - if (is_load_warp && lane == 0) { - tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)my_tma_k, mbar_addr, 0, 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 += 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]; - } - - // Load Q for kt=0 - if (is_load_warp) { - for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; - for (int r = 0; r < T; r++) { - for (int d = lane; d < MMA_K_BF16; d += 32) { - int full_d = d; - if (full_d < HD) { - int ck = d/8, lc = d%8, cm = r/8, lr = r%8; - sQ0[ck*CORES_MN*64 + cm*64 + lr*8 + lc] = q_head[r * HD + full_d]; - } - } - } - } - __syncthreads(); - - // Pipeline loop - for (int kt = 0; kt < NKT_QK; kt++) { - bf16_t* cur_sK = sK_bufs[cur_buf]; - bool has_next = (kt + 1 < NKT_QK); - - // 1. Issue TMA for kt+1 (DMA runs in parallel with MMA) - if (has_next && is_load_warp && lane == 0) { - tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)my_tma_k, - mbar_addr, (kt+1) * MMA_K_BF16, 0); - tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); - } - - // 2. MMA QK - 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(cur_sK), 128); - if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - } - - // 3. Wait for TMA kt+1 - if (has_next) { - tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; - } - __syncthreads(); - - // 4. Convert → next buffer - 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 kt+1 - if (has_next && is_load_warp) { - for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; - for (int r = 0; r < T; r++) { - 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, cm = r/8, lr = r%8; - sQ0[ck*CORES_MN*64 + cm*64 + lr*8 + lc] = q_head[r * HD + full_d]; - } - } - } - } - - cur_buf = 1 - cur_buf; - __syncthreads(); - } - } - - asm volatile("fence.sc.gpu;" ::: "memory"); - __syncthreads(); - - // ================================================================ - // SOFTMAX — 2-pass, P in registers - // ================================================================ - float my_row_max = -INFINITY; - if (my_warp_active) { - 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) my_row_max = fmaxf(my_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; - } - } - } - } - } - if (my_row_active) sRowSum[my_row] = my_row_sum; - __syncthreads(); - - // ================================================================ - // PV GEMM — P→sPk + V TMA → O in TMEM - // ================================================================ - for (int n_sub = 0; n_sub < N_NSUB; n_sub++) { - int d_base = n_sub * 16; - for (int pv_kt = 0; pv_kt < NKT_PV; pv_kt++) { - const int col_start = pv_kt * MMA_K_BF16; - - if (is_load_warp) { - for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0; - } - __syncthreads(); - - if (my_row_active) { - for (int c = 0; c < MMA_K_BF16; c++) { - int gc = col_start + c; - int ck = c/8, lc = c%8; - int core_mn = my_row/8, local_r = my_row%8; - sPk[ck*CORES_MN*64 + core_mn*64 + local_r*8 + lc] = f32_to_bf16(my_p_vals[gc]); - } - } - __syncthreads(); - - // V via TMA - if (is_load_warp && lane == 0) { - tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)my_tma_v, - 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; - __syncthreads(); - - for (int i = tid; i < V_SUB_SZ; i += 192) sV[i] = 0; - for (int i = tid; i < 16 * MMA_K_BF16; i += 192) { - int dd = i / MMA_K_BF16, lr = i % MMA_K_BF16; - int g_mn = dd/8, g_k = lr/8, llr = dd%8, lc = lr%8; - sV[g_k*2*64 + g_mn*64 + llr*8 + lc] = sTmaBuf[i]; - } - __syncthreads(); - - if (is_mma_warp) { - uint32_t idesc_pv = make_idesc(128, 16); - uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), 128); - uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16); - if (tid == 128) umma_ss_f16(tb + n_sub*16, dp, dv, idesc_pv, pv_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++) { - 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 d = n * 8 + c; - if (d < HD) o_head[my_row * HD + d] = f32_to_bf16(tmp[c] * inv_rs); - } - } - } - if (my_row_active && lse_head) lse_head[my_row] = logf(rs) + rm; - } - __syncthreads(); - - if (is_mma_warp) tmem_dealloc(tb, TMEM_N); -} - -} // namespace diff --git a/dsv4/kernels/attention/fmha_6warp_tma_multitile.cuh b/dsv4/kernels/attention/fmha_6warp_tma_multitile.cuh deleted file mode 100644 index 712bb7da..00000000 --- a/dsv4/kernels/attention/fmha_6warp_tma_multitile.cuh +++ /dev/null @@ -1,315 +0,0 @@ -/** - * DSV4 FMHA — 6-warp TMA kernel, multi-tile KV, T=1 decode. - * - * Register-resident O accumulator for in-kernel multi-tile KV merge. - * Avoids the D1.5 TMEM round-trip problem. - * - * ONLY supports T=1 decode. For T>1, use fmha_6warp_tma_multirow - * with Python KV merge. - * - * ================================================================== - * MULTI-TILE KV DESIGN - * ================================================================== - * - * Outer loop: kv_tile = 0..n_kv_tiles-1 - * Inner loop: kt = 0..NKT_QK-1 (K sub-tiles within a KV tile) - * - * For each KV tile: - * 1. TMA load K sub-tile → sTmaBuf → canonical sK - * 2. QK MMA → S in TMEM - * 3. Softmax (with online rescale of running O) - * 4. PV MMA → O_tile in TMEM - * 5. Read O_tile from TMEM → registers - * 6. Rescale running O: o_acc *= exp(running_max - tile_max) - * 7. Add O_tile to o_acc - * 8. Update running_max, running_sum - * - * Final: o_acc /= running_sum → BF16 → GMEM + LSE - * - * Register pressure (T=1): o_acc[HD] = 256-1024 bytes per thread. - * Only lane 0 of warp 0 holds the accumulator. - */ - -#pragma once - -#include "fmha_common.cuh" -#include "fmha_umma_desc.cuh" -#include "fmha_tma.cuh" - -namespace dsv4::kernels::attention { - -struct FmhaTmaMultiTileParams { - const bf16_t* __restrict__ q; - CUtensorMap* __restrict__ tma_k; // Array of [n_h] TMA descriptors for K: (s_k, HD) tile (128,16) - CUtensorMap* __restrict__ tma_v; // Array of [n_h] TMA descriptors for V: (HD, s_k) tile (16,16) - const bf16_t* __restrict__ v; // (HD, s_k) — fallback direct GMEM - bf16_t* __restrict__ o; - float* __restrict__ lse; - int s_k, n_h; - float scale; - int q_head_stride, q_batch_stride; - int v_head_stride, v_batch_stride; - int o_head_stride, o_batch_stride; - int lse_head_stride, lse_batch_stride; -}; - -template -__global__ void __launch_bounds__(192) -fmha_6warp_tma_multitile_kernel(FmhaTmaMultiTileParams params) { - 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; - static constexpr int TMEM_N = (HD <= 128) ? 128 : (HD <= 256) ? 256 : 512; - 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); - - 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 bool is_mma_warp = (wid == 4); - const bool is_load_warp = (wid == 5); - const int s_k = params.s_k; - const int n_kv_tiles = (s_k + SK_TILE - 1) / SK_TILE; - const float scale = params.scale; - - const bf16_t* __restrict__ q_head = params.q + head_idx * params.q_head_stride + batch_idx * params.q_batch_stride; - // v_head kept for reference; V loaded via TMA - // const bf16_t* __restrict__ v_head = params.v + head_idx * params.v_head_stride + batch_idx * params.v_batch_stride; - bf16_t* __restrict__ o_head = params.o + head_idx * params.o_head_stride + batch_idx * params.o_batch_stride; - float* __restrict__ lse_head = params.lse ? params.lse + head_idx * params.lse_head_stride + batch_idx * params.lse_batch_stride : nullptr; - CUtensorMap* __restrict__ my_tma_k = params.tma_k + batch_idx * params.n_h + head_idx; - CUtensorMap* __restrict__ my_tma_v = params.tma_v + batch_idx * params.n_h + head_idx; - - // ================================================================ - // SMEM allocation - // ================================================================ - 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; - 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* 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* s_p_vals = (float*)(sbuf + off); off += SK_TILE * sizeof(float); - float* sTileMax = (float*)(sbuf + off); off += sizeof(float); - float* sTileSum = (float*)(sbuf + off); off += sizeof(float); - - // Init - 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"); - } - __syncthreads(); - uint32_t tb = *sTmemBase; - const uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar); - int phase = 0; - - // Q loaded per-K-sub-tile inside the kt loop (same as single-tile kernel) - // Register accumulator - float o_acc[HD]; - float running_max = -INFINITY; - float running_sum = 0.0f; - if (wid == 0 && lane == 0) { - for (int d = 0; d < HD; d++) o_acc[d] = 0.0f; - } - - // ================================================================ - // Multi-tile KV loop - // ================================================================ - for (int kv_tile = 0; kv_tile < n_kv_tiles; kv_tile++) { - int kv_start = kv_tile * SK_TILE; - int kv_len = min(SK_TILE, s_k - kv_start); - - // ---- QK GEMM ---- - for (int kt = 0; kt < NKT_QK; kt++) { - // Load Q sub-tile (same as single-tile kernel: local d → canonical at column 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) { - 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_head[full_d]; - } - } - } - - // TMA load K sub-tile at (kt*16, kv_start) - if (is_load_warp && lane == 0) { - tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)my_tma_k, - mbar_addr, kt * MMA_K_BF16, kv_start); - tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); - } - tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; - __syncthreads(); - - // Convert sTmaBuf → canonical sK0 - for (int i = tid; i < TILE_SZ; i += 192) sK0[i] = 0; - for (int i = tid; i < kv_len * 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(); - - // 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 (warp 0, lane 0, T=1) ---- - 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)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane == 0) for (int c=0;c<8;c++) { - int col = n*8+c; - if (col < kv_len) { - s_vals[col] = tmp[c] * scale; - row_max = fmaxf(row_max, tmp[c] * scale); - } - } - } - row_max = wmax(row_max); - float row_sum = 0.0f; - if (lane == 0) { - for (int j=0;j 0); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - } - __syncthreads(); - } - } - asm volatile("fence.sc.gpu;" ::: "memory"); - __syncthreads(); - - // ---- Read O_tile from TMEM, rescale, accumulate ---- - if (wid == 0) { - float tile_max = *sTileMax; - float tile_sum = *sTileSum; - float o_tile[HD]; - - // Read O_tile from TMEM (warp-collective, but only lane 0 uses data) - for (int n = 0; n < N_NSUB * 2; 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 (lane == 0) for (int c=0;c<8;c++) { - int d = n*8+c; - if (d < HD) o_tile[d] = tmp[c]; - } - } - - // Online softmax rescale + accumulate - if (lane == 0) { - if (kv_tile == 0) { - running_max = tile_max; - running_sum = tile_sum; - for (int d = 0; d < HD; d++) o_acc[d] = o_tile[d]; - } else { - float old_max = running_max; - running_max = fmaxf(running_max, tile_max); - float rescale = expf(old_max - running_max); - for (int d = 0; d < HD; d++) o_acc[d] = o_acc[d] * rescale + o_tile[d] * expf(tile_max - running_max); - running_sum = running_sum * rescale + tile_sum * expf(tile_max - running_max); - } - } - } - __syncthreads(); - } - - // ================================================================ - // Final epilogue - // ================================================================ - if (wid == 0 && lane == 0) { - float inv_rs = 1.0f / running_sum; - for (int d = 0; d < HD; d++) o_head[d] = f32_to_bf16(o_acc[d] * inv_rs); - if (lse_head) lse_head[0] = logf(running_sum) + running_max; - } - __syncthreads(); - if (is_mma_warp) tmem_dealloc(tb, TMEM_N); -} - -} // namespace diff --git a/dsv4/kernels/attention/fmha_multihead_capi.cu b/dsv4/kernels/attention/fmha_multihead_capi.cu deleted file mode 100644 index 4cb17385..00000000 --- a/dsv4/kernels/attention/fmha_multihead_capi.cu +++ /dev/null @@ -1,105 +0,0 @@ -/** - * 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_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; // 2048 BF16 - constexpr int V_SUB_SZ = 256; - // 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) - // With 128B alignment between sections - int total = base + sQ0 + sK0 + (sPk + 127) + (sV + 127) + sp + (sO + 127) + 256; - // Round up to 128B - return (total + 127) & ~127; -} - -int fmha_multihead_decode_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 -) { - 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; - - 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); - } - - cudaError_t err; - 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 return -1; - - err = cudaGetLastError(); - if (err != cudaSuccess) return (int)err; - return 0; -} - -} // extern "C" diff --git a/dsv4/kernels/attention/fmha_multihead_op.py b/dsv4/kernels/attention/fmha_multihead_op.py deleted file mode 100644 index f30e8c7c..00000000 --- a/dsv4/kernels/attention/fmha_multihead_op.py +++ /dev/null @@ -1,267 +0,0 @@ -"""DSV4 FMHA — 6-warp multi-head decode kernel loader. - -Precompiles the raw CUDA kernel with nvcc (sm_100a) on first use, -then loads the .so via ctypes. This bypasses torch.utils.cpp_extension -which compiles with -arch=sm_100 (missing tcgen05 support). - -Decode-only: T=1, single KV segment (N <= 128). -Supports MHA, MQA, and GQA attention patterns. -""" - -import torch -import logging -import os -import subprocess -import ctypes -from typing import Optional - -logger = logging.getLogger(__name__) - -KERNEL_DIR = os.path.dirname(os.path.abspath(__file__)) -REPO_ROOT = os.path.normpath(os.path.join(KERNEL_DIR, "..", "..")) -SOURCE = os.path.join(KERNEL_DIR, "fmha_multihead_capi.cu") -BUILD_DIR = os.path.join(REPO_ROOT, "build", "fmha_multihead") -SO_NAME = "libfmha_multihead_decode.so" - -_lib = None -_lib_lock = False - - -def _find_nvcc(): - """Find nvcc on the system.""" - for c in ["/usr/local/cuda-13.2/bin/nvcc", "/usr/local/cuda/bin/nvcc"]: - if os.path.isfile(c): - return c - # Try PATH - import shutil - nvcc = shutil.which("nvcc") - if nvcc: - return nvcc - raise RuntimeError("nvcc not found — required for tcgen05 kernel compilation") - - -def _ensure_built(): - """Build the shared library with nvcc if needed. Returns .so path.""" - global _lib - if _lib is not None: - return _lib - - so_path = os.path.join(BUILD_DIR, SO_NAME) - - # Check if rebuild needed - need_build = True - if os.path.isfile(so_path): - src_mtime = os.path.getmtime(SOURCE) - for dep in ["fmha_common.cuh", "fmha_umma_desc.cuh", - "fmha_6warp_multihead.cuh", "fmha_multihead_capi.cu"]: - dep_path = os.path.join(KERNEL_DIR, dep) - if os.path.isfile(dep_path): - src_mtime = max(src_mtime, os.path.getmtime(dep_path)) - need_build = src_mtime > os.path.getmtime(so_path) - - if not need_build: - logger.info(f"Using cached {so_path}") - _lib = ctypes.CDLL(so_path) - return _lib - - logger.info(f"Building {SO_NAME} with nvcc (sm_100a)...") - os.makedirs(BUILD_DIR, exist_ok=True) - - nvcc = _find_nvcc() - cmd = [ - nvcc, - "-std=c++20", - "-shared", - "-Xcompiler", "-fPIC", - "-gencode=arch=compute_100a,code=sm_100a", - "-gencode=arch=compute_100a,code=compute_100a", - f"-I{KERNEL_DIR}", - f"-I{REPO_ROOT}", - "-O3", - "--expt-relaxed-constexpr", - SOURCE, - "-o", so_path, - "-lcudart", - "-lcuda", - ] - - result = subprocess.run(cmd, capture_output=True, text=True) - if result.returncode != 0: - raise RuntimeError(f"nvcc compilation failed:\n{result.stderr}") - if result.stderr: - logger.debug(f"nvcc warnings:\n{result.stderr}") - - _lib = ctypes.CDLL(so_path) - logger.info(f"Built and loaded {so_path}") - return _lib - - -def _get_lib(): - """Get or build the shared library.""" - global _lib_lock - if _lib is not None: - return _lib - if _lib_lock: - raise RuntimeError("Recursive build") - _lib_lock = True - try: - return _ensure_built() - finally: - _lib_lock = False - - -# --------------------------------------------------------------------------- -# Kernel launch via ctypes -# --------------------------------------------------------------------------- - -def fmha_multihead_decode_raw( - q: torch.Tensor, # (batch, n_h, 1, hd) BF16, contiguous - k: torch.Tensor, # (batch, n_kv, N, hd) BF16, contiguous - v: torch.Tensor, # (batch, n_kv, hd, N) BF16, contiguous - scale: float, - n_comp: int, - swa_len: int, - is_causal: bool, - attn_sink: torch.Tensor, # (batch, n_h) FP32 — unused by kernel currently -) -> tuple[torch.Tensor, torch.Tensor]: - """Launch the 6-warp multi-head FMHA kernel. Returns (O, LSE). - - O: (batch, n_h, 1, hd) BF16 - LSE: (batch, n_h, 1) FP32 - """ - lib = _get_lib() - - B = q.shape[0] - n_h = q.shape[1] - hd = q.shape[3] - n_kv = k.shape[1] - N = k.shape[2] - - assert q.shape[2] == 1, f"Decode requires T=1, got T={q.shape[2]}" - assert hd in (64, 128, 256), f"Unsupported hd={hd}" - assert N > 0, f"N must be positive, got N={N}" - - q_per_kv = n_h // n_kv - - # GQA: expand K/V to (1, n_h, ...) so head_idx * stride = correct data - if n_kv < n_h: - k = k.repeat_interleave(q_per_kv, dim=1) - v = v.repeat_interleave(q_per_kv, dim=1) - - # The kernel template has SK_TILE=128 hardcoded in the softmax loop. - # When N < 128, pad K and V to 128 so the kernel processes zeros for - # the extra positions (correctly gets zero attention weight after softmax). - # When N > 128, the kernel loops over KV tiles internally (FA2 online softmax). - if N < 128: - pad_len = 128 - N - k = torch.cat([k, - torch.zeros(k.shape[0], k.shape[1], pad_len, hd, dtype=torch.bfloat16, device=k.device)], dim=2) - v = torch.cat([v, - torch.zeros(v.shape[0], v.shape[1], hd, pad_len, dtype=torch.bfloat16, device=v.device)], dim=3) - k = k.contiguous() - v = v.contiguous() - N = 128 - else: - k = k.contiguous() - v = v.contiguous() - - # Pad N to multiple of 128 for the KV tile loop - n_kv_tiles = (N + 127) // 128 - N_padded = n_kv_tiles * 128 - if N < N_padded: - pad_len = N_padded - N - k = torch.cat([k, - torch.zeros(k.shape[0], k.shape[1], pad_len, hd, dtype=torch.bfloat16, device=k.device)], dim=2) - v = torch.cat([v, - torch.zeros(v.shape[0], v.shape[1], hd, pad_len, dtype=torch.bfloat16, device=v.device)], dim=3) - k = k.contiguous() - v = v.contiguous() - N = N_padded - - q = q.contiguous() - - o = torch.zeros(B, n_h, 1, hd, dtype=torch.bfloat16, device=q.device) - lse = torch.zeros(B, n_h, 1, dtype=torch.float32, device=q.device) - - # Compute strides in BF16 elements - q_hs = q.stride(1) - q_bs = q.stride(0) - k_hs = k.stride(1) - k_bs = k.stride(0) - v_hs = v.stride(1) - v_bs = v.stride(0) - o_hs = o.stride(1) - o_bs = o.stride(0) - lse_hs = lse.stride(1) - lse_bs = lse.stride(0) - - # Call the C API - ret = lib.fmha_multihead_decode_launch( - ctypes.c_void_p(q.data_ptr()), - ctypes.c_void_p(k.data_ptr()), - ctypes.c_void_p(v.data_ptr()), - ctypes.c_void_p(o.data_ptr()), - ctypes.c_void_p(lse.data_ptr()), - ctypes.c_int(B), - ctypes.c_int(n_h), - ctypes.c_int(n_kv), - ctypes.c_int(N), - ctypes.c_int(hd), - ctypes.c_int(q_hs), - ctypes.c_int(q_bs), - ctypes.c_int(k_hs), - ctypes.c_int(k_bs), - ctypes.c_int(v_hs), - ctypes.c_int(v_bs), - ctypes.c_int(o_hs), - ctypes.c_int(o_bs), - ctypes.c_int(lse_hs), - ctypes.c_int(lse_bs), - ctypes.c_float(scale), - ) - - if ret != 0: - raise RuntimeError(f"Kernel launch failed with code {ret}") - - return o, lse - - -# --------------------------------------------------------------------------- -# Custom op registration -# --------------------------------------------------------------------------- - -@torch.library.custom_op("dsv4::fmha_multihead_decode", mutates_args=()) -def fmha_multihead_decode( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float, - n_comp: int, - swa_len: int, - is_causal: bool, - attn_sink: torch.Tensor, -) -> torch.Tensor: - o, _ = fmha_multihead_decode_raw( - q, k, v, scale, n_comp, swa_len, is_causal, attn_sink - ) - return o - - -@fmha_multihead_decode.register_fake -def _(q, k, v, scale, n_comp, swa_len, is_causal, attn_sink): - return torch.empty_like(q) - - -def fmha_multihead_decode_with_lse( - q, k, v, scale, n_comp=0, swa_len=0, is_causal=False, attn_sink=None, -) -> tuple[torch.Tensor, torch.Tensor]: - if attn_sink is None: - attn_sink = torch.zeros(q.shape[0], q.shape[1], - dtype=torch.float32, device=q.device) - return fmha_multihead_decode_raw( - q, k, v, scale, n_comp, swa_len, is_causal, attn_sink - ) - - -def can_use_6warp_decode(T: int, N: int, hd: int, n_segments: int) -> bool: - return T == 1 and hd in (64, 128, 256) diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py index 2cb45815..736c6085 100644 --- a/dsv4/kernels/attention/production.py +++ b/dsv4/kernels/attention/production.py @@ -300,57 +300,6 @@ def _run_fmha_segmented( # Fast path: 6-warp multi-head decode kernel # --------------------------------------------------------------------------- -def _dsv4_attention_fast_decode( - q: torch.Tensor, # (n_q_heads, T, hd) BF16 — T must be 1 - k: torch.Tensor, # (n_kv_heads, N, hd) BF16 - v: torch.Tensor, # (n_kv_heads, N, hd) BF16 - scale: float, - n_comp: int = 0, - sink_bias: Optional[torch.Tensor] = None, # (n_q_heads,) FP32 -) -> torch.Tensor: - """Fast decode path using 6-warp multi-head FMHA kernel. - - Single kernel launch for all heads. No Python KV merge. - No cudaDeviceSynchronize on the hot path. - - Q layout: (n_q_heads, 1, hd) — reshape to (1, n_q_heads, 1, hd) for grid - K layout: (n_kv, N, hd) — reshape to (1, n_kv, N, hd) - V layout: (n_kv, hd, N) — TRANSPOSED from input (n_kv, N, hd) - """ - from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw - - n_q, T, hd = q.shape - n_kv = k.shape[0] if k.dim() == 3 else 1 - N = k.shape[-2] if k.dim() == 3 else k.shape[0] - - # Reshape Q: (n_q, 1, hd) -> (1, n_q, 1, hd) - q_4d = q.unsqueeze(0).contiguous() - - # Reshape K: (n_kv, N, hd) -> (1, n_kv, N, hd) - if k.dim() == 2: - k_4d = k.unsqueeze(0).unsqueeze(0).contiguous() # (N,hd) -> (1,1,N,hd) - else: - k_4d = k.unsqueeze(0).contiguous() # (n_kv,N,hd) -> (1,n_kv,N,hd) - - # V must be (batch, n_kv, hd, N) for the kernel — transpose the last two dims - if v.dim() == 2: - v_4d = v.unsqueeze(0).unsqueeze(0).transpose(-1, -2).contiguous() - else: - v_4d = v.unsqueeze(0).transpose(-1, -2).contiguous() # (1,n_kv,hd,N) - - # Sink bias: (n_q,) -> (1, n_q) — zeros if not provided - if sink_bias is not None: - sb = sink_bias.unsqueeze(0).contiguous() - else: - sb = torch.zeros(1, n_q, dtype=torch.float32, device=q.device) - - o_4d, _lse = fmha_multihead_decode_raw( - q_4d, k_4d, v_4d, scale, n_comp, 0, False, sb - ) # (1, n_q, 1, hd) - - return o_4d.squeeze(0) # (n_q, 1, hd) - - def _dsv4_attention_multitile( q: torch.Tensor, k: torch.Tensor,