biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rewrite multi-row softmax using 16x256b.x1 TMEM reads for proper multi-row access

Browse Source
This commit is contained in:
biondizzle
2026-05-28 20:08:30 +00:00
parent 5e3c61184c
commit aaa76c1af1
1 changed files with 115 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

263
dsv4/kernels/attention/fmha_6warp_multirow.cuh
View File

@@ -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<int HD, int SK_TILE = 128>
@@ -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];
}
}
}