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

Fix multi-row softmax v3: 32x32b.x8 with per-lane per-row (no wmax/wsum), per-row sRowMax/sRowSum arrays

Browse Source
This commit is contained in:
biondizzle
2026-05-28 20:10:13 +00:00
parent aaa76c1af1
commit 08694b8136
2 changed files with 62 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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;
}
}
}

2
tests/unit/test_fmha_6warp_multirow.cu
View File

@@ -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) {