From aaa76c1af13b952ac4cf09b641d3730360990aaa Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 20:08:30 +0000 Subject: [PATCH] Rewrite multi-row softmax using 16x256b.x1 TMEM reads for proper multi-row access --- .../kernels/attention/fmha_6warp_multirow.cuh | 263 ++++++++---------- 1 file changed, 115 insertions(+), 148 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_multirow.cuh b/dsv4/kernels/attention/fmha_6warp_multirow.cuh index c5bdc17a..a0696c18 100644 --- a/dsv4/kernels/attention/fmha_6warp_multirow.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multirow.cuh @@ -4,39 +4,24 @@ * ================================================================== * MULTI-ROW SOFTMAX (Milestone 4) * ================================================================== - * Extends the 6-warp kernel to handle T>1 rows (prefill). + * 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). * - * Architecture: - * Warp 0-3 (softmax/epilogue): Each warp handles up to 32 rows. - * Warp w processes rows [w*32, min(w*32+32, T)). - * For T<=32, only warp 0 is active. - * Warp 4 (MMA): QK + PV GEMM via tcgen05.mma SS - * Warp 5 (data staging): Load Q/K/V from GMEM to SMEM + * 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. * - * TMEM multi-row addressing: - * tcgen05.ld.32x32b.x8 reads 32 rows from TMEM. - * Address: tmem_base + (row_offset << 16) + col - * Warp w uses row_offset = w*32 to read its group of 32 rows. - * Lane l in warp w reads row (w*32 + l) from TMEM. - * - * Q layout for prefill: - * Q is (T, hd) per head, loaded as (128, hd) in SMEM canonical layout. - * Rows 0..T-1 have data; rows T..127 are zero. - * - * P layout for PV GEMM: - * After per-row softmax, P values are stored in s_p_vals[T][SK_TILE]. - * For PV, P is loaded as (128, 16) canonical sub-tiles. - * - * Output: - * O: (T, hd) per head, normalized. - * LSE: (T,) per head — one float per row. + * Epilogue: Read O from TMEM with 16x256b.x1, normalize, write GMEM. * * ================================================================== - * CONSTRAINTS + * TMEM LANE MAPPING (16x256b.x1 format, verified on B200) * ================================================================== - * - T <= 128 (fits in one 128-row MMA tile; longer → D5 KV merge) - * - s_k must be a multiple of 16 (MMA K-tile size) - * - head_dim must be one of: 16, 64, 128, 256 + * 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). * ================================================================== */ @@ -48,27 +33,21 @@ namespace dsv4::kernels::attention { struct FmhaMultiRowParams { - const bf16_t* __restrict__ q; // [batch, n_h, T, hd] - const bf16_t* __restrict__ k; // [batch, n_kv, N, hd] - const bf16_t* __restrict__ v; // [batch, n_kv, hd, N] - bf16_t* __restrict__ o; // [batch, n_h, T, hd] - float* __restrict__ lse; // [batch, n_h, T] + const bf16_t* __restrict__ q; + const bf16_t* __restrict__ k; + const bf16_t* __restrict__ v; + bf16_t* __restrict__ o; + float* __restrict__ lse; - int s_k; // KV sequence length - int T; // Query sequence length (1 for decode, >1 for prefill) - float scale; // 1/sqrt(hd) - int head_dim; // hd + int s_k, T; + float scale; + int head_dim; - int q_head_stride; // T * hd - int q_batch_stride; // n_h * T * hd - int k_head_stride; // N * hd (0 for MQA) - int k_batch_stride; // n_kv * N * hd - int v_head_stride; // hd * N (0 for MQA) - int v_batch_stride; // n_kv * hd * N - int o_head_stride; // T * hd - int o_batch_stride; // n_h * T * hd - int lse_head_stride; // T - int lse_batch_stride; // n_h * T + 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 @@ -80,8 +59,8 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { static constexpr int TILE_SZ = 128 * MMA_K_BF16; static constexpr int V_SUB_SZ = 256; static constexpr int TMEM_N = (HD <= 128) ? 128 : 256; - static constexpr int ROWS_PER_WARP = 32; 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; @@ -89,7 +68,7 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { const int wid = tid / 32; const int lane = tid % 32; - const bool is_softmax_warp = (wid < 4); + const bool is_softmax_warp = (wid == 0); const bool is_mma_warp = (wid == 4); const bool is_load_warp = (wid == 5); @@ -97,11 +76,8 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { const int s_k = params.s_k; const float scale = params.scale; - // Per-warp row range - const int warp_row_start = wid * ROWS_PER_WARP; - const int warp_row_end = min(warp_row_start + ROWS_PER_WARP, T); - const int warp_num_rows = is_softmax_warp ? max(0, warp_row_end - warp_row_start) : 0; - const bool warp_has_rows = (warp_num_rows > 0); + // Per-lane row range + const int lane_row_start = lane * ROWS_PER_LANE; // ================================================================== // Per-head GMEM pointers @@ -128,18 +104,16 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { // ================================================================ extern __shared__ char sbuf[]; uint32_t* sTmemBase = (uint32_t*)sbuf; - // Per-row max/sum: 128 floats each (one per possible row) float* sRowMax = (float*)(sbuf + 4); float* sRowSum = sRowMax + 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); - // s_p_vals: [MAX_ROWS][SK_TILE] = 128 × 128 = 16384 floats = 64KB - float* s_p_vals = (float*)(sV + V_SUB_SZ); + float* s_p_vals = (float*)(sV + V_SUB_SZ); // [MAX_ROWS][SK_TILE] // ================================================================ - // TMEM allocation (warp 4) + // TMEM allocation // ================================================================ if (is_mma_warp) { uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); @@ -149,26 +123,21 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { uint32_t tb = *sTmemBase; // ================================================================ - // QK GEMM loop + // QK GEMM // ================================================================ for (int kt = 0; kt < NKT_QK; kt++) { if (is_load_warp) { constexpr int CORES_MN = 128 / 8; - // Zero Q SMEM for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; - // Load Q (T, hd): write T rows to canonical layout 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 core_k = full_d / 8; - int local_c = full_d % 8; - int core_mn = r / 8; - int local_r = r % 8; + int core_k = full_d / 8, local_c = full_d % 8; + int core_mn = r / 8, local_r = r % 8; sQ0[core_k * CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c] = q_head[r * HD + full_d]; } } - // Load K (s_k, hd) for (int i = lane; i < TILE_SZ; i += 32) sK0[i] = 0; for (int r = 0; r < s_k; r++) { for (int d = lane; d < MMA_K_BF16; d += 32) { @@ -192,100 +161,97 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } // ================================================================ - // Multi-row softmax (Warps 0-3) - // ================================================================ - // Each warp reads its 32 rows from TMEM using row-offset addressing. - // 32x32b.x8: addr = tb + (row_base << 16) + col_group*8 - // Lane l reads row (row_base + l) from TMEM. - // Only active lanes (l < warp_num_rows) do actual computation, - // but ALL 32 lanes must participate in the collective TMEM load. + // Multi-row softmax (16x256b.x1, two-pass) // ================================================================ if (is_softmax_warp) { - uint32_t row_base_addr = tb + (warp_row_start << 16); + float my_row_max[ROWS_PER_LANE]; + for (int i = 0; i < ROWS_PER_LANE; i++) my_row_max[i] = -INFINITY; - // Per-lane: one row of S values - float s_vals[SK_TILE]; - float row_max = -INFINITY; - - // Read S from TMEM: 16 iterations × 8 columns = 128 columns - 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"(row_base_addr + n * 8)); + // 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)); asm volatile("tcgen05.wait::ld.sync.aligned;"); - // All lanes get their row's values; only active lanes compute - if (lane < (unsigned)warp_num_rows) { - 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]); + 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]); } } } - // Per-row max: each lane has its own row's values. - // NO cross-lane reduction — each row's max is independent. - // For T=1, only lane 0 has data; we still skip wmax because - // the single row's max is already in lane 0. - // However, for T=1 backward compat: lane 0 needs its value - // written to sRowMax. No broadcast needed for multi-row. - - if (lane < (unsigned)warp_num_rows) { - sRowMax[warp_row_start + lane] = row_max; - } - - // Compute exp and sum per row (independent per lane) - float row_sum = 0.0f; - if (lane < (unsigned)warp_num_rows) { - for (int j = 0; j < SK_TILE; j++) { - s_vals[j] = expf(s_vals[j] - row_max); - row_sum += s_vals[j]; + // 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]; } } - // No cross-lane sum reduction — each row's sum is in its lane + __syncwarp(); - if (lane < (unsigned)warp_num_rows) { - sRowSum[warp_row_start + lane] = row_sum; + // 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; + } + } } - // Normalize and write P to s_p_vals - if (lane < (unsigned)warp_num_rows) { - float inv_sum = 1.0f / row_sum; - int my_row = warp_row_start + lane; - for (int j = 0; j < SK_TILE; j++) { - s_p_vals[my_row * SK_TILE + j] = s_vals[j] * inv_sum; + // 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; + } } } } __syncthreads(); // ================================================================ - // PV GEMM loop: N=16 sub-tiles × K-tiles - // P (128, s_k) × V (s_k, hd) → O (128, hd) in TMEM + // PV GEMM // ================================================================ 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; - // Fill sPk from s_p_vals: all T rows × 16 cols 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 global_col = kt * MMA_K_BF16 + c; float pval = s_p_vals[r * SK_TILE + global_col]; - int core_mn = r / 8; - int local_r = r % 8; - int core_k = c / 8; - int local_c = c % 8; + int core_mn = r / 8, local_r = r % 8; + int core_k = c / 8, local_c = c % 8; sPk[core_k * CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c] = f32_to_bf16(pval); } } - // Load V sub-tile 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++) { @@ -311,35 +277,36 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } // ================================================================ - // Multi-row epilogue: read O from TMEM, normalize, write to GMEM - // Same row-offset addressing as softmax. + // 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 // ================================================================ if (is_softmax_warp) { - uint32_t row_base_addr = tb + (warp_row_start << 16); - - 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"(row_base_addr + n * 8)); + 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)); asm volatile("tcgen05.wait::ld.sync.aligned;"); - if (lane < (unsigned)warp_num_rows) { - for (int c = 0; c < 8; c++) o_vals[n * 8 + c] = tmp[c]; + + 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 < (unsigned)warp_num_rows) { - int my_row = warp_row_start + lane; - float row_max = sRowMax[my_row]; - float row_sum = sRowSum[my_row]; - float inv_row_sum = 1.0f / row_sum; - for (int d = 0; d < HD; d++) { - o_head[my_row * HD + d] = f32_to_bf16(o_vals[d] * inv_row_sum); - } - if (lse_head) { - lse_head[my_row] = logf(row_sum) + row_max; + // 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]; } } }