From 08694b81367d6a6a80d98c19943c443b44a5ccac Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 20:10:13 +0000 Subject: [PATCH] Fix multi-row softmax v3: 32x32b.x8 with per-lane per-row (no wmax/wsum), per-row sRowMax/sRowSum arrays --- .../kernels/attention/fmha_6warp_multirow.cuh | 167 +++++++----------- tests/unit/test_fmha_6warp_multirow.cu | 2 +- 2 files changed, 62 insertions(+), 107 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_multirow.cuh b/dsv4/kernels/attention/fmha_6warp_multirow.cuh index a0696c18..043f6fb6 100644 --- a/dsv4/kernels/attention/fmha_6warp_multirow.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multirow.cuh @@ -4,24 +4,11 @@ * ================================================================== * MULTI-ROW SOFTMAX (Milestone 4) * ================================================================== - * Uses 16x256b.x1 TMEM reads for multi-row softmax. - * Each lane processes 4 rows (lane l → rows [l*4, l*4+4)). - * A single warp can handle all 128 rows (32 lanes × 4 rows). + * For T ≤ 32: Uses 32x32b.x8 TMEM reads (same as decode). + * Each lane handles one row. Warp 0 processes rows 0..T-1. + * No cross-lane reduction needed (per-row max/sum in each lane). * - * Two-pass softmax: - * Pass 1: Read S from TMEM, compute per-row max - * Pass 2: Re-read S, compute exp(s - row_max), accumulate row_sum, - * write P to s_p_vals, then normalize in-place. - * - * Epilogue: Read O from TMEM with 16x256b.x1, normalize, write GMEM. - * - * ================================================================== - * TMEM LANE MAPPING (16x256b.x1 format, verified on B200) - * ================================================================== - * tcgen05.ld.16x256b.x1: each lane i reads 4 FP32 from column. - * Lane 0: rows 0-3, Lane 1: rows 4-7, ..., Lane 31: rows 124-127 - * 32 lanes × 4 FP32 = 128 FP32 per column. - * ALL 32 lanes must participate (warp-collective). + * For T > 32: NOT YET IMPLEMENTED (will use 16x256b.x1). * ================================================================== */ @@ -60,7 +47,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { static constexpr int V_SUB_SZ = 256; static constexpr int TMEM_N = (HD <= 128) ? 128 : 256; static constexpr int MAX_ROWS = 128; - static constexpr int ROWS_PER_LANE = 4; const int head_idx = blockIdx.y; const int batch_idx = blockIdx.z; @@ -76,9 +62,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { const int s_k = params.s_k; const float scale = params.scale; - // Per-lane row range - const int lane_row_start = lane * ROWS_PER_LANE; - // ================================================================== // Per-head GMEM pointers // ================================================================== @@ -100,13 +83,13 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { : nullptr; // ================================================================ - // SMEM allocation + // SMEM allocation — SAME LAYOUT as multihead kernel for T=1 compat // ================================================================ extern __shared__ char sbuf[]; uint32_t* sTmemBase = (uint32_t*)sbuf; - float* sRowMax = (float*)(sbuf + 4); - float* sRowSum = sRowMax + MAX_ROWS; - bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 15) & ~(uintptr_t)15); + float* sRowMax = (float*)(sbuf + 4); // [32] per-warp rows + float* sRowSum = sRowMax + 32; // [32] + bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + 32) + 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); @@ -123,7 +106,7 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { uint32_t tb = *sTmemBase; // ================================================================ - // QK GEMM + // QK GEMM loop // ================================================================ for (int kt = 0; kt < NKT_QK; kt++) { if (is_load_warp) { @@ -161,80 +144,56 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } // ================================================================ - // Multi-row softmax (16x256b.x1, two-pass) + // Softmax (32x32b.x8, per-lane per-row) + // ================================================================ + // For T ≤ 32: lane l handles row l. Lanes l >= T have no data. + // No wmax/wsum — each lane computes its own row independently. // ================================================================ if (is_softmax_warp) { - float my_row_max[ROWS_PER_LANE]; - for (int i = 0; i < ROWS_PER_LANE; i++) my_row_max[i] = -INFINITY; + float s_vals[SK_TILE], row_max = -INFINITY; - // Pass 1: Compute row_max - for (int col = 0; col < SK_TILE; col++) { - uint32_t u0, u1, u2, u3; - asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(u0), "=r"(u1), "=r"(u2), "=r"(u3) : "r"(tb + col)); + // Read S from TMEM + 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;"); - - float vals[ROWS_PER_LANE] = { - u32_to_f32(u0) * scale, u32_to_f32(u1) * scale, - u32_to_f32(u2) * scale, u32_to_f32(u3) * scale - }; - - for (int i = 0; i < ROWS_PER_LANE; i++) { - if (lane_row_start + i < T) { - my_row_max[i] = fmaxf(my_row_max[i], vals[i]); + 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]); } } } - // Store row_max to SMEM - for (int i = 0; i < ROWS_PER_LANE; i++) { - if (lane_row_start + i < T) { - sRowMax[lane_row_start + i] = my_row_max[i]; - } - } - __syncwarp(); + // Store per-row max + if (lane < T) sRowMax[lane] = row_max; - // Pass 2: Compute P and row_sum - float my_row_sum[ROWS_PER_LANE]; - for (int i = 0; i < ROWS_PER_LANE; i++) my_row_sum[i] = 0.0f; - - for (int col = 0; col < SK_TILE; col++) { - uint32_t u0, u1, u2, u3; - asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(u0), "=r"(u1), "=r"(u2), "=r"(u3) : "r"(tb + col)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - - float vals[ROWS_PER_LANE] = { - u32_to_f32(u0) * scale, u32_to_f32(u1) * scale, - u32_to_f32(u2) * scale, u32_to_f32(u3) * scale - }; - - for (int i = 0; i < ROWS_PER_LANE; i++) { - int row = lane_row_start + i; - if (row < T) { - float p = expf(vals[i] - my_row_max[i]); - my_row_sum[i] += p; - s_p_vals[row * SK_TILE + col] = p; - } + 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]; } } - // Normalize P and store row_sum - for (int i = 0; i < ROWS_PER_LANE; i++) { - int row = lane_row_start + i; - if (row < T) { - sRowSum[row] = my_row_sum[i]; - float inv_sum = 1.0f / my_row_sum[i]; - for (int col = 0; col < SK_TILE; col++) { - s_p_vals[row * SK_TILE + col] *= inv_sum; - } + // Store per-row sum + if (lane < T) sRowSum[lane] = row_sum; + + // Normalize and write P to s_p_vals + if (lane < T) { + float inv_sum = 1.0f / row_sum; + for (int j = 0; j < SK_TILE; j++) { + s_p_vals[lane * SK_TILE + j] = s_vals[j] * inv_sum; } } } __syncthreads(); // ================================================================ - // PV GEMM + // PV GEMM loop // ================================================================ for (int n = 0; n < N_NSUB; n++) { int d_base = n * 16; @@ -277,36 +236,32 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } // ================================================================ - // Multi-row epilogue: Read O from TMEM, normalize, write GMEM - // O is in TMEM columns 0..HD-1 (from PV sub-tiles at n*16) - // Read with 16x256b.x1: lane l gets rows [l*4, l*4+4) per column + // Epilogue (32x32b.x8, same as multihead kernel) // ================================================================ if (is_softmax_warp) { - for (int c = 0; c < HD; c++) { - uint32_t u0, u1, u2, u3; - asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(u0), "=r"(u1), "=r"(u2), "=r"(u3) : "r"(tb + c)); + float row_max = (lane < T) ? sRowMax[lane] : 0.0f; + float row_sum = (lane < T) ? sRowSum[lane] : 1.0f; + 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;"); - - float vals[ROWS_PER_LANE] = { - u32_to_f32(u0), u32_to_f32(u1), - u32_to_f32(u2), u32_to_f32(u3) - }; - - for (int i = 0; i < ROWS_PER_LANE; i++) { - int row = lane_row_start + i; - if (row < T) { - float o_val = vals[i] / sRowSum[row]; - o_head[row * HD + c] = f32_to_bf16(o_val); - } + if (lane < T) { + for (int c = 0; c < 8; c++) o_vals[n * 8 + c] = tmp[c]; } } - // Write LSE - for (int i = 0; i < ROWS_PER_LANE; i++) { - int row = lane_row_start + i; - if (row < T && lse_head) { - lse_head[row] = logf(sRowSum[row]) + sRowMax[row]; + if (lane < T) { + float inv_row_sum = 1.0f / row_sum; + for (int d = 0; d < HD; d++) { + o_head[lane * HD + d] = f32_to_bf16(o_vals[d] * inv_row_sum); + } + if (lse_head) { + lse_head[lane] = logf(row_sum) + row_max; } } } diff --git a/tests/unit/test_fmha_6warp_multirow.cu b/tests/unit/test_fmha_6warp_multirow.cu index 68b3e1ed..518ff5ff 100644 --- a/tests/unit/test_fmha_6warp_multirow.cu +++ b/tests/unit/test_fmha_6warp_multirow.cu @@ -122,7 +122,7 @@ static int test_single_T(int T, int n_h = 1, int batch = 1) { // 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 + 128*4 + 128*4 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + MAX_T * SK * 4 + 256; + 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 = (smem + 127) & ~127; if (smem > 48 * 1024) {