diff --git a/dsv4/kernels/attention/fmha_6warp_multihead.cuh b/dsv4/kernels/attention/fmha_6warp_multihead.cuh index c0433bc1..0b9beed7 100644 --- a/dsv4/kernels/attention/fmha_6warp_multihead.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multihead.cuh @@ -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 @@ -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(); diff --git a/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh b/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh index a20c414c..9c62dec3 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh @@ -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 @@ -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(); diff --git a/dsv4/kernels/attention/fmha_multihead_capi.cu b/dsv4/kernels/attention/fmha_multihead_capi.cu index 47bcb442..4cb17385 100644 --- a/dsv4/kernels/attention/fmha_multihead_capi.cu +++ b/dsv4/kernels/attention/fmha_multihead_capi.cu @@ -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(q_ptr); - params.k = reinterpret_cast(k_ptr); - params.v = reinterpret_cast(v_ptr); - params.o = reinterpret_cast(o_ptr); - params.lse = reinterpret_cast(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><<>>(params); - else if (hd == 128) fmha_6warp_multihead_kernel<128, 128><<>>(params); - else if (hd == 256) fmha_6warp_multihead_kernel<256, 128><<>>(params); - else { - cudaFree(d_tma_o); - return -1; - } - - err = cudaGetLastError(); - cudaFree(d_tma_o); - if (err != cudaSuccess) return (int)err; - return 0; -} - } // extern "C" diff --git a/dsv4/kernels/attention/fmha_multitile_capi.cu b/dsv4/kernels/attention/fmha_multitile_capi.cu index 676f7655..cc3fdeed 100644 --- a/dsv4/kernels/attention/fmha_multitile_capi.cu +++ b/dsv4/kernels/attention/fmha_multitile_capi.cu @@ -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);