diff --git a/dsv4/kernels/attention/fmha_common.cuh b/dsv4/kernels/attention/fmha_common.cuh new file mode 100644 index 00000000..c9089a6d --- /dev/null +++ b/dsv4/kernels/attention/fmha_common.cuh @@ -0,0 +1,61 @@ +/** + * DSV4 FMHA shared definitions — base header. + * BF16 type, TMEM ops, warp reductions, constants. + */ + +#pragma once + +#include +#include +#include + +namespace dsv4::kernels::attention { + +typedef unsigned short bf16_t; + +__device__ __forceinline__ bf16_t f32_to_bf16(float f) { + bf16_t h; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); return h; +} +__device__ __forceinline__ float bf16_to_f32(bf16_t h) { + float f; asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); return f; +} + +constexpr int WARP = 32; +constexpr int NTHREADS = 192; +constexpr int NWARPS = 6; + +__device__ __forceinline__ float wmax(float v) { + for(int o=16;o>0;o>>=1) v=fmaxf(v,__shfl_xor_sync(0xFFFFFFFF,v,o)); return v; +} +__device__ __forceinline__ float wsum(float v) { + for(int o=16;o>0;o>>=1) v+=__shfl_xor_sync(0xFFFFFFFF,v,o); return v; +} + +// TMEM operations +__device__ uint32_t tmem_alloc(int n) { + uint32_t b = 0; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 %0, [%1], %2;" + : "=r"(b) : "r"(0), "r"(n)); + return b; +} +__device__ void tmem_dealloc(uint32_t b, int n) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(b), "r"(n)); +} +__device__ void tmem_load_col(uint32_t col, int row_group, + float& r0, float& r1, float& r2, float& r3) { + uint32_t addr = col + row_group; + asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" + : "=f"(r0), "=f"(r1), "=f"(r2), "=f"(r3) : "r"(addr)); +} +__device__ void tmem_store_col(uint32_t col, int row_group, + float r0, float r1, float r2, float r3) { + uint32_t addr = col + row_group; + asm volatile("tcgen05.st.sync.aligned.16x256b.x1.b32 [%0], {%1, %2, %3, %4};" + :: "r"(addr), "f"(r0), "f"(r1), "f"(r2), "f"(r3)); +} +__device__ void tmem_fence() { + asm volatile("tcgen05.fence.cta_group::1.sync.aligned;" ::: "memory"); +} + +} // namespace dsv4::kernels::attention diff --git a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh index f7d86e94..851d8d49 100644 --- a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh +++ b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh @@ -1,382 +1,133 @@ /** - * DSV4 FMHA One-Way Correction Epilogue — Raw CUDA C++ for Blackwell SM100 + * DSV4 FMHA Phase 2 — TMEM accumulator + one-way correction epilogue. * - * This is the Priority 2 implementation from ROADMAP.md: - * Replace epilogue_tma_store with the MoE-style one-way pipeline: + * Priority 2 from ROADMAP: + * TMEM → regs (tcgen05.ld) → normalize → BF16 → GMEM * - * TMEM → registers (tcgen05.ld) → normalize/cast → SMEM → GMEM (TMA/bulk) - * - * This unblocks: - * - D2 multi-CTA grid (flat_divide + cpasync.tma_partition works with this pattern) - * - NVFP4-1.2 (register slot for FP4 amax + pack between t2r and r2s) - * - In-kernel normalize (O / row_sum in registers) - * - * The MoE kernel (fused_swiglu.py) uses this exact pattern successfully: - * epilogue_tmem_copy_and_partition → SwiGLU/clamp → epilogue_smem_copy_and_partition - * - * We do the same but with normalize instead of SwiGLU. + * D1.5 fix: O rescale in REGISTERS between KV tiles. + * Unblocks: D2 multi-CTA, NVFP4-1.2 (register slot for FP4 pack). */ - #pragma once - -#include -#include -#include +#include "fmha_common.cuh" namespace dsv4::kernels::attention { -typedef unsigned short bf16_t; - -__device__ __forceinline__ bf16_t f32_to_bf16(float f) { - bf16_t h; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); return h; -} -__device__ __forceinline__ float bf16_to_f32(bf16_t h) { - float f; asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); return f; -} - -// ===================================================================== -// TMEM operations -// ===================================================================== - -__device__ uint32_t tmem_alloc(int n) { - uint32_t b = 0; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 %0, [%1], %2;" - : "=r"(b) : "r"(0), "r"(n)); - return b; -} - -__device__ void tmem_dealloc(uint32_t b, int n) { - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" - :: "r"(b), "r"(n)); -} - -/** - * TMEM load: 16 rows × 256 bits from one column. - * Returns 4 FP32 values per calling thread. - * 16 threads (half-warp) cooperate per column load. - * For 128-row tile: 8 column-loads per column, each covering 16 rows. - */ -__device__ void tmem_load_col(uint32_t col, int row_group, - float& r0, float& r1, float& r2, float& r3) { - // Each column in TMEM is addressed as col_index + row_offset - // tcgen05.ld reads 16 rows × 256 bits from one column - // 256 bits = 8 FP32 values, delivered as 4 per thread (2 threads per 16-row group) - // The instruction signature: tcgen05.ld.sync.aligned.16x256b.x1.b32 {r0,r1,r2,r3}, [col] - uint32_t addr = col + row_group; - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" - : "=f"(r0), "=f"(r1), "=f"(r2), "=f"(r3) - : "r"(addr) - ); -} - -/** - * TMEM store: 16 rows × 256 bits to one column. - */ -__device__ void tmem_store_col(uint32_t col, int row_group, - float r0, float r1, float r2, float r3) { - uint32_t addr = col + row_group; - asm volatile( - "tcgen05.st.sync.aligned.16x256b.x1.b32 [%0], {%1, %2, %3, %4};" - :: "r"(addr), "f"(r0), "f"(r1), "f"(r2), "f"(r3) - ); -} - -__device__ void tmem_fence() { - asm volatile("tcgen05.fence.cta_group::1.sync.aligned;" ::: "memory"); -} - -// ===================================================================== -// One-way Correction Epilogue -// ===================================================================== - -/** - * FMHA one-way correction epilogue: - * Read O from TMEM → normalize (O/row_sum) → cast to BF16 → write to GMEM. - * - * This is the exact analog of the MoE epilogue pattern: - * TMEM → regs (tcgen05.ld) → [normalize + BF16 cast] → SMEM → GMEM - * - * Template params: - * HD: head dimension - * TILE_M: number of query rows (128 for decode head-packed, 1 for single-head) - * NORMALIZE: if true, divide by row_sum before writing - * - * Thread mapping: - * - All 192 threads (6 warps) participate in the epilogue - * - Each thread handles a subset of the (row, d) output elements - * - TMEM layout: 128 rows × (HD/2) FP32 columns - * Each column holds 4 FP32 values (16 rows × 256 bits per tcgen05.ld) - * For HD=64: 32 columns, 8 row-groups per column - * Total: 128 rows × 64 values = 8192 FP32 values - */ -template -__device__ void fmha_epilogue( - uint32_t tmem_o_base, // TMEM base column for O - float* row_sums, // (TILE_ROWS,) row sums for normalization - bf16_t* __restrict__ gmem_o, // (TILE_ROWS, HD) output in GMEM - int gmem_stride, // stride between rows in GMEM (in bf16_t elements) - int smem_size_bytes // SMEM buffer size for intermediate BF16 -) { - // SMEM buffer for BF16 output (written by t2r+normalize, read by s2g) - extern __shared__ char smem_epilogue[]; - bf16_t* smem_o = reinterpret_cast(smem_epilogue); - - // TMEM layout for O accumulator: - // 128 rows, HD/2 FP32 columns - // Each tcgen05.ld reads 16 rows × 4 FP32 from one column - // So: 128/16 = 8 row-groups per column, HD/2 columns - // - // Thread mapping: each thread loads one (row_group, col) tile - // 8 row-groups × (HD/2) columns = 4*HD total load operations - // With 192 threads: each thread handles ~4*HD/192 loads - // For HD=64: 256 loads, ~1.3 per thread → 1 per thread (some threads idle) - // For HD=128: 512 loads, ~2.7 per thread → 2-3 per thread - - const int n_cols = HD / 2; // TMEM columns for O (2 BF16 per FP32 column in pack::16b mode) - const int n_row_groups = TILE_ROWS / 16; // 8 row-groups (16 rows per tcgen05.ld) - - // Each thread processes a range of (row_group, col) pairs - const int total_tiles = n_row_groups * n_cols; - const int tid = threadIdx.x; - - // Phase 1: Load from TMEM → normalize → cast to BF16 → write to SMEM - for (int tile = tid; tile < total_tiles; tile += NTHREADS) { - int rg = tile / n_cols; // row group (0-7) - int col = tile % n_cols; // TMEM column (0 to HD/2-1) - - // Load 4 FP32 values from TMEM - float r0, r1, r2, r3; - tmem_load_col(tmem_o_base + col, rg, r0, r1, r2, r3); - - // Normalize by row_sum - // Each row group covers 16 rows. We need the row_sum for each row. - // For decode (T=1), only row 0 matters. - // Row mapping: row_group r covers rows [r*16, r*16+15] - // For T=1, only row 0 has a valid row_sum. - // TODO: For T>1, load per-row row_sums from SMEM. - float inv_sum = 1.0f; - if (NORMALIZE && row_sums[0] > 0.0f) { - inv_sum = 1.0f / row_sums[0]; - } - - // Normalize + cast to BF16 - // Each FP32 value maps to one output element - // TMEM column col, row group rg, values r0-r3 → output positions - // Row: rg*16 + (value index within the 16 rows) - // Col in output: col*2 + (0 or 1) — since 2 BF16 per FP32 in pack mode - // - // Actually, for the un-normalized O output (which is what the CuTeDSL - // kernel produces), the TMEM layout packs 2 BF16 per FP32 column. - // But in our reference kernel, O is in FP32 SMEM, not TMEM. - // For the TMEM-based kernel, we'll need to understand the exact layout. - // - // For now, write the normalized values directly to GMEM (skip SMEM staging). - // This is the "correct but not optimal" path — SMEM staging would allow - // TMA bulk copy which is faster for large outputs. - - if (rg == 0) { // Only row 0 for decode T=1 - // Write to GMEM directly - int d0 = col * 4 + 0; - int d1 = col * 4 + 1; - int d2 = col * 4 + 2; - int d3 = col * 4 + 3; - if (d0 < HD) gmem_o[d0] = f32_to_bf16(r0 * inv_sum); - if (d1 < HD) gmem_o[d1] = f32_to_bf16(r1 * inv_sum); - if (d2 < HD) gmem_o[d2] = f32_to_bf16(r2 * inv_sum); - if (d3 < HD) gmem_o[d3] = f32_to_bf16(r3 * inv_sum); - } - } -} - -// ===================================================================== -// FMHA Decode Kernel with TMEM + Correction Epilogue -// ===================================================================== - -constexpr int WARP = 32; -constexpr int NTHREADS = 192; -constexpr int NWARPS = 6; - -__device__ __forceinline__ float wmax(float v) { - for(int o=16;o>0;o>>=1) v=fmaxf(v,__shfl_xor_sync(0xFFFFFFFF,v,o)); return v; -} -__device__ __forceinline__ float wsum(float v) { - for(int o=16;o>0;o>>=1) v+=__shfl_xor_sync(0xFFFFFFFF,v,o); return v; -} - -/** - * FMHA decode with TMEM accumulator and one-way correction epilogue. - * - * Phase 2: Uses TMEM for O accumulation, correction epilogue for normalize. - * QK and PV still computed in registers (scalar) — tcgen05.mma comes in Phase 3. - * - * The key innovation: O rescale happens in REGISTERS between KV tiles, - * loading from TMEM → registers → multiply → store back to TMEM. - * This is the D1.5 fix that CuTeDSL couldn't do (TMEM round-trip broken). - */ template __global__ void __launch_bounds__(NTHREADS) fmha_decode_tmem( - const bf16_t* __restrict__ q, - const bf16_t* __restrict__ k, - const bf16_t* __restrict__ v, - bf16_t* __restrict__ o, + const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, + const bf16_t* __restrict__ v, bf16_t* __restrict__ o, int bstride_q, int bstride_kv, int bstride_o, - int s_k, int n_comp, int swa_len, - float scale, - const float* __restrict__ attn_sink, - float* __restrict__ lse_out + int s_k, int n_comp, int swa_len, float scale, + const float* __restrict__ attn_sink, float* __restrict__ lse_out ) { - const int head = blockIdx.y; - const int batch = blockIdx.z; - const int tid = threadIdx.x; - const int wid = tid / WARP; - const int lane = tid % WARP; + const int head = blockIdx.y, batch = blockIdx.z, tid = threadIdx.x; + const int wid = tid / WARP, lane = tid % WARP; - const bf16_t* qh = q + batch * bstride_q + head * HD; - const bf16_t* kb = k + batch * bstride_kv; - const bf16_t* vb = v + batch * bstride_kv; - bf16_t* oh = o + batch * bstride_o + head * HD; + const bf16_t* qh = q + batch*bstride_q + head*HD; + const bf16_t* kb = k + batch*bstride_kv; + const bf16_t* vb = v + batch*bstride_kv; + bf16_t* oh = o + batch*bstride_o + head*HD; // TMEM allocation for O accumulator - // O needs HD FP32 values (for T=1 decode) - // TMEM columns: each holds 128 FP32 values (128 rows × 1 FP32 per row per column) - // For HD=64: 64 columns needed, but TMEM loads 4 FP32 per column per row-group - // So we need ceil(HD/4) = 16 columns for HD=64 - const int tmem_o_cols = (HD + 3) / 4; // 4 FP32 per tcgen05.ld per column - int tmem_n = 1; while(tmem_n < tmem_o_cols + 4) tmem_n *= 2; // round to power of 2 + // Each tcgen05.ld reads 4 FP32 per column per row-group + // For HD=64: need ceil(64/4)=16 columns + const int tmem_o_cols = (HD + 3) / 4; + int tmem_n = 1; while(tmem_n < tmem_o_cols + 4) tmem_n *= 2; uint32_t tb = 0; - if (wid == 0 && lane == 0) tb = tmem_alloc(tmem_n); + if (wid==0 && lane==0) tb = tmem_alloc(tmem_n); tb = __shfl_sync(0xFFFFFFFF, tb, 0); - const uint32_t to = tb; // O starts at TMEM base + const uint32_t to = tb; - // SMEM for Q, row_sums + // SMEM for Q and row_sums extern __shared__ char sbuf[]; - float* sQ = (float*)sbuf; // HD floats - float* sRowSums = (float*)(sbuf + HD * sizeof(float)); // 1 float (row_sum for T=1) + float* sQ = (float*)sbuf; + float* sRowSums = (float*)(sbuf + HD*sizeof(float)); - for (int d = tid; d < HD; d += NTHREADS) sQ[d] = bf16_to_f32(qh[d]); + for (int d=tid; d0 && c>=n_comp+swa_len) s_val = -INFINITY; - if (swa_len > 0 && c >= n_comp + swa_len) s_val = -INFINITY; - - // Online softmax with O rescale in TMEM float new_max = fmaxf(row_max, s_val); if (new_max > row_max) { float rescale = expf(row_max - new_max); - // D1.5 FIX: Rescale O in TMEM - // Load O from TMEM → multiply by rescale → store back - // This is the one-way path: TMEM → regs → multiply → TMEM - // (NOT a round-trip with mismatched atoms — we use the SAME - // tcgen05.ld + tcgen05.st pair, which IS correct for same-column ops) - for (int col = 0; col < tmem_o_cols; col++) { - float r0, r1, r2, r3; - tmem_load_col(to + col, 0, r0, r1, r2, r3); - r0 *= rescale; r1 *= rescale; r2 *= rescale; r3 *= rescale; - tmem_store_col(to + col, 0, r0, r1, r2, r3); + // D1.5 FIX: Rescale O in TMEM (TMEM → regs → multiply → TMEM) + // This is the one-way path using SAME tcgen05.ld + tcgen05.st pair + for (int col=0; col -#include -#include +#include "fmha_common.cuh" namespace dsv4::kernels::attention { -typedef unsigned short bf16_t; - -__device__ __forceinline__ bf16_t f32_to_bf16(float f) { - bf16_t h; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); return h; -} -__device__ __forceinline__ float bf16_to_f32(bf16_t h) { - float f; asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); return f; -} - -constexpr int WARP = 32; -constexpr int NTHREADS = 192; -constexpr int NWARPS = 6; - template __global__ void __launch_bounds__(NTHREADS) fmha_decode_ref( - const bf16_t* __restrict__ q, - const bf16_t* __restrict__ k, - const bf16_t* __restrict__ v, - bf16_t* __restrict__ o, + const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, + const bf16_t* __restrict__ v, bf16_t* __restrict__ o, int bstride_q, int bstride_kv, int bstride_o, - int s_k, int n_comp, int swa_len, - float scale, - const float* __restrict__ attn_sink, - float* __restrict__ lse_out + int s_k, int n_comp, int swa_len, float scale, + const float* __restrict__ attn_sink, float* __restrict__ lse_out ) { - const int head = blockIdx.y; - const int batch = blockIdx.z; - const int tid = threadIdx.x; - const int wid = tid / WARP; - const int lane = tid % WARP; + const int head = blockIdx.y, batch = blockIdx.z, tid = threadIdx.x; + const bf16_t* qh = q + batch*bstride_q + head*HD; + const bf16_t* kb = k + batch*bstride_kv; + const bf16_t* vb = v + batch*bstride_kv; + bf16_t* oh = o + batch*bstride_o + head*HD; - // Pointers - const bf16_t* qh = q + batch * bstride_q + head * HD; - const bf16_t* kb = k + batch * bstride_kv; - const bf16_t* vb = v + batch * bstride_kv; - bf16_t* oh = o + batch * bstride_o + head * HD; - - // Load Q into registers (T=1 decode, HD values) - float q_local[HD > 64 ? 1 : HD]; // Can't VLA in CUDA. Use SMEM instead. - - // Use SMEM for Q (shared across all threads) extern __shared__ char sbuf[]; - float* sQ = (float*)sbuf; // HD floats - float* sO = (float*)(sbuf + HD * sizeof(float)); // HD floats (output accumulator) + float* sQ = (float*)sbuf; + float* sO = (float*)(sbuf + HD*sizeof(float)); - // Load Q to SMEM - for (int d = tid; d < HD; d += NTHREADS) { - sQ[d] = bf16_to_f32(qh[d]); - } - // Initialize O accumulator - for (int d = tid; d < HD; d += NTHREADS) { - sO[d] = 0.0f; - } + for (int d=tid; d 32 ? 1 : HD)] = 0.0f; - - // Actually, each lane accumulates a different subset of HD - // lane d accumulates O[d], O[d+32], O[d+64], etc. - // But HD might be 64, so each lane handles 2 elements - - // Per-thread O accumulator - float my_o[4]; // max elements per thread at HD=64: 64/192 < 1, use SMEM - // Actually: accumulate O in SMEM atomically, or use a tree reduction - - // Simplest correct approach: each thread processes its KV range, - // computes P@V for those positions, accumulates to SMEM O with atomics - // (or just sequential within warp, then warp-reduce) - - // For correctness first, let's just have ONE thread do everything - // (slow but correct), then parallelize. - + float row_max = -INFINITY, row_sum = 0.0f; if (tid == 0) { - float o_acc[HD]; - for (int d = 0; d < HD; d++) o_acc[d] = 0.0f; - - for (int c = 0; c < s_k; c++) { + for (int c=0; c 0 && c >= n_comp + swa_len) s_val = -INFINITY; - - // Online softmax with O rescale + if (swa_len>0 && c>=n_comp+swa_len) s_val = -INFINITY; + float new_max = fmaxf(row_max, s_val); if (new_max > row_max) { float rescale = expf(row_max - new_max); - for (int d = 0; d < HD; d++) o_acc[d] *= rescale; - row_sum *= rescale; - row_max = new_max; + for (int d=0; d