From 863a030c3b28128377cacd8a97a739c16d7c31e2 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 22:52:52 +0000 Subject: [PATCH] fmha_multirow: rewrite with 32x32b.x8 only, no s_p_vals, row_page addressing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Kill 64KB s_p_vals buffer — P is streamed per K-tile through sPk - All TMEM ops use 32x32b.x8 exclusively (16x256b.x1 crashes on 2nd call) - T>32: 4 softmax warps use row_page offset in TMEM address (row<<16) - Lane l in warp w handles row w*32+l - Two-pass softmax: pass 1 row_max, pass 2 exp/sum interleaved with PV - PV: N=16 sub-tiles, SS MMA sPk(128,16) × sV(16,16) → TMEM - Epilogue: 32x32b.x8 TMEM read, normalize, BF16 → GMEM - SMEM budget: ~14KB (well within 232KB) --- .../kernels/attention/fmha_6warp_multirow.cuh | 415 +++++++++++------- tests/unit/test_fmha_6warp_multirow.cu | 24 +- 2 files changed, 265 insertions(+), 174 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_multirow.cuh b/dsv4/kernels/attention/fmha_6warp_multirow.cuh index 4b2bf684..6a6ff54f 100644 --- a/dsv4/kernels/attention/fmha_6warp_multirow.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multirow.cuh @@ -1,11 +1,57 @@ /** * DSV4 FMHA — 6-warp specialized kernel, multi-row softmax (prefill T>1). * - * T <= 32: Warp 0 only, 32x32b.x8, lane l = row l. - * T > 32: 4 softmax warps, 16x256b.x1, two-pass online softmax. - * Warp w handles rows [w*32, min((w+1)*32, T)). - * Lane j reads rows j*4+0..3 per TMEM column. - * Active lanes for warp w: j in [w*8, w*8+8). + * ================================================================== + * 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. + * + * TMEM operations ALL use 32x32b.x8 format exclusively. + * The 16x256b.x1 format crashes on the 2nd call — verified on B200. + * + * TMEM ADDRESSING for 32x32b.x8: + * addr = tmem_base + (row_page << 16) + col_group * 8 + * - col_group selects 8 consecutive TMEM columns + * - row_page selects a 32-row "page" (row_page=0 → rows 0-31, =1 → rows 32-63, etc.) + * - Lane l within the warp reads row (row_page*32 + l) across the 8 columns + * - ALL 32 lanes must call with the same address (warp-collective) + * + * SOFTMAX (two regimes): + * - T <= 32: warp 0 only. row_page=0. Lane l processes row l. 32x32b.x8 reads. + * - T > 32: 4 softmax warps. Warp w uses row_page=w. + * Lane l in warp w processes row w*32+l. Two-pass: + * (1) row_max, (2) exp+sum, stream P to sPk per K-tile. + * + * P STAGING FOR PV (no s_p_vals buffer): + * P is written one K-tile at a time to sPk (128, 16) SMEM in canonical + * layout. Softmax warps compute P values in registers, write to sPk. + * Then MMA warp does PV with sPk × sV. + * The softmax and PV are interleaved K-tile-by-K-tile (required because + * we can't hold the full 128×128 P matrix in SMEM). + * + * PV GEMM: + * N=16 sub-tiles (avoids Layout D N=64 bug). SS MMA: sPk(128,16) × sV(16,16). + * Accumulate into TMEM across K-tiles and N-sub-tiles. + * + * EPILOGUE: + * 32x32b.x8 TMEM read → normalize (1/row_sum) → BF16 → GMEM. + * LSE = ln(row_sum) + row_max for D5 KV merge. + * + * ================================================================== + * SMEM BUDGET (HD=64) + * ================================================================== + * sTmemBase: 8 B + * sRowMax: 128 × 4 = 512 B + * sRowSum: 128 × 4 = 512 B + * padding: 128 B (128B alignment for UMMA) + * sQ0: 128 × 16 × 2 = 4096 B (one K-tile) + * sK0: 128 × 16 × 2 = 4096 B (one K-tile) + * sPk: 128 × 16 × 2 = 4096 B (one K-tile P buffer) + * sV: 16 × 16 × 2 = 512 B (one N-sub-tile, one K-tile) + * Total: ~14 KB — well within 232 KB budget + * ================================================================== */ #pragma once @@ -34,13 +80,14 @@ struct FmhaMultiRowParams { 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 = 256; + static constexpr int NKT_QK = HD / MMA_K_BF16; // K-tiles for QK GEMM + static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; // K-tiles for PV GEMM (= 8) + static constexpr int N_NSUB = HD / 16; // N-sub-tiles for PV (avoids N=64 bug) + static constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16 elements + static constexpr int V_SUB_SZ = 16 * MMA_K_BF16; // 256 BF16 elements (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; // 16 core matrices along M const int head_idx = blockIdx.y; const int batch_idx = blockIdx.z; @@ -53,56 +100,81 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { const int T = params.T; const int s_k = params.s_k; const float scale = params.scale; - const bool full_tile = (T > 32); + // 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; - // SMEM + // ================================================================ + // SMEM allocation — NO s_p_vals buffer + // ================================================================ extern __shared__ char sbuf[]; - uint32_t* sTmemBase = (uint32_t*)sbuf; - float* sRowMax = (float*)(sbuf + 4); // [MAX_ROWS] - float* sRowSum = sRowMax + MAX_ROWS; // [MAX_ROWS] - bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 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); // [MAX_ROWS][SK_TILE] + uint32_t* sTmemBase = (uint32_t*)sbuf; // 4 B (8B-aligned) + float* sRowMax = (float*)(sbuf + 8); // [MAX_ROWS] = 512 B + float* sRowSum = sRowMax + MAX_ROWS; // [MAX_ROWS] = 512 B + // Align sQ0 to 128B for UMMA descriptors + bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 127) & ~(uintptr_t)127); + bf16_t* sK0 = sQ0 + TILE_SZ; // (128, 16) BF16 canonical + bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + TILE_SZ) + 127) & ~(uintptr_t)127); // (128, 16) BF16 canonical + bf16_t* sV = (bf16_t*)(((uintptr_t)(sPk + TILE_SZ) + 127) & ~(uintptr_t)127); // (16, 16) BF16 canonical - // TMEM alloc - if (is_mma_warp) { uint32_t p = __cvta_generic_to_shared(sTmemBase); tmem_alloc(p, TMEM_N); } + // ================================================================ + // 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 + // Row assignment for softmax + // ================================================================ + // T <= 32: only warp 0, row_page = 0, lane l = row l + // T > 32: 4 softmax warps, warp w uses row_page = w + // lane l in warp w handles row w*32 + l + const bool my_warp_active = (T <= 32) ? (wid == 0) : is_softmax_warp; + const int my_row_page = (T <= 32) ? 0 : wid; + const int my_row = my_row_page * 32 + lane; + const bool my_row_active = my_warp_active && (my_row < T); + const uint32_t my_tmem_row_offset = (uint32_t)(my_row_page * 32) << 16; + + // ================================================================ + // QK GEMM: accumulate all K-tiles into TMEM // ================================================================ for (int kt = 0; kt < NKT_QK; kt++) { if (is_load_warp) { - constexpr int CORES_MN = 128 / 8; + // Load Q K-tile: Q is (T, HD), columns [kt*16, kt*16+16) 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; - int ck = d / 8, lc = d % 8; - int core_mn = r / 8, local_r = r % 8; - sQ0[ck * CORES_MN * 64 + core_mn * 64 + local_r * 8 + lc] = q_head[r * HD + full_d]; + if (full_d < HD) { + int ck = d / 8, lc = d % 8; + int core_mn = r / 8, local_r = r % 8; + sQ0[ck * CORES_MN * 64 + core_mn * 64 + local_r * 8 + lc] = q_head[r * HD + full_d]; + } } } + // Load K K-tile: K is (s_k, HD), columns [kt*16, kt*16+16) 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; - 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 + full_d]; + if (full_d < HD) { + 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 + 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); @@ -113,189 +185,198 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { __syncthreads(); } + // TMEM now contains S[128, s_k] — the raw QK dot products. + // S[row, col] needs scaling by `scale` before softmax. + // ================================================================ - // SOFTMAX + // SOFTMAX PASS 1: compute row_max for each row // ================================================================ - if (!full_tile) { - // T <= 32: warp 0, 32x32b.x8, lane l = row l - 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 < T) { - for (int c = 0; c < 8; c++) { - s_vals[n * 8 + c] = tmp[c] * scale; - row_max = fmaxf(row_max, s_vals[n * 8 + c]); + // Read all SK_TILE columns (16 groups of 8 columns each). + // 32x32b.x8: addr = tb + (row_page*32 << 16) + n*8 + // Each lane reads its assigned row across 8 columns per read. + float my_row_max = -INFINITY; + const int NUM_READS = SK_TILE / 8; // = 16 + + if (my_warp_active) { + for (int n = 0; n < NUM_READS; n++) { + float tmp[8]; + uint32_t addr = tb + my_tmem_row_offset + n * 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"(addr)); + 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 val = tmp[c] * scale; + my_row_max = fmaxf(my_row_max, val); } } } - if (lane < T) sRowMax[lane] = row_max; - float row_sum = 0.0f; - if (lane < T) { - for (int j = 0; j < SK_TILE; j++) { s_vals[j] = expf(s_vals[j] - row_max); row_sum += s_vals[j]; } - } - if (lane < T) sRowSum[lane] = row_sum; - if (lane < T) { - float inv = 1.0f / row_sum; - for (int j = 0; j < SK_TILE; j++) s_p_vals[lane * SK_TILE + j] = s_vals[j] * inv; - } - } - } else { - // T > 32: 4 warps, 16x256b.x1, two-pass softmax - if (is_softmax_warp) { - const int warp_row_start = wid * 32; - const int warp_row_end = min((wid + 1) * 32, T); - const int lane_row_base = lane * 4; - const bool active = (lane_row_base >= warp_row_start) && (lane_row_base < warp_row_end); - const int vrows = active ? min(4, warp_row_end - lane_row_base) : 0; - - float rmax[4], rsum[4]; - for (int r = 0; r < 4; r++) rmax[r] = -INFINITY; - - // Pass 1: row_max - for (int col = 0; col < SK_TILE; col++) { - float v[4]; - asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0,%1,%2,%3},[%4,%5];" - : "=f"(v[0]), "=f"(v[1]), "=f"(v[2]), "=f"(v[3]) : "r"(tb), "r"(col)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (active) for (int r = 0; r < vrows; r++) rmax[r] = fmaxf(rmax[r], v[r] * scale); - } - - // Pass 2: exp, sum, write P - for (int r = 0; r < 4; r++) rsum[r] = 0.0f; - for (int col = 0; col < SK_TILE; col++) { - float v[4]; - asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0,%1,%2,%3},[%4,%5];" - : "=f"(v[0]), "=f"(v[1]), "=f"(v[2]), "=f"(v[3]) : "r"(tb), "r"(col)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (active) { - for (int r = 0; r < vrows; r++) { - int gr = lane_row_base + r; - float p = expf(v[r] * scale - rmax[r]); - rsum[r] += p; - s_p_vals[gr * SK_TILE + col] = p; - } - } - } - // Normalize and store max/sum - if (active) { - for (int r = 0; r < vrows; r++) { - int gr = lane_row_base + r; - sRowMax[gr] = rmax[r]; - sRowSum[gr] = rsum[r]; - float inv = 1.0f / rsum[r]; - for (int j = 0; j < SK_TILE; j++) s_p_vals[gr * SK_TILE + j] *= inv; - } - } } } + + // Store row_max to SMEM + if (my_row_active) sRowMax[my_row] = my_row_max; __syncthreads(); // ================================================================ - // PV GEMM + // SOFTMAX PASS 2 + PV GEMM (interleaved per PV K-tile) // ================================================================ - for (int n = 0; n < N_NSUB; n++) { - int d_base = n * 16; - for (int kt = 0; kt < NKT_PV; kt++) { - if (is_load_warp) { - constexpr int CORES_MN = 128 / 8; - for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0; - for (int r = 0; r < T; r++) { - for (int c = lane; c < MMA_K_BF16; c += 32) { - int gc = kt * MMA_K_BF16 + c; - float pv = s_p_vals[r * SK_TILE + gc]; - int cmn = r / 8, lr = r % 8, ck = c / 8, lc = c % 8; - sPk[ck * CORES_MN * 64 + cmn * 64 + lr * 8 + lc] = f32_to_bf16(pv); + // For each PV K-tile (16 columns of S): + // 1. Softmax warps re-read TMEM, compute exp(S*scale - row_max) + // 2. Write P to sPk in canonical layout + // 3. Load V sub-tile to sV + // 4. MMA does PV: sPk(128,16) × sV(16,16) → TMEM + // + // We accumulate row_sum across K-tiles in registers. + // P written to sPk is UN-normalized (exp(S*scale - row_max)). + // Normalization happens in the epilogue (divide by row_sum). + // ================================================================ + float my_row_sum = 0.0f; + + for (int pv_kt = 0; pv_kt < NKT_PV; pv_kt++) { + const int col_start = pv_kt * MMA_K_BF16; + const int col_end = col_start + MMA_K_BF16; + // TMEM column groups for this K-tile: [col_start/8, col_end/8) + const int rg_start = col_start / 8; + const int rg_end = col_end / 8; + + // P values for this K-tile (un-normalized) + float p_vals[16]; // MMA_K_BF16 values + for (int c = 0; c < MMA_K_BF16; c++) p_vals[c] = 0.0f; + + if (my_warp_active) { + float rm = my_row_active ? sRowMax[my_row] : 0.0f; + + // Read TMEM for this K-tile's columns + for (int n = rg_start; n < rg_end; n++) { + float tmp[8]; + uint32_t addr = tb + my_tmem_row_offset + n * 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"(addr)); + 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 && col >= col_start && col < col_end) { + float val = expf(tmp[c] * scale - rm); + p_vals[col - col_start] = val; + my_row_sum += val; + } } } + } + } + + // Zero sPk first (load warp does this for all rows) + if (is_load_warp) { + for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0; + } + __syncthreads(); + + // Write P to sPk in canonical layout (softmax warps) + if (my_row_active) { + for (int c = 0; c < MMA_K_BF16; 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(p_vals[c]); + } + } + __syncthreads(); + + // PV GEMM: N=16 sub-tiles + for (int n_sub = 0; n_sub < N_NSUB; n_sub++) { + int d_base = n_sub * 16; + + if (is_load_warp) { + // Load V sub-tile: V is (HD, s_k) in GMEM + // Need V[d_base..d_base+15, col_start..col_end-1] + // Canonical layout for (16, 16): + // g_mn = d/8, g_k = lr/8, llr = d%8, lc = lr%8 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 = kt * MMA_K_BF16 + lr; - int gmn = dd / 8, gk = lr / 8, llr = dd % 8, lc = lr % 8; - sV[gk * 2 * 64 + gmn * 64 + llr * 8 + lc] = v_head[(d_base + dd) * s_k + r]; + int r = col_start + lr; + if (r < s_k && (d_base + dd) < HD) { + int g_mn = dd / 8, g_k = lr / 8; + int 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(); + if (is_mma_warp) { - uint32_t idesc16 = make_idesc(128, 16); + 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 * 16, dp, dv, idesc16, kt > 0); + // accumulate = true for pv_kt > 0 (K-tile accumulation) + // For n_sub > 0, the PV result goes to different TMEM columns (n_sub*16) + // so first call for each n_sub should NOT accumulate + 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(); } } + // Store row_sum to SMEM + if (my_row_active) sRowSum[my_row] = my_row_sum; + __syncthreads(); + // ================================================================ - // EPILOGUE + // EPILOGUE: TMEM → regs → normalize → BF16 → GMEM // ================================================================ - if (!full_tile) { - // T <= 32: warp 0, 32x32b.x8 - if (wid == 0) { - float rm = (lane < T) ? sRowMax[lane] : 0.0f; - float rs = (lane < T) ? sRowSum[lane] : 1.0f; - float o_vals[HD]; - for (int n = 0; n < HD / 8; n++) { + if (my_warp_active) { + if (my_row_active) { + float rm = sRowMax[my_row]; + float rs = sRowSum[my_row]; + float inv_rs = 1.0f / rs; + + // Read O from TMEM: HD values in HD/8 groups of 8 + // O is in TMEM at columns 0..N_NSUB*16-1 + // After PV GEMM, TMEM layout is: + // N-sub-tile n writes to columns n*16..n*16+15 + // 32x32b.x8 reads 8 columns at a time + // So we read (N_NSUB*16)/8 = N_NSUB*2 groups + // But we need to read ALL N_NSUB*2 groups with correct row + for (int n = 0; n < N_NSUB * 2; n++) { float tmp[8]; + uint32_t addr = tb + my_tmem_row_offset + n * 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"(addr)); asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane < T) for (int c = 0; c < 8; c++) o_vals[n * 8 + c] = tmp[c]; - } - if (lane < T) { - float inv = 1.0f / rs; - for (int d = 0; d < HD; d++) o_head[lane * HD + d] = f32_to_bf16(o_vals[d] * inv); - if (lse_head) lse_head[lane] = logf(rs) + rm; - } - } - } else { - // T > 32: 4 warps, 16x256b.x1 - if (is_softmax_warp) { - const int warp_row_start = wid * 32; - const int warp_row_end = min((wid + 1) * 32, T); - const int lane_row_base = lane * 4; - const bool active = (lane_row_base >= warp_row_start) && (lane_row_base < warp_row_end); - const int vrows = active ? min(4, warp_row_end - lane_row_base) : 0; - // Read O from TMEM one column at a time - for (int d = 0; d < HD; d++) { - float v[4]; - asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0,%1,%2,%3},[%4,%5];" - : "=f"(v[0]), "=f"(v[1]), "=f"(v[2]), "=f"(v[3]) : "r"(tb), "r"(d)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (active) { - for (int r = 0; r < vrows; r++) { - int gr = lane_row_base + r; - if (gr < T) { - float inv = 1.0f / sRowSum[gr]; - o_head[gr * HD + d] = f32_to_bf16(v[r] * inv); - } + 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); } } } + // LSE - if (active) { - for (int r = 0; r < vrows; r++) { - int gr = lane_row_base + r; - if (gr < T && lse_head) lse_head[gr] = logf(sRowSum[gr]) + sRowMax[gr]; - } + if (lse_head) { + lse_head[my_row] = logf(rs) + rm; } } } __syncthreads(); - // TMEM dealloc - if (is_mma_warp) tmem_dealloc(tb, TMEM_N); + // TMEM dealloc (warp 4) + if (is_mma_warp) { + tmem_dealloc(tb, TMEM_N); + } } } // namespace dsv4::kernels::attention diff --git a/tests/unit/test_fmha_6warp_multirow.cu b/tests/unit/test_fmha_6warp_multirow.cu index 518ff5ff..160c4a16 100644 --- a/tests/unit/test_fmha_6warp_multirow.cu +++ b/tests/unit/test_fmha_6warp_multirow.cu @@ -4,8 +4,8 @@ * * Tests: * 1. T=1 decode (regression — must match single-row results) - * 2. T=2,4,8,16,32 (small prefill — only warp 0) - * 3. T=64,128 (multi-warp prefill — all 4 softmax warps) + * 2. T=2,4,8,16,32 (small prefill — warp 0 only) + * 3. T=64,128 (multi-warp prefill — 4 softmax warps) * 4. LSE correctness for multi-row * 5. Multi-head + batched with T>1 */ @@ -120,9 +120,9 @@ static int test_single_T(int T, int n_h = 1, int batch = 1) { params.lse_head_stride = T; params.lse_batch_stride = n_h * T; - // SMEM: tmemBase(4) + sRowMax(128*4) + sRowSum(128*4) + padding + sQ0 + sK0 + sPk + sV + s_p_vals - // s_p_vals = 128 * 128 * 4 = 65536 bytes - int smem = 4 + 32*4 + 32*4 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + MAX_T * SK * 4 + 256; + // SMEM: sTmemBase(8) + sRowMax(512) + sRowSum(512) + padding(128) + sQ0(4096) + sK0(4096) + sPk(4096) + sV(512) + slack(256) + // NO s_p_vals buffer — P is streamed per K-tile through sPk + int smem = 8 + MAX_T*4 + MAX_T*4 + 128 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + 256; smem = (smem + 127) & ~127; if (smem > 48 * 1024) { @@ -140,7 +140,7 @@ static int test_single_T(int T, int n_h = 1, int batch = 1) { } sync_err = cudaDeviceSynchronize(); if (sync_err != cudaSuccess) { - printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err)); + printf("CUDA ERROR: ~%s\n", cudaGetErrorString(sync_err)); pass = 0; goto cleanup; } @@ -173,6 +173,16 @@ static int test_single_T(int T, int n_h = 1, int batch = 1) { if (cs < 0.999f) { printf(" FAIL batch=%d head=%d row=%d: cos=%.6f lse_err=%.6f\n", b, h, t, cs, lse_err); + if (failed < 3) { + // Print first few values for debugging + printf(" kernel[0..7]="); + for (int d2 = 0; d2 < min(8, HD); d2++) + printf(" %.4f", bf16_to_f32_host(h_o[(idx * T + t) * HD + d2])); + printf("\n ref [0..7]="); + for (int d2 = 0; d2 < min(8, HD); d2++) + printf(" %.4f", o_ref[t * HD + d2]); + printf("\n"); + } failed++; } } @@ -198,7 +208,7 @@ int main() { // Test 1: T=1 decode (regression — must match single-row) all_pass &= test_single_T(1); - // Test 2: Small prefill (only warp 0) + // Test 2: Small prefill (warp 0 only) all_pass &= test_single_T(2); all_pass &= test_single_T(4); all_pass &= test_single_T(8);