feat: double-buffer TMA pipeline for K loads in single-tile kernel
This commit is contained in:
@@ -1,39 +1,28 @@
|
||||
/**
|
||||
* DSV4 FMHA — 6-warp specialized kernel with TMA async loads for Blackwell SM100.
|
||||
* DSV4 FMHA — 6-warp specialized kernel with double-buffer TMA pipeline.
|
||||
*
|
||||
* ==================================================================
|
||||
* WARP SPECIALIZATION (same as fmha_6warp.cuh)
|
||||
* DOUBLE-BUFFER TMA PIPELINE
|
||||
* ==================================================================
|
||||
* Warp 0-3 (tid 0-127): Softmax + correction + epilogue
|
||||
* Warp 4 (tid 128-159): MMA (QK + PV)
|
||||
* Warp 5 (tid 160-191): TMA loads (K, V) + Q direct load + P fill
|
||||
*
|
||||
* ==================================================================
|
||||
* TMA LOAD PIPELINE
|
||||
* ==================================================================
|
||||
* - Q: direct GMEM load (T=1 decode, only 1 row; small enough to not benefit from TMA)
|
||||
* - K: TMA async load via cp.async.bulk.tensor.2d with mbarrier completion
|
||||
* - V: TMA async load (future; currently direct GMEM load)
|
||||
* - mbarrier: init once, arrive.expect_tx after TMA issue, phase parity tracking
|
||||
* - ALL threads wait on mbarrier (works correctly at 192 threads, verified on B200)
|
||||
* K is loaded via TMA with double-buffering:
|
||||
* - Two canonical SMEM buffers: sK0 and sK1
|
||||
* - Preload kt=0 into sK0
|
||||
* - Loop: issue TMA for kt+1 → sTmaBuf (while MMA runs on current sK),
|
||||
* then convert sTmaBuf → next sK, swap buffers
|
||||
* - TMA DMA transfer overlaps with MMA compute
|
||||
*
|
||||
* ==================================================================
|
||||
* SMEM LAYOUT
|
||||
* ==================================================================
|
||||
* All SMEM regions are 128-byte aligned for TMA compatibility.
|
||||
* sTmaBuf is used as a staging area: TMA writes row-major data here,
|
||||
* then it's converted to canonical layout for MMA.
|
||||
* The pipeline stages for kt=N:
|
||||
* 1. Issue TMA load for kt=N+1 (if exists) → sTmaBuf
|
||||
* 2. MMA QK using sK[current] (runs while TMA DMA is in flight)
|
||||
* 3. Wait for TMA completion (usually already done after MMA)
|
||||
* 4. Convert sTmaBuf → sK[next] (canonical layout)
|
||||
* 5. Swap current/next buffer pointers
|
||||
*
|
||||
* sTmemBase: 4 bytes
|
||||
* sMbar: 16 bytes (128-byte aligned)
|
||||
* sTmaBuf: TILE_SZ * 2 bytes (128-byte aligned) — TMA staging buffer
|
||||
* sQ0: TILE_SZ * 2 bytes (128-byte aligned) — canonical for MMA
|
||||
* sK0: TILE_SZ * 2 bytes (128-byte aligned) — canonical for MMA
|
||||
* sPk: TILE_SZ * 2 bytes (128-byte aligned) — canonical for PV
|
||||
* sV: V_SUB_SZ bytes (128-byte aligned) — canonical for PV
|
||||
* sRowMax: MAX_ROWS * 4 bytes
|
||||
* sRowSum: MAX_ROWS * 4 bytes
|
||||
* s_p_vals: SK_TILE * 4 bytes
|
||||
* V is also loaded via TMA but NOT double-buffered (V tiles are tiny:
|
||||
* 16×16 = 512 bytes, TMA latency is negligible).
|
||||
*
|
||||
* Q is loaded directly from GMEM (T=1 decode, only 1 row).
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
@@ -52,17 +41,16 @@ fmha_6warp_tma_kernel(
|
||||
CUtensorMap* __restrict__ tma_v,
|
||||
bf16_t* __restrict__ o,
|
||||
float* __restrict__ lse,
|
||||
int s_k,
|
||||
float scale
|
||||
int s_k, float scale
|
||||
) {
|
||||
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; // (16,16) canonical
|
||||
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
|
||||
static constexpr int CORES_MN = 128 / 8;
|
||||
static constexpr int NUM_READS = SK_TILE / 8;
|
||||
static constexpr int TMA_TILE_BYTES = TILE_SZ * sizeof(bf16_t);
|
||||
|
||||
@@ -70,35 +58,31 @@ fmha_6warp_tma_kernel(
|
||||
const int wid = tid / 32;
|
||||
const int lane = tid % 32;
|
||||
|
||||
// Warp role predicates
|
||||
const bool is_softmax_warp = (wid < 4);
|
||||
const bool is_mma_warp = (wid == 4);
|
||||
const bool is_load_warp = (wid == 5);
|
||||
|
||||
// ================================================================
|
||||
// SMEM allocation — all 128-byte aligned
|
||||
// SMEM allocation — 128-byte aligned
|
||||
// ================================================================
|
||||
extern __shared__ __align__(128) char sbuf[];
|
||||
size_t off = 0;
|
||||
|
||||
uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off += 4;
|
||||
off = (off + 127) & ~(size_t)127; // 128-byte align
|
||||
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
uint64_t* sMbar = (uint64_t*)(sbuf + off); off += 16;
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
bf16_t* sTmaBuf = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
bf16_t* sQ0 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
bf16_t* sK0 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
bf16_t* sK1 = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
|
||||
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 = (bf16_t*)(sbuf + off); off += V_SUB_SZ * sizeof(bf16_t);
|
||||
|
||||
float* sRowMax = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float);
|
||||
@@ -106,11 +90,9 @@ fmha_6warp_tma_kernel(
|
||||
float* s_p_vals = (float*)(sbuf + off); off += SK_TILE * sizeof(float);
|
||||
|
||||
// ================================================================
|
||||
// Initialize TMEM + mbarrier
|
||||
// Init
|
||||
// ================================================================
|
||||
if (is_mma_warp) {
|
||||
tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N);
|
||||
}
|
||||
if (is_mma_warp) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N);
|
||||
if (tid == 0) {
|
||||
tma_mbarrier_init((uint32_t)__cvta_generic_to_shared(sMbar), 1);
|
||||
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
|
||||
@@ -121,51 +103,98 @@ fmha_6warp_tma_kernel(
|
||||
int phase = 0;
|
||||
|
||||
// ================================================================
|
||||
// QK GEMM loop: TMA for K, direct load for Q
|
||||
// QK GEMM — double-buffer TMA pipeline
|
||||
// ================================================================
|
||||
for (int kt = 0; kt < NKT_QK; kt++) {
|
||||
// ---- Warp 5: Load Q K-tile directly from GMEM ----
|
||||
if (is_load_warp) {
|
||||
for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0;
|
||||
for (int d = lane; d < MMA_K_BF16; d += 32) {
|
||||
int full_d = kt * MMA_K_BF16 + d;
|
||||
if (full_d < HD) {
|
||||
int ck = d / 8, lc = d % 8;
|
||||
sQ0[ck * CORES_MN * 64 + lc] = q[full_d];
|
||||
}
|
||||
}
|
||||
}
|
||||
{
|
||||
uint32_t idesc = make_idesc(128, 128);
|
||||
bf16_t* sK_bufs[2] = {sK0, sK1};
|
||||
int cur_buf = 0; // alternates 0,1,0,1,...
|
||||
|
||||
// ---- TMA load K sub-tile ----
|
||||
// Only 1 thread issues TMA (warp 5, lane 0)
|
||||
// --- Preload kt=0: TMA → sTmaBuf → convert → sK0 ---
|
||||
if (is_load_warp && lane == 0) {
|
||||
tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_k,
|
||||
mbar_addr, kt * MMA_K_BF16, 0);
|
||||
mbar_addr, 0, 0);
|
||||
tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES);
|
||||
}
|
||||
// ALL threads wait for TMA completion
|
||||
tma_mbarrier_wait(mbar_addr, phase); phase ^= 1;
|
||||
__syncthreads();
|
||||
|
||||
// ---- Convert sTmaBuf (row-major) → sK0 (canonical) ----
|
||||
// All threads participate in the conversion
|
||||
// Convert sTmaBuf → sK0 (canonical)
|
||||
for (int i = tid; i < TILE_SZ; i += 192) sK0[i] = 0;
|
||||
for (int i = tid; i < s_k * MMA_K_BF16; i += 192) {
|
||||
int r = i / MMA_K_BF16, c = i % MMA_K_BF16;
|
||||
int ck = c / 8, lc = c % 8, tmn = r / 8, lr = r % 8;
|
||||
sK0[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = sTmaBuf[i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ---- Warp 4: QK MMA ----
|
||||
if (is_mma_warp) {
|
||||
uint32_t idesc = make_idesc(128, 128);
|
||||
uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ0), 128);
|
||||
uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK0), 128);
|
||||
if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0);
|
||||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||||
// Also load Q for kt=0
|
||||
if (is_load_warp) {
|
||||
for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0;
|
||||
for (int d = lane; d < MMA_K_BF16; d += 32) {
|
||||
if (d < HD) {
|
||||
int ck = d / 8, lc = d % 8;
|
||||
sQ0[ck * CORES_MN * 64 + lc] = q[d];
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// --- Pipeline loop ---
|
||||
for (int kt = 0; kt < NKT_QK; kt++) {
|
||||
bf16_t* cur_sK = sK_bufs[cur_buf];
|
||||
|
||||
// 1. Issue TMA for kt+1 (if exists) → sTmaBuf
|
||||
bool has_next = (kt + 1 < NKT_QK);
|
||||
if (has_next && is_load_warp && lane == 0) {
|
||||
tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_k,
|
||||
mbar_addr, (kt + 1) * MMA_K_BF16, 0);
|
||||
tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES);
|
||||
}
|
||||
|
||||
// 2. MMA QK using cur_sK (TMA DMA runs in parallel!)
|
||||
if (is_mma_warp) {
|
||||
uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ0), 128);
|
||||
uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(cur_sK), 128);
|
||||
if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0);
|
||||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||||
}
|
||||
// Note: no __syncthreads() here — MMA and TMA can overlap across warps.
|
||||
// But we need the MMA to complete before reading TMEM (softmax).
|
||||
// The softmax warps will sync later.
|
||||
|
||||
// 3. Wait for TMA kt+1 completion (usually already done after MMA)
|
||||
if (has_next) {
|
||||
tma_mbarrier_wait(mbar_addr, phase); phase ^= 1;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 4. Convert sTmaBuf → next sK buffer (canonical)
|
||||
if (has_next) {
|
||||
bf16_t* next_sK = sK_bufs[1 - cur_buf];
|
||||
for (int i = tid; i < TILE_SZ; i += 192) next_sK[i] = 0;
|
||||
for (int i = tid; i < s_k * MMA_K_BF16; i += 192) {
|
||||
int r = i / MMA_K_BF16, c = i % MMA_K_BF16;
|
||||
int ck = c / 8, lc = c % 8, tmn = r / 8, lr = r % 8;
|
||||
next_sK[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = sTmaBuf[i];
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Load Q for next kt (if exists)
|
||||
if (has_next && is_load_warp) {
|
||||
for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0;
|
||||
for (int d = lane; d < MMA_K_BF16; d += 32) {
|
||||
int full_d = (kt + 1) * MMA_K_BF16 + d;
|
||||
if (full_d < HD) {
|
||||
int ck = d / 8, lc = d % 8;
|
||||
sQ0[ck * CORES_MN * 64 + lc] = q[full_d];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 6. Swap buffers
|
||||
cur_buf = 1 - cur_buf;
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
// ================================================================
|
||||
@@ -200,27 +229,29 @@ fmha_6warp_tma_kernel(
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// PV GEMM loop: N=16 sub-tiles × K-tiles
|
||||
// PV GEMM loop: N=16 sub-tiles × K-tiles, V via TMA
|
||||
// ================================================================
|
||||
for (int n = 0; n < N_NSUB; n++) {
|
||||
int d_base = n * 16;
|
||||
|
||||
for (int kt = 0; kt < NKT_PV; kt++) {
|
||||
// ---- Fill sPk ----
|
||||
const int col_start = kt * MMA_K_BF16;
|
||||
|
||||
// Fill sPk
|
||||
if (is_load_warp) {
|
||||
for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0;
|
||||
if (lane < 16) {
|
||||
int c = lane;
|
||||
int ck = c / 8, lc = c % 8;
|
||||
sPk[ck * CORES_MN * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[kt * MMA_K_BF16 + c]);
|
||||
sPk[ck * CORES_MN * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[col_start + c]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ---- Load V via TMA ----
|
||||
// Load V via TMA
|
||||
if (is_load_warp && lane == 0) {
|
||||
tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_v,
|
||||
mbar_addr, kt * MMA_K_BF16, d_base);
|
||||
mbar_addr, col_start, d_base);
|
||||
tma_mbarrier_arrive_expect_tx(mbar_addr, V_SUB_SZ * sizeof(bf16_t));
|
||||
}
|
||||
tma_mbarrier_wait(mbar_addr, phase); phase ^= 1;
|
||||
@@ -235,7 +266,7 @@ fmha_6warp_tma_kernel(
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ---- Warp 4: PV MMA ----
|
||||
// MMA
|
||||
if (is_mma_warp) {
|
||||
uint32_t idesc_pv16 = make_idesc(128, 16);
|
||||
uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), 128);
|
||||
@@ -248,7 +279,7 @@ fmha_6warp_tma_kernel(
|
||||
}
|
||||
|
||||
// ================================================================
|
||||
// Epilogue: TMEM → regs → normalize → BF16 → GMEM (warp 0)
|
||||
// Epilogue: TMEM → regs → BF16 → GMEM (warp 0)
|
||||
// ================================================================
|
||||
if (wid == 0) {
|
||||
float o_vals[HD];
|
||||
@@ -266,10 +297,7 @@ fmha_6warp_tma_kernel(
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// TMEM dealloc (warp 4)
|
||||
if (is_mma_warp) {
|
||||
tmem_dealloc(tb, TMEM_N);
|
||||
}
|
||||
if (is_mma_warp) tmem_dealloc(tb, TMEM_N);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
@@ -44,6 +44,8 @@ static size_t compute_smem() {
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
off += TILE_SZ * sizeof(bf16_t); // sK0
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
off += TILE_SZ * sizeof(bf16_t); // sK1 (double buffer)
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
off += TILE_SZ * sizeof(bf16_t); // sPk
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
off += V_SUB_SZ * sizeof(bf16_t); // sV
|
||||
|
||||
Reference in New Issue
Block a user