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

fix: load full Q (128,HD) once before QK loop — not per K-sub-tile

Browse Source 
The MMA expects Q sub-tiles from a full (128,HD) canonical buffer,
but we were only loading (128,16) sub-tiles into a (128,16) buffer.
The MMA descriptor with block_mn=128 describes a (128,128) matrix,
reading 128 columns from SMEM but only 16 had real data.

Now: load all HD/16 TMA tiles of Q into a full (128,HD) canonical
buffer before the QK loop. The MMA reads the kt-th sub-tile via
descriptor offset kt * 128 * 32 bytes.

Also: share single sTmaBuf staging buffer for all TMA loads (Q, K, V).
Removed separate sQ_tma, sK_tma, sV_tma buffers.
This commit is contained in:
biondizzle
2026-05-29 18:28:45 +00:00
parent bf7cf54a51
commit 204cc90808
2 changed files with 47 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
dsv4/kernels/attention/fmha_6warp_tma.cuh
View File

@@ -99,13 +99,10 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
float* sRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float);
off = (off + 127) & ~(size_t)127;
bf16_t* sQ_tma = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t);
bf16_t* sTmaBuf = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t); // TMA staging buffer (128, 16) row-major
off = (off + 127) & ~(size_t)127;
bf16_t* sK_tma = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t);
off = (off + 127) & ~(size_t)127;
bf16_t* sQ = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
bf16_t* sQ = (bf16_t*)(sbuf + off); off += 128 * HD * sizeof(bf16_t); // full Q (128, HD) canonical
off = (off + 127) & ~(size_t)127;
bf16_t* sK = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
@@ -113,9 +110,6 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
off = (off + 127) & ~(size_t)127;
bf16_t* sPk = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
off = (off + 127) & ~(size_t)127;
bf16_t* sV_tma = (bf16_t*)(sbuf + off); off += 16 * 128 * sizeof(bf16_t);
off = (off + 127) & ~(size_t)127;
bf16_t* sV = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
@@ -143,31 +137,52 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
const bool my_row_active = my_warp_active && (my_row < T);
// ==================================================================
// QK GEMM → S in TMEM (loop over K sub-tiles)
// Load full Q into SMEM (128, HD) canonical via TMA
// ==================================================================
// mbarrier phase parity — tracks which phase we're waiting on.
// The mbarrier was initialized with count=1 before the loop.
// After each TMA + arrive.expect_tx + wait, we flip the phase.
// We do NOT re-init the mbarrier inside the loop.
int phase = 0;
for (int kt = 0; kt < NKT_QK; kt++) {
// --- TMA load Q sub-tile ---
// Zero Q canonical buffer first
if (is_load_warp) {
for (int i = lane; i < 128 * HD; i += 32) sQ[i] = 0;
}
__syncthreads();
for (int qkt = 0; qkt < NKT_QK; qkt++) {
if (is_load_warp && lane == 0) {
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sQ_tma);
tma_load_2d(smem_dst, (uint64_t)tma_q, mbar_addr, kt * MMA_K_BF16, 0);
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sTmaBuf);
tma_load_2d(smem_dst, (uint64_t)tma_q, mbar_addr, qkt * MMA_K_BF16, 0);
tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES);
}
tma_mbarrier_wait(mbar_addr, phase);
phase ^= 1;
__syncthreads();
if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sQ, sQ_tma);
// Write (128, 16) row-major TMA buffer into the right position in (128, HD) canonical Q
// The qkt-th (128, 16) sub-tile in canonical = columns [qkt*16, qkt*16+16)
// canonical offset for core_k = qkt*16/8 = qkt*2, same core_mn layout
if (is_load_warp) {
constexpr int CORES_MN = 128 / 8; // 16
constexpr int CORES_K_SUB = 16 / 8; // 2
constexpr int SUB_TOTAL = 128 * 16;
for (int i = lane; i < SUB_TOTAL; i += 32) {
int r = i / 16, c = i % 16;
int core_mn = r / 8, local_r = r % 8;
int core_k_sub = c / 8, local_c = c % 8;
int core_k_full = qkt * 2 + core_k_sub;
int dst_idx = core_k_full * CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c;
sQ[dst_idx] = sTmaBuf[i];
}
}
__syncthreads();
}
// ==================================================================
// QK GEMM → S in TMEM (loop over K sub-tiles)
// ==================================================================
for (int kt = 0; kt < NKT_QK; kt++) {
// --- TMA load K sub-tile ---
if (is_load_warp && lane == 0) {
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sK_tma);
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sTmaBuf);
tma_load_2d(smem_dst, (uint64_t)tma_k, mbar_addr, kt * MMA_K_BF16, 0);
tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES);
}
@@ -175,13 +190,15 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
phase ^= 1;
__syncthreads();
if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sK, sK_tma);
if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sK, sTmaBuf);
__syncthreads();
// MMA: sQ × sK → TMEM
// MMA: Q sub-tile × K sub-tile → TMEM
// Q's kt-th sub-tile starts at offset kt * 128 * 32 bytes in canonical SMEM
if (is_mma_warp) {
uint32_t idesc = make_idesc(128, 128);
uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ), 128);
uint32_t sq_kt = (uint32_t)__cvta_generic_to_shared(sQ) + kt * 128 * 32;
uint64_t dq = make_umma_desc_kmajor_none(sq_kt, 128);
uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK), 128);
if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0);
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
@@ -262,9 +279,8 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
}
__syncthreads();
// --- TMA load V sub-tile ---
if (is_load_warp && lane == 0) {
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sV_tma);
uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sTmaBuf);
// V is (HD, s_k). TMA 2D: coord {col_start, d_base}
tma_load_2d(smem_dst, (uint64_t)tma_v, mbar_addr, col_start, d_base);
tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES);
@@ -273,7 +289,7 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
phase ^= 1;
__syncthreads();
// Transpose sV_tma (16, 128) → sV (128, 16) canonical
// Transpose sTmaBuf (16, 128) → sV (128, 16) canonical
if (is_load_warp) {
constexpr int SV_CORES_MN = 128 / 8;
for (int i = lane; i < TILE_SZ; i += 32) sV[i] = 0;
@@ -282,7 +298,7 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
int core_mn = r / 8, local_r = r % 8;
int core_k = d / 8, local_c = d % 8;
int dst_idx = core_k * SV_CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c;
sV[dst_idx] = sV_tma[i];
sV[dst_idx] = sTmaBuf[i];
}
}
__syncthreads();

14
tests/unit/test_fmha_tma.cu
View File

@@ -40,19 +40,15 @@ static int compute_smem_tma() {
off += MAX_T * sizeof(float); // sRowMax
off += MAX_T * sizeof(float); // sRowSum
off = (off + 127) & ~(size_t)127;
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sQ_tma
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sTmaBuf (TMA staging)
off = (off + 127) & ~(size_t)127;
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sK_tma
off += 128 * HD * sizeof(bf16_t); // sQ full (128, HD) canonical
off = (off + 127) & ~(size_t)127;
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sQ canonical
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sK canonical (128, 16)
off = (off + 127) & ~(size_t)127;
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sK canonical
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sPk canonical (128, 16)
off = (off + 127) & ~(size_t)127;
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sPk canonical
off = (off + 127) & ~(size_t)127;
off += 16 * 128 * sizeof(bf16_t); // sV_tma
off = (off + 127) & ~(size_t)127;
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sV canonical
off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sV canonical (128, 16)
return (int)off;
}