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

fix(tmem): correct lane-to-position mapping for tcgen05.ld/st

Browse Source 
CRITICAL FIX: tcgen05.st 16x256b.x1.b32 is warp-collective where:
- Lane i writes to positions i*4+0..i*4+3 within the column
- 32 lanes × 4 FP32 = 128 FP32 per column
- For row 0: lane 0 = positions 0-3, lane 1 = 4-7, ..., lane 31 = 124-127

Old code iterated col = lane; col < N; col += 32, treating each lane
as owning a separate column. That was WRONG — all 32 lanes share each
column, each owning 4 positions within it.

New code: HD values need ceil(HD/128) columns. Lane i writes
sPvBuf[i*4+0..3] to column 0 (or column 1 for HD > 128).

Verified via test_tmem_lane_mapping.cu on B200.
This commit is contained in:
biondizzle
2026-05-28 07:43:40 +00:00
parent 593bc25afa
commit 278f1b34af
1 changed files with 39 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
dsv4/kernels/attention/fmha_epilogue_sm100.cuh
View File

@@ -70,14 +70,16 @@ fmha_decode_tmem(
const bf16_t* vb = v + batch*bstride_kv;
bf16_t* oh = o + batch*bstride_o + head*HD;
// TMEM column layout: each column holds 4 FP32 values for row-group 0
// (lane 0 gets rows 0-3, lane 1 gets rows 4-7, etc. — but for T=1 decode
// only row 0 matters, so only lane 0's 4 values are meaningful).
constexpr int TMEM_O_COLS = (HD + 3) / 4;
// TMEM layout (verified on B200 via test_tmem_lane_mapping.cu):
// tcgen05.st 16x256b.x1.b32: each lane i writes 4 FP32 to positions
// i*4+0..i*4+3 within the column. 32 lanes × 4 = 128 FP32 per column.
// For row 0, lane 0 = positions 0-3, lane 1 = 4-7, ..., lane 31 = 124-127.
// So HD values need ceil(HD/128) columns.
constexpr int TMEM_COLS_NEEDED = (HD + 127) / 128; // 1 for HD<=128, 2 for HD<=256
// tcgen05.alloc requires power-of-2 columns, minimum 32
constexpr int TMEM_N = TMEM_O_COLS <= 32 ? 32 :
(TMEM_O_COLS <= 64 ? 64 :
(TMEM_O_COLS <= 128 ? 128 : 256));
constexpr int TMEM_N = TMEM_COLS_NEEDED <= 32 ? 32 :
(TMEM_COLS_NEEDED <= 64 ? 64 :
(TMEM_COLS_NEEDED <= 128 ? 128 : 256));
// SMEM layout:
// [0..3] tmem_base (written by tcgen05.alloc)
@@ -105,11 +107,13 @@ fmha_decode_tmem(
// Read tmem_base from SMEM (written by alloc)
uint32_t tmem_base = *sTmemBase;
if (tid == 0) printf("[tmem] base=%u, n=%d, o_cols=%d\n", tmem_base, TMEM_N, TMEM_O_COLS);
if (tid == 0) printf("[tmem] base=%u, alloc_n=%d, cols_needed=%d\n", tmem_base, TMEM_N, TMEM_COLS_NEEDED);
// Initialize TMEM O to zero — warp-collective
// Use TMEM_COLS_NEEDED columns, each zeroed by all 32 lanes writing 0.
constexpr int TMEM_COLS_NEEDED = (HD + 127) / 128;
if (wid == 0) {
for (int col = lane; col < TMEM_N; col += WARP) {
for (int col = 0; col < TMEM_COLS_NEEDED; col++) {
tmem_store(tmem_base + col, 0, 0, 0, 0);
}
tmem_fence_store();
@@ -140,7 +144,7 @@ fmha_decode_tmem(
// D1.5: Rescale O in TMEM
// Write rescale factor to SMEM so warp 0 can use it collectively
sPvBuf[0] = rescale;
sPvBuf[1] = (float)TMEM_O_COLS; // signal: >0 means rescale
sPvBuf[1] = 1.0f; // signal: rescale needed
// Wake warp 0 to do the rescale
__threadfence_block(); // ensure SMEM writes visible
@@ -219,16 +223,24 @@ fmha_decode_tmem(
// ================================================================
// Step 1: Write SMEM accumulator to TMEM (warp 0, warp-collective)
// IMPORTANT: ALL 32 lanes must call tmem_store (warp-collective).
// With TMEM_O_COLS=16 and WARP=32, only lanes 0-15 would enter the loop.
// Fix: loop over TMEM_N (always >= 32) so all lanes participate.
// Lanes writing beyond TMEM_O_COLS write don't-care data to don't-care columns.
//
// TMEM lane mapping (verified on B200 via test_tmem_lane_mapping.cu):
// tcgen05.st 16x256b.x1.b32: each lane i writes 4 FP32 to positions
// i*4+0..i*4+3 within the column. 32 lanes × 4 = 128 FP32 per column.
// For row 0, lane 0 = positions 0-3, lane 1 = 4-7, ..., lane 31 = 124-127.
//
// So for HD<=128, a single column suffices (128 FP32 per col, row 0).
// Lane i writes sPvBuf[i*4+0..3] to column 0.
// Lanes with i*4 >= HD write zeros (don't-care, but must participate).
//
constexpr int TMEM_COLS_NEEDED = (HD + 127) / 128; // 1 for HD<=128, 2 for HD<=256
if (wid == 0) {
for (int col = lane; col < TMEM_N; col += WARP) {
int d0 = col * 4 + 0;
int d1 = col * 4 + 1;
int d2 = col * 4 + 2;
int d3 = col * 4 + 3;
for (int col = 0; col < TMEM_COLS_NEEDED; col++) {
int base = col * 128; // FP32 offset for this column
int d0 = base + lane * 4 + 0;
int d1 = base + lane * 4 + 1;
int d2 = base + lane * 4 + 2;
int d3 = base + lane * 4 + 3;
uint32_t u0 = (d0 < HD) ? f32_to_u32(sPvBuf[d0]) : 0;
uint32_t u1 = (d1 < HD) ? f32_to_u32(sPvBuf[d1]) : 0;
@@ -243,11 +255,12 @@ fmha_decode_tmem(
if (tid == 0) printf("[tmem] wrote to TMEM OK\n");
// Step 2: Read from TMEM to registers (warp 0, warp-collective)
// Same warp-collective constraint: ALL 32 lanes must call tmem_load.
// Same lane mapping: lane i reads positions i*4+0..3 from the column.
constexpr int TMEM_COLS_NEEDED = (HD + 127) / 128;
if (wid == 0) {
float inv_sum = 1.0f / sRowSums[0];
for (int col = lane; col < TMEM_N; col += WARP) {
for (int col = 0; col < TMEM_COLS_NEEDED; col++) {
uint32_t u0, u1, u2, u3;
tmem_load(tmem_base + col, u0, u1, u2, u3);
@@ -258,10 +271,11 @@ fmha_decode_tmem(
float r3 = u32_to_f32(u3) * inv_sum;
// Step 4: Cast to BF16 and write to GMEM
int d0 = col * 4 + 0;
int d1 = col * 4 + 1;
int d2 = col * 4 + 2;
int d3 = col * 4 + 3;
int base = col * 128;
int d0 = base + lane * 4 + 0;
int d1 = base + lane * 4 + 1;
int d2 = base + lane * 4 + 2;
int d3 = base + lane * 4 + 3;
if (d0 < HD) oh[d0] = f32_to_bf16(r0);
if (d1 < HD) oh[d1] = f32_to_bf16(r1);