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

fmha_multirow: use natural 4-warp TMEM partitioning after UMMA

Browse Source 
After UMMA (QK GEMM), 4 warps reading TMEM with 32x32b.x8 each
see a different 32-row partition (verified on B200):
  Warp 0 → rows 0-31, Warp 1 → rows 32-63, etc.
Lane l in warp w reads row w*32 + l.

This eliminates the broken row_page<<16 addressing and allows:
- T<=32: warp 0 only, 32x32b.x8, each lane = one row
- T>32: 4 warps, each reads its natural 32-row partition
- Epilogue: same partitioning for reading O from TMEM

No s_p_vals buffer. P streamed per K-tile through sPk.
This commit is contained in:
biondizzle
2026-05-28 23:07:31 +00:00
parent ff8c677486
commit be2685e9e3
1 changed files with 43 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9,27 +9,27 @@
* 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.
* 16x256b.x1 stores crash on 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)
* KEY DISCOVERY (verified on B200):
* After UMMA (QK GEMM), 4 warps reading TMEM with 32x32b.x8 each
* see a DIFFERENT 32-row partition of the S matrix:
* Warp 0 → rows 0-31, Warp 1 → rows 32-63,
* Warp 2 → rows 64-95, Warp 3 → rows 96-127.
* Lane l in warp w reads row w*32 + l across 8 columns.
* This is because the MMA writes Layout D to TMEM, and each warp's
* 32x32b.x8 read naturally accesses its partition.
*
* 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.
* - T <= 32: only warp 0 active. Lane l handles row l.
* - T > 32: 4 softmax warps. Warp w handles rows [w*32, min((w+1)*32, T)).
* Lane l in warp w handles row w*32 + l.
* Two-pass: (1) row_max over all S columns, (2) exp/sum, stream P to sPk.
*
* 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).
* P written one K-tile at a time to sPk (128, 16) SMEM in canonical layout.
* Softmax warps compute P in registers, write to sPk. Then MMA does PV.
* Softmax and PV interleaved per K-tile (can't hold full 128×128 P in SMEM).
*
* PV GEMM:
* N=16 sub-tiles (avoids Layout D N=64 bug). SS MMA: sPk(128,16) × sV(16,16).
@@ -37,20 +37,8 @@
*
* EPILOGUE:
* 32x32b.x8 TMEM read → normalize (1/row_sum) → BF16 → GMEM.
* Each softmax warp reads its own 32-row partition of O from TMEM.
* 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
* ==================================================================
*/
@@ -80,14 +68,14 @@ struct FmhaMultiRowParams {
template<int HD, int SK_TILE = 128>
__global__ void __launch_bounds__(192)
fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
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 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 = 16 * MMA_K_BF16;
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
static constexpr int CORES_MN = 128 / 8;
const int head_idx = blockIdx.y;
const int batch_idx = blockIdx.z;
@@ -101,7 +89,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
const int s_k = params.s_k;
const float scale = params.scale;
// 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;
@@ -112,18 +99,15 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
// SMEM allocation — NO s_p_vals buffer
// ================================================================
extern __shared__ char sbuf[];
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
uint32_t* sTmemBase = (uint32_t*)sbuf;
float* sRowMax = (float*)(sbuf + 8);
float* sRowSum = sRowMax + MAX_ROWS;
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
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);
// ================================================================
// TMEM allocation (warp 4)
// ================================================================
if (is_mma_warp) {
uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
tmem_alloc(smem_ptr, TMEM_N);
@@ -134,21 +118,17 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
// ================================================================
// 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
// After UMMA, warp w's 32x32b.x8 read gives rows [w*32, w*32+31].
// Lane l in warp w reads 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 int my_row = my_warp_active ? (wid * 32 + lane) : 0;
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) {
// 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) {
@@ -160,7 +140,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
}
}
}
// 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) {
@@ -185,26 +164,19 @@ 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 PASS 1: compute row_max for each row
// ================================================================
// 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
const int NUM_READS = SK_TILE / 8;
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));
: "r"(tb + n * 8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (my_row_active) {
@@ -219,47 +191,32 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
}
}
// Store row_max to SMEM
if (my_row_active) sRowMax[my_row] = my_row_max;
__syncthreads();
// ================================================================
// SOFTMAX PASS 2 + PV GEMM (interleaved per PV K-tile)
// ================================================================
// 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
float p_vals[16];
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));
: "r"(tb + n * 8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (my_row_active) {
@@ -275,13 +232,13 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
}
}
// Zero sPk first (load warp does this for all rows)
// Zero sPk
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)
// Write P to sPk in canonical layout
if (my_row_active) {
for (int c = 0; c < MMA_K_BF16; c++) {
int ck = c / 8, lc = c % 8;
@@ -296,10 +253,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
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++) {
@@ -318,9 +271,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
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);
// 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");
}
@@ -328,33 +278,27 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
}
}
// Store row_sum to SMEM
if (my_row_active) sRowSum[my_row] = my_row_sum;
__syncthreads();
// ================================================================
// EPILOGUE: TMEM → regs → normalize → BF16 → GMEM
// ================================================================
// After PV GEMM, O[128, HD] is in TMEM. Each softmax warp reads
// its own 32-row partition (same as after QK GEMM).
// HD/8 = N_NSUB*2 read groups.
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"(addr));
: "r"(tb + n * 8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
for (int c = 0; c < 8; c++) {
@@ -365,7 +309,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
}
}
// LSE
if (lse_head) {
lse_head[my_row] = logf(rs) + rm;
}
@@ -373,10 +316,7 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
}
__syncthreads();
// TMEM dealloc (warp 4)
if (is_mma_warp) {
tmem_dealloc(tb, TMEM_N);
}
if (is_mma_warp) tmem_dealloc(tb, TMEM_N);
}
} // namespace dsv4::kernels::attention