From 278f1b34af1bed5418a98bec7906a1858a41498d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 07:43:40 +0000 Subject: [PATCH] fix(tmem): correct lane-to-position mapping for tcgen05.ld/st MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../kernels/attention/fmha_epilogue_sm100.cuh | 64 +++++++++++-------- 1 file changed, 39 insertions(+), 25 deletions(-) diff --git a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh index 73d6b28d..181eacbc 100644 --- a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh +++ b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh @@ -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);