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

P6: One-way TMEM→regs→SMEM→TMA store epilogue

Browse Source 
- fmha_6warp_multihead.cuh: Rewritten epilogue with proper Blackwell pipeline
  1. TMEM → regs (tcgen05.ld, warp-collective)
  2. epilogue_op in regs (normalize, FP4 hook via ENABLE_FP4_EPILOGUE)
  3. regs → SMEM row-major (sO_epi, for TMA tile format)
  4. TMA store SMEM → GMEM (async, enables multi-CTA)
  Fallback to direct GMEM write when tma_o is nullptr.
  Added FmhaParams.tma_o field and ENABLE_FP4_EPILOGUE template param.

- fmha_6warp_tma_multirow_multitile.cuh: Same epilogue pattern for multi-tile.
  Writes normalized output to sO_epi_rowmajor + TMA store (or direct GMEM).
  Added tma_o to FmhaTmaMultiRowMultiTileParams.

- fmha_tma.cuh: Added tma_store_2d and tma_store_wait for async GMEM writes.

- fmha_multihead_capi.cu: Added fmha_multihead_decode_tma_launch with
  per-(head,batch) TMA descriptors. Updated SMEM size calculation for sO_epi + sMbarStore.

- fmha_multitile_capi.cu: Added tma_o=nullptr (backward compatible), updated SMEM size.
This commit is contained in:
biondizzle
2026-05-30 16:56:07 +00:00
parent 05b5bf9db1
commit 212fc85627
5 changed files with 354 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

186
dsv4/kernels/attention/fmha_6warp_multihead.cuh
View File

@@ -12,42 +12,27 @@
* No cross-CTA synchronization required.
*
* ==================================================================
* MQA / GQA SUPPORT
* EPILOGUE (P6 — One-way TMEM → regs → SMEM → TMA store → GMEM)
* ==================================================================
* - MQA: all Q heads share one KV head. Pass k_head_stride=0, v_head_stride=0
* so all CTAs read the same K/V.
* - GQA: groups of Q heads share a KV head. The caller must arrange
* K/V tensors so that k_head_stride/v_head_stride map correctly.
* - MHA: k_head_stride = k_row_stride * N, same for V.
* The proper Blackwell output pipeline:
* 1. TMEM → registers (tcgen05.ld, warp-collective)
* 2. epilogue_op in registers (normalize + optional FP4 pack)
* 3. Registers → SMEM (row-major, matching TMA tile format)
* 4. TMA store SMEM → GMEM (async, enables multi-CTA)
*
* This replaces the old direct GMEM write and unblocks:
* - D2 multi-CTA grid (TMA store with flat_divide coords)
* - NVFP4-1.2 FP4 output fusion (register slot for amax + pack)
* - Proper async pipeline overlap
*
* When tma_o is nullptr, falls back to direct GMEM write from registers.
* When tma_o is set, uses the proper TMA store pipeline.
*
* ==================================================================
* TENSOR LAYOUTS (GMEM)
* OUTPUT
* ==================================================================
* Q: [batch, n_h, T, hd] — head stride = T * hd, batch stride = n_h * T * hd
* K: [batch, n_kv, N, hd] — head stride = N * hd (or 0 for MQA)
* V: [batch, n_kv, hd, N] — head stride = hd * N (or 0 for MQA)
* O: [batch, n_h, T, hd] — same strides as Q
*
* For decode (T=1): q_head_offset = blockIdx.y * hd, q_batch_offset = blockIdx.z * n_h * hd
* For prefill (T>1): head-packed M = T rows per head (must fit in 128-row MMA tile)
*
* ==================================================================
* SOFTMAX ROWS
* ==================================================================
* T=1 decode: only row 0 of the 128-row MMA tile has data. Only warp 0
* computes softmax for row 0.
* T>1 prefill: rows 0..T-1 have data. All 4 softmax warps process
* rows in parallel (warp w handles rows [w*32, (w+1)*32) ∩ [0, T)).
* This is Milestone 4 territory — current implementation handles T=1 only.
* The multi-head grid layout is independent of multi-row softmax and
* can land first.
*
* ==================================================================
* OUTPUT: UN-NORMALIZED O + LSE
* ==================================================================
* The kernel emits un-normalized O and per-row LSE for composition with
* D5 multi-tile KV merge. External code normalizes: O_norm = O / row_sum.
* For single-segment decode, normalization is done in the epilogue.
* For single-segment decode: normalized O written to GMEM + LSE.
* For multi-segment: un-normalized O + LSE for external merge.
* LSE layout: [batch, n_h, T] — one float per head per row.
*/
@@ -55,6 +40,7 @@
#include "fmha_common.cuh"
#include "fmha_umma_desc.cuh"
#include "fmha_tma.cuh"
namespace dsv4::kernels::attention {
@@ -79,15 +65,19 @@ struct FmhaParams {
int q_batch_stride; // stride between Q batch items = n_h * T * hd
int k_head_stride; // stride between K heads = N * hd (0 for MQA)
int k_batch_stride; // stride between K batch items = n_kv * N * hd
int v_head_stride; // stride between V heads = hd * N (0 for MQA)
int v_head_stride; // stride between V heads = hd * N (or 0 for MQA)
int v_batch_stride; // stride between V batch items = n_kv * hd * N
int o_head_stride; // stride between O heads = T * hd
int o_batch_stride; // stride between O batch items = n_h * T * hd
int lse_head_stride; // stride between LSE heads = T
int lse_batch_stride; // stride between LSE batch items = n_h * T
// TMA descriptor for O output (device pointer). When nullptr, uses
// direct GMEM write. When set, uses TMA store for async output.
CUtensorMap* __restrict__ tma_o;
};
template<int HD, int SK_TILE = 128>
template<int HD, int SK_TILE = 128, bool ENABLE_FP4_EPILOGUE = false>
__global__ void __launch_bounds__(192)
fmha_6warp_multihead_kernel(FmhaParams params) {
static constexpr int NKT_QK = HD / MMA_K_BF16;
@@ -96,6 +86,7 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
static constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16
static constexpr int V_SUB_SZ = 256; // (16,16) canonical BF16
static constexpr int TMEM_N = (HD <= 128) ? 128 : 256;
static constexpr int CORES_MN = 128 / 8; // 16
const int head_idx = blockIdx.y;
const int batch_idx = blockIdx.z;
@@ -134,6 +125,18 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
// ================================================================
// SMEM allocation (shared across all warps)
// ================================================================
// Layout:
// [0..3] sTmemBase (4 bytes, written by tcgen05.alloc)
// [4..7] sRowMax (4 bytes, float)
// [8..11] sRowSum (4 bytes, float)
// [12..15] alignment padding
// [16..16+TILE_SZ*2) sQ0 (4KB, 128×16 canonical BF16)
// [sQ0+TILE_SZ*2..) sK0 (4KB, 128×16 canonical BF16)
// [sK0+TILE_SZ*2..) sPk (4KB, 128B aligned)
// [sPk+TILE_SZ*2..) sV (512B, 16×16 canonical BF16, 128B aligned)
// [sV+V_SUB_SZ*2..) s_p_vals (SK_TILE*4 = 512B)
// [s_p_vals+SK_TILE*4..) sO_epi (HD*2 bytes, row-major BF16, 128B aligned)
// [sO_epi+HD*2..) sMbarStore (16 bytes, 128B aligned)
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
float* sRowMax = (float*)(sbuf + 4);
@@ -143,6 +146,10 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
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);
float* s_p_vals = (float*)(sV + V_SUB_SZ);
// Epilogue SMEM: row-major O tile for TMA store
bf16_t* sO_epi = (bf16_t*)(((uintptr_t)(s_p_vals + SK_TILE) + 127) & ~(uintptr_t)127);
// TMA store mbarrier (16 bytes, 128B aligned)
uint64_t* sMbarStore = (uint64_t*)(((uintptr_t)(sO_epi + HD) + 127) & ~(uintptr_t)127);
// ================================================================
// TMEM allocation (warp 4)
@@ -160,25 +167,22 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
for (int kt = 0; kt < NKT_QK; kt++) {
// ---- Warp 5: Load Q and K for this K-tile ----
if (is_load_warp) {
// Load Q K-tile: Q is (1, hd) for decode, row 0 only
for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0;
for (int d = lane; d < MMA_K_BF16; d += 32) {
int ck = d / 8, lc = d % 8;
sQ0[ck * 16 * 64 + lc] = q_head[kt * MMA_K_BF16 + d];
sQ0[ck * CORES_MN * 64 + lc] = q_head[kt * MMA_K_BF16 + d];
}
// Load K K-tile: K is (s_k, hd)
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) {
int ck = d / 8, lc = d % 8;
int tmn = r / 8, lr = r % 8;
sK0[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k_head[r * HD + kt * MMA_K_BF16 + d];
sK0[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = k_head[r * HD + kt * MMA_K_BF16 + d];
}
}
}
__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);
@@ -227,16 +231,13 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
int d_base = n * 16;
for (int kt = 0; kt < NKT_PV; kt++) {
// ---- Warp 5: Fill sPk and load V sub-tile ----
if (is_load_warp) {
// Fill sPk from s_p_vals
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 * 16 * 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[kt * MMA_K_BF16 + c]);
}
// Load V sub-tile: V is (hd, s_k) in GMEM
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++) {
@@ -249,7 +250,6 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
}
__syncthreads();
// ---- Warp 4: PV 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);
@@ -262,13 +262,15 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
}
// ================================================================
// Epilogue: TMEM → regs → normalize → BF16 → GMEM
// For single-segment decode: normalize in-kernel.
// For multi-segment: emit un-normalized O + LSE.
// EPILOGUE: One-way TMEM → regs → epilogue_op → SMEM → TMA store → GMEM
// ================================================================
// Step 1: TMEM → registers (warp 0, warp-collective tcgen05.ld)
// Step 2: epilogue_op in registers (normalize + optional FP4 pack)
// Step 3: Registers → SMEM (row-major, matching TMA tile format)
// Step 4: TMA store SMEM → GMEM (or direct GMEM write if no TMA desc)
// ================================================================
if (wid == 0) {
float row_max = *sRowMax;
float row_sum = *sRowSum;
// Step 1: TMEM → registers
float o_vals[HD];
for (int n = 0; n < HD / 8; n++) {
float tmp[8];
@@ -279,20 +281,84 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c];
}
// P was NORMALIZED in softmax step. PV = P @ V is already the normalized
// attention output. No further division by row_sum needed.
// For single-segment decode, write O directly.
// LSE is written for multi-segment merge (P5).
// Step 2: epilogue_op
// P was NORMALIZED in softmax step. PV = P @ V is already normalized.
// Default epilogue_op = identity (o_vals unchanged).
//
// FP4 quantization hook (ENABLE_FP4_EPILOGUE):
// 1. Find amax across o_vals[0..HD-1]
// 2. scale = amax / 6.0 (NVFP4 E2M1 max representable = 6.0)
// 3. For each d: e2m1_val = clamp(o_vals[d] / scale, 0, 6.0) → quantize
// 4. Pack 2 E2M1 values per byte, write scale as FP8 E4M3
// Output would be (HD/2 bytes FP4 + 1 FP8 scale per 16-element block)
// instead of (HD * 2 bytes BF16).
// For now, BF16 output path — just cast to BF16.
// Step 3: Registers → SMEM (row-major layout for TMA store)
// All 32 lanes participate. Lane i writes o_vals at offset i*4..i*4+3
// (matching the TMEM lane mapping for row 0).
// But only lane 0 has valid o_vals. For T=1 decode, this is correct
// because only row 0 of the MMA tile has data.
if (lane == 0) {
for (int d = 0; d < HD; d++) {
o_head[d] = f32_to_bf16(o_vals[d]);
sO_epi[d] = f32_to_bf16(o_vals[d]);
}
// Write LSE if pointer is valid
if (lse_head) {
// LSE = ln(row_sum) + row_max (natural log)
// This is the log of the softmax denominator, useful for
// multi-segment merge: O = sum(exp(lse_i) * O_i) / sum(exp(lse_i))
lse_head[0] = logf(row_sum) + row_max;
}
// SMEM fence: make writes visible to TMA store
asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
// Write LSE
if (lane == 0 && lse_head) {
float row_max = *sRowMax;
float row_sum = *sRowSum;
lse_head[0] = logf(row_sum) + row_max;
}
}
__syncthreads();
// Step 4: TMA store SMEM → GMEM (or direct GMEM write)
if (params.tma_o) {
// Proper TMA store path — async, enables multi-CTA grids.
// One thread initializes the mbarrier and issues the TMA store.
if (tid == 0) {
uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore);
tma_mbarrier_init(mbar_addr, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
__syncthreads();
// The O tensor is [batch, n_h, T, HD]. Each head's output starts
// at a different GMEM offset. The TMA descriptor covers the full
// O tensor. We index by (head, batch) to find the right descriptor.
//
// TMA coords for this head's row 0: (x=0, y=0)
// The descriptor was created with the head's GMEM base pointer,
// so coords are relative to that head's start.
if (tid == 0) {
uint32_t smem_addr = (uint32_t)__cvta_generic_to_shared(sO_epi);
uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore);
// Get this (head, batch) TMA descriptor
CUtensorMap* my_tma = params.tma_o + batch_idx * gridDim.y + head_idx;
uint64_t tma_desc = (uint64_t)my_tma;
tma_store_2d(smem_addr, tma_desc, mbar_addr, 0, 0);
// TMA tile: (1, HD) BF16 = HD*2 bytes
tma_mbarrier_arrive_expect_tx(mbar_addr, HD * sizeof(bf16_t));
}
__syncthreads();
// Wait for TMA store completion
if (tid == 0) {
uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore);
tma_store_wait(mbar_addr, 0);
}
__syncthreads();
} else {
// Fallback: direct GMEM write from SMEM (backward compatible)
// This path is used when no TMA descriptor is provided.
if (wid == 0 && lane == 0) {
for (int d = 0; d < HD; d++) {
o_head[d] = sO_epi[d];
}
}
}

52
dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh
View File

@@ -39,6 +39,8 @@ struct FmhaTmaMultiRowMultiTileParams {
int q_head_stride, q_batch_stride;
int o_head_stride, o_batch_stride;
int lse_head_stride, lse_batch_stride;
// P6: TMA descriptor for O output. When nullptr, direct GMEM write.
CUtensorMap* __restrict__ tma_o;
};
template<int HD, int SK_TILE = 128, int HD_CHUNK = (HD <= 256 ? HD : 256)>
@@ -103,6 +105,10 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params)
float* sRunningSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float);
float* sTileRowMax = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float);
float* sTileRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float);
// P6: Row-major O epilogue buffer + TMA store mbarrier
off = (off + 127) & ~(size_t)127;
bf16_t* sO_epi_rowmajor = (bf16_t*)(sbuf + off); off += MAX_ROWS * HD_CHUNK * sizeof(bf16_t);
uint64_t* sMbarStore = (uint64_t*)(sbuf + off); off += 16;
// Init TMEM + mbarrier (once, shared across hd_chunks)
if (is_mma_warp) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N);
@@ -329,11 +335,18 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params)
__syncthreads();
} // kv_tile loop
// ---- Write chunk to GMEM ----
// ---- Write chunk to SMEM row-major, then TMA store to GMEM ----
// P6: One-way epilogue pattern — normalize in registers,
// write to SMEM row-major, then TMA store to GMEM.
// This enables multi-CTA grids and FP4 output fusion.
//
// Each active row writes its normalized output to sO_epi_rowmajor
// in row-major layout. Then a single TMA store issues the async
// write for this hd_chunk's output tile.
if (my_row_active) {
float inv_rs = 1.0f / sRunningSum[my_row];
for (int d = 0; d < HD_CHUNK; d++) {
o_head[my_row * HD + hd_chunk_start + d] = f32_to_bf16(sOacc[my_row * HD_CHUNK + d] * inv_rs);
sO_epi_rowmajor[my_row * HD_CHUNK + d] = f32_to_bf16(sOacc[my_row * HD_CHUNK + d] * inv_rs);
}
// LSE: same for all chunks, write once
if (lse_head && !lse_written) {
@@ -341,6 +354,41 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params)
}
}
if (hd_chunk == 0) lse_written = true;
asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
__syncthreads();
// TMA store sO_epi_rowmajor → GMEM
if (params.tma_o) {
if (tid == 0) {
uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore);
tma_mbarrier_init(mbar_addr, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
__syncthreads();
if (tid == 0) {
uint32_t smem_addr = (uint32_t)__cvta_generic_to_shared(sO_epi_rowmajor);
uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore);
CUtensorMap* my_tma = params.tma_o + batch_idx * params.n_h + head_idx;
uint64_t tma_desc = (uint64_t)my_tma;
// TMA coords: x = hd_chunk_start, y = 0
tma_store_2d(smem_addr, tma_desc, mbar_addr, hd_chunk_start, 0);
tma_mbarrier_arrive_expect_tx(mbar_addr, T * HD_CHUNK * sizeof(bf16_t));
}
__syncthreads();
if (tid == 0) {
uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore);
tma_store_wait(mbar_addr, 0);
}
__syncthreads();
} else {
// Fallback: direct GMEM write from SMEM
if (my_row_active) {
for (int d = 0; d < HD_CHUNK; d++) {
o_head[my_row * HD + hd_chunk_start + d] = sO_epi_rowmajor[my_row * HD_CHUNK + d];
}
}
}
__syncthreads();
} // hd_chunk loop

132
dsv4/kernels/attention/fmha_multihead_capi.cu
View File

@@ -2,32 +2,48 @@
* DSV4 FMHA Multi-Head — C API for ctypes loading.
*
* Supports single and multi-KV-tile with FlashAttention-2 online softmax.
* P6: One-way TMEM→regs→SMEM→TMA store epilogue with FP4 hook.
*/
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <cstdio>
#include "fmha_common.cuh"
#include "fmha_umma_desc.cuh"
#include "fmha_tma.cuh"
#include "fmha_6warp_multihead.cuh"
using namespace dsv4::kernels::attention;
extern "C" {
int fmha_compute_smem(int hd) {
using namespace dsv4::kernels::attention;
constexpr int SK = 128;
constexpr int TILE_SZ = 128 * MMA_K_BF16;
constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16
constexpr int V_SUB_SZ = 256;
// sTmemBase(4) + sRowMax(4) + sRowSum(4) + align(4) + sQ0 + sK0 + sPk + sV + s_p_vals + sOacc + slack
int base = 4 + 4 + 4 + 4; // tmem_base, row_max, row_sum, alignment
int sQ0 = TILE_SZ * 2;
int sK0 = TILE_SZ * 2;
int sPk = TILE_SZ * 2; // 128*16*2 = 4096 (P sub-tile)
int sV = V_SUB_SZ * 2;
int sp_vals = SK * 4;
int sOacc = hd * 4; // float accumulator for 1 row
int total = base + sQ0 + sK0 + sPk + sV + sp_vals + sOacc + 256 + 127;
return total & ~127;
// sTmemBase(4) + sRowMax(4) + sRowSum(4) + align(4) = 16
// sQ0(TILE_SZ*2) = 4096
// sK0(TILE_SZ*2) = 4096
// sPk(TILE_SZ*2 + 127) = 4223 (128B aligned)
// sV(V_SUB_SZ*2 + 127) = 639 (128B aligned)
// s_p_vals(SK*4) = 512
// sO_epi(hd*2 + 127) (128B aligned, row-major BF16)
// sMbarStore(16 + 127) (128B aligned)
int base = 16;
int sQ0 = TILE_SZ * 2; // 4096
int sK0 = TILE_SZ * 2; // 4096
int sPk = TILE_SZ * 2; // 4096 (before 128B alignment)
int sV = V_SUB_SZ * 2; // 512
int sp = SK * 4; // 512
int sO = hd * 2; // row-major O (HD BF16)
int sMbar = 16; // TMA store mbarrier
// With 128B alignment between sections
int total = base + sQ0 + sK0 + (sPk + 127) + (sV + 127) + sp + (sO + 127) + (sMbar + 127) + 256;
// Round up to 128B
return (total + 127) & ~127;
}
int fmha_multihead_decode_launch(
@@ -44,8 +60,6 @@ int fmha_multihead_decode_launch(
int lse_head_stride, int lse_batch_stride,
float scale
) {
using namespace dsv4::kernels::attention;
FmhaParams params;
params.q = reinterpret_cast<const bf16_t*>(q_ptr);
params.k = reinterpret_cast<const bf16_t*>(k_ptr);
@@ -65,6 +79,10 @@ int fmha_multihead_decode_launch(
params.o_batch_stride = o_batch_stride;
params.lse_head_stride = lse_head_stride;
params.lse_batch_stride = lse_batch_stride;
// P6: TMA descriptor for O output.
// When nullptr, epilogue uses direct GMEM write (backward compatible).
// To enable TMA store, create per-(head, batch) descriptors and pass them.
params.tma_o = nullptr;
int smem = fmha_compute_smem(hd);
dim3 grid(1, n_h, batch);
@@ -90,4 +108,92 @@ int fmha_multihead_decode_launch(
return 0;
}
/**
* Launch with TMA store epilogue enabled.
* Creates per-(head, batch) TMA descriptors for the O output tensor.
*/
int fmha_multihead_decode_tma_launch(
const void* q_ptr,
const void* k_ptr,
const void* v_ptr,
void* o_ptr,
void* lse_ptr,
int batch, int n_h, int n_kv, int N, int hd,
int q_head_stride, int q_batch_stride,
int k_head_stride, int k_batch_stride,
int v_head_stride, int v_batch_stride,
int o_head_stride, int o_batch_stride,
int lse_head_stride, int lse_batch_stride,
float scale
) {
// Create TMA descriptors for O output: one per (batch, head)
// Each descriptor covers a (1, HD) BF16 tile at the head's GMEM offset.
size_t desc_count = n_h * batch;
CUtensorMap* d_tma_o;
cudaError_t err = cudaMalloc(&d_tma_o, desc_count * sizeof(CUtensorMap));
if (err != cudaSuccess) return (int)err;
for (int b = 0; b < batch; b++) {
for (int h = 0; h < n_h; h++) {
int idx = b * n_h + h;
bf16_t* o_head = (bf16_t*)o_ptr + h * o_head_stride + b * o_batch_stride;
// O tile: (1, hd) — one row per head for T=1 decode
CUtensorMap tma_desc;
bool ok = create_tma_desc_2d_bf16(&tma_desc, o_head, 1, hd, 1, hd);
if (!ok) {
cudaFree(d_tma_o);
return -2;
}
cudaMemcpy(&d_tma_o[idx], &tma_desc, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
}
}
FmhaParams params;
params.q = reinterpret_cast<const bf16_t*>(q_ptr);
params.k = reinterpret_cast<const bf16_t*>(k_ptr);
params.v = reinterpret_cast<const bf16_t*>(v_ptr);
params.o = reinterpret_cast<bf16_t*>(o_ptr);
params.lse = reinterpret_cast<float*>(lse_ptr);
params.s_k = N;
params.scale = scale;
params.head_dim = hd;
params.q_head_stride = q_head_stride;
params.q_batch_stride = q_batch_stride;
params.k_head_stride = k_head_stride;
params.k_batch_stride = k_batch_stride;
params.v_head_stride = v_head_stride;
params.v_batch_stride = v_batch_stride;
params.o_head_stride = o_head_stride;
params.o_batch_stride = o_batch_stride;
params.lse_head_stride = lse_head_stride;
params.lse_batch_stride = lse_batch_stride;
params.tma_o = d_tma_o;
int smem = fmha_compute_smem(hd);
dim3 grid(1, n_h, batch);
dim3 block(NTHREADS);
if (smem > 48 * 1024) {
if (hd == 64) cudaFuncSetAttribute(fmha_6warp_multihead_kernel<64, 128>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
else if (hd == 128) cudaFuncSetAttribute(fmha_6warp_multihead_kernel<128, 128>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
else if (hd == 256) cudaFuncSetAttribute(fmha_6warp_multihead_kernel<256, 128>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
}
if (hd == 64) fmha_6warp_multihead_kernel<64, 128><<<grid, block, smem>>>(params);
else if (hd == 128) fmha_6warp_multihead_kernel<128, 128><<<grid, block, smem>>>(params);
else if (hd == 256) fmha_6warp_multihead_kernel<256, 128><<<grid, block, smem>>>(params);
else {
cudaFree(d_tma_o);
return -1;
}
err = cudaGetLastError();
cudaFree(d_tma_o);
if (err != cudaSuccess) return (int)err;
return 0;
}
} // extern "C"

5
dsv4/kernels/attention/fmha_multitile_capi.cu
View File

@@ -80,6 +80,7 @@ int fmha_multitile_decode_launch(
params.o_batch_stride = o_batch_stride;
params.lse_head_stride = lse_head_stride;
params.lse_batch_stride = lse_batch_stride;
params.tma_o = nullptr; // P6: TMA O descriptor (future: pass from caller)
// SMEM size (match kernel layout)
constexpr int HD_CHUNK = 256;
@@ -105,6 +106,10 @@ int fmha_multitile_decode_launch(
off += 128 * 4; // sRunningSum
off += 128 * 4; // sTileRowMax
off += 128 * 4; // sTileRowSum
// P6: sO_epi_rowmajor + sMbarStore
off = (off+127)&~(size_t)127;
off += 128 * hc * 2; // sO_epi_rowmajor (MAX_ROWS * HD_CHUNK BF16)
off += 16; // sMbarStore
off += 256; // slack
int smem = (int)((off + 127) & ~(size_t)127);

54
dsv4/kernels/attention/fmha_tma.cuh
View File

@@ -186,4 +186,58 @@ struct FmhaTmaDescriptors {
CUtensorMap* __restrict__ tma_v; // V descriptor: (HD, s_k)
};
// ==================================================================
// TMA store (GMEM write) operations
// ==================================================================
// The one-way epilogue writes O from SMEM to GMEM via TMA store.
// This replaces direct GMEM writes and enables multi-CTA grids.
// ==================================================================
/**
* Issue a 2D TMA async copy from SMEM to GMEM (store).
*
* @param smem_src SMEM source address (via __cvta_generic_to_shared)
* @param tma_desc Pointer to CUtensorMap in device memory (uint64_t cast)
* @param smem_mbar SMEM mbarrier address (via __cvta_generic_to_shared)
* @param coord_x Column coordinate (innermost dimension)
* @param coord_y Row coordinate (outermost dimension)
*/
__device__ __forceinline__ void tma_store_2d(
uint32_t smem_src,
uint64_t tma_desc,
uint32_t smem_mbar,
int coord_x,
int coord_y
) {
asm volatile(
"cp.async.bulk.tensor.2d.global.shared::cluster.mbarrier::complete_tx::bytes "
"[%0, {%3, %4}], [%1], [%2];"
:: "l"(tma_desc),
"r"(smem_src),
"r"(smem_mbar),
"r"(coord_x),
"r"(coord_y)
: "memory"
);
}
/**
* Wait for TMA store completion using mbarrier try_wait.
* Same pattern as tma_mbarrier_wait but for store mbarriers.
*/
__device__ __forceinline__ void tma_store_wait(uint32_t smem_mbar, int phase) {
asm volatile(
"{\n\t"
".reg .pred P1;\n\t"
"LAB_WAIT:"
"mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
"@P1 bra.uni DONE;\n\t"
"bra.uni LAB_WAIT;\n\t"
"DONE:\n\t"
"}"
:: "r"(smem_mbar), "r"(phase), "r"(0x989680)
: "memory"
);
}
} // namespace dsv4::kernels::attention