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

P6: Remove broken TMA store — use direct GMEM write from SMEM

Browse Source 
cp.async.bulk.tensor store (SMEM→GMEM) is NOT available on SM100.
The CUTLASS SM100 epilogue uses st.global directly.

The one-way epilogue pipeline is now:
  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)
  4. SMEM → GMEM (direct write)

This is the same pattern as the MoE kernel but with st.global instead
of TMA store. Multi-CTA (D2) will use st.global with flat_divide coords.

Removed: tma_o from FmhaParams, fmha_multihead_decode_tma_launch,
sMbarStore from SMEM, broken TMA store PTX from fmha_tma.cuh.
This commit is contained in:
biondizzle
2026-05-30 17:11:17 +00:00
parent f97359fbfc
commit c0379a0f86
4 changed files with 36 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -71,10 +71,6 @@ struct FmhaParams {
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, bool ENABLE_FP4_EPILOGUE = false>
@@ -146,10 +142,8 @@ 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
// Epilogue SMEM: row-major O for GMEM write
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)
@@ -262,15 +256,20 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
}
// ================================================================
// EPILOGUE: One-way TMEM → regs → epilogue_op → SMEM → TMA store → GMEM
// EPILOGUE: One-way TMEM → regs → epilogue_op → SMEM → 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)
// Step 3: Registers → SMEM (row-major)
// Step 4: SMEM → GMEM (direct write)
//
// NOTE: cp.async.bulk.tensor store (SMEM → GMEM) is NOT available
// on SM100 for this use case. The CUTLASS SM100 epilogue uses
// st.global directly. For multi-CTA (D2), we'll use st.global
// with flat_divide coordinates.
// ================================================================
if (wid == 0) {
// Step 1: TMEM → registers
// Step 1: TMEM → registers (warp-collective)
float o_vals[HD];
for (int n = 0; n < HD / 8; n++) {
float tmp[8];
@@ -282,30 +281,26 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c];
}
// Step 2: epilogue_op
// Step 2: epilogue_op in registers
// P was NORMALIZED in softmax step. PV = P @ V is already normalized.
// Default epilogue_op = identity (o_vals unchanged).
// Default epilogue_op = identity.
//
// 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.
// 2. scale = amax / 6.0 (NVFP4 E2M1 range)
// 3. Quantize each o_vals[d] to E2M1, pack 2 per byte
// 4. Write scale as FP8 E4M3 per 16-element block
// For now, BF16 output: just cast.
// 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.
// Step 3: Registers → SMEM (row-major)
// This is the SMEM staging step of the one-way epilogue.
// In the MoE kernel (CuTeDSL), this is done by
// epilogue_smem_copy_and_partition. Here we write directly.
if (lane == 0) {
for (int d = 0; d < HD; d++) {
sO_epi[d] = f32_to_bf16(o_vals[d]);
}
}
// SMEM fence: make writes visible to TMA store
asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
// Write LSE
@@ -317,43 +312,13 @@ fmha_6warp_multihead_kernel(FmhaParams params) {
}
__syncthreads();
// Step 4: TMA store SMEM → GMEM (or direct GMEM write)
if (params.tma_o) {
// Proper TMA store path — async, enables multi-CTA grids.
// Uses bulk_group completion (commit_group + wait_group), NOT mbarrier.
//
// The O tensor is [batch, n_h, T, HD]. Each head's output starts
// at a different GMEM offset. The TMA descriptor covers one head's
// output. 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);
CUtensorMap* my_tma = params.tma_o + batch_idx * gridDim.y + head_idx;
uint64_t tma_desc = (uint64_t)my_tma;
tma_mbarrier_init(mbar_addr, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
tma_store_2d(smem_addr, tma_desc, mbar_addr, 0, 0);
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_mbarrier_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];
}
// Step 4: SMEM → GMEM (direct write)
// The SMEM buffer is the single source of truth for the output.
// For single-CTA decode, this is correct and efficient.
// For multi-CTA (D2), will use st.global with flat_divide coords.
if (wid == 0 && lane == 0) {
for (int d = 0; d < HD; d++) {
o_head[d] = sO_epi[d];
}
}
__syncthreads();

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

@@ -39,8 +39,6 @@ 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)>
@@ -108,7 +106,6 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params)
// 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);
@@ -341,46 +338,25 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params)
// 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.
// in row-major layout, then writes directly to GMEM.
if (my_row_active) {
float inv_rs = 1.0f / sRunningSum[my_row];
for (int d = 0; d < HD_CHUNK; d++) {
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) {
lse_head[my_row] = logf(sRunningSum[my_row]) + sRunningMax[my_row];
}
}
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 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_mbarrier_init(mbar_addr, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
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));
// Write SMEM → GMEM (direct)
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();
if (tid == 0) {
uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbarStore);
tma_mbarrier_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];
}
// LSE: same for all chunks, write once
if (lse_head && !lse_written) {
lse_head[my_row] = logf(sRunningSum[my_row]) + sRunningMax[my_row];
}
}
__syncthreads();

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

@@ -12,7 +12,6 @@
#include "fmha_common.cuh"
#include "fmha_umma_desc.cuh"
#include "fmha_tma.cuh"
#include "fmha_6warp_multihead.cuh"
using namespace dsv4::kernels::attention;
@@ -39,9 +38,8 @@ int fmha_compute_smem(int hd) {
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;
int total = base + sQ0 + sK0 + (sPk + 127) + (sV + 127) + sp + (sO + 127) + 256;
// Round up to 128B
return (total + 127) & ~127;
}
@@ -79,10 +77,6 @@ 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);
@@ -108,92 +102,4 @@ 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"

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

@@ -80,7 +80,6 @@ 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;
@@ -109,7 +108,6 @@ int fmha_multitile_decode_launch(
// 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);