FMHA SM100: Refactor into common + reference + TMEM epilogue headers
- fmha_common.cuh: BF16, TMEM ops, warp reductions (shared) - fmha_sm100.cuh: Phase 1 reference (SMEM-based, cos 0.999999) - fmha_epilogue_sm100.cuh: Phase 2 TMEM+correction epilogue (Priority 2) - Test both kernels at hd=64 and hd=128
This commit is contained in:
61
dsv4/kernels/attention/fmha_common.cuh
Normal file
61
dsv4/kernels/attention/fmha_common.cuh
Normal file
@@ -0,0 +1,61 @@
|
||||
/**
|
||||
* DSV4 FMHA shared definitions — base header.
|
||||
* BF16 type, TMEM ops, warp reductions, constants.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
|
||||
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
|
||||
@@ -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 <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
#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<int HD, int TILE_ROWS = 128, bool NORMALIZE = true>
|
||||
__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<bf16_t*>(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<int HD>
|
||||
__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; d<HD; d+=NTHREADS) sQ[d] = bf16_to_f32(qh[d]);
|
||||
__syncthreads();
|
||||
|
||||
// Online softmax with O rescale in TMEM
|
||||
float row_max = -INFINITY;
|
||||
float row_sum = 0.0f;
|
||||
|
||||
// Initialize TMEM O to zero
|
||||
for (int col = tid; col < tmem_o_cols; col += NTHREADS) {
|
||||
for (int rg = 0; rg < 8; rg++) { // 8 row-groups of 16 rows each
|
||||
tmem_store_col(to + col, rg, 0.0f, 0.0f, 0.0f, 0.0f);
|
||||
}
|
||||
// Init TMEM O to zero
|
||||
for (int col=tid; col<tmem_o_cols; col+=NTHREADS) {
|
||||
for (int rg=0; rg<8; rg++) tmem_store_col(to+col, rg, 0,0,0,0);
|
||||
}
|
||||
tmem_fence();
|
||||
__syncthreads();
|
||||
tmem_fence(); __syncthreads();
|
||||
|
||||
// Process KV positions (single-thread for Phase 2 correctness)
|
||||
float row_max = -INFINITY, row_sum = 0.0f;
|
||||
|
||||
// Single-thread softmax + P@V with TMEM O accumulation
|
||||
if (tid == 0) {
|
||||
for (int c = 0; c < s_k; c++) {
|
||||
for (int c=0; c<s_k; c++) {
|
||||
float s_val = 0.0f;
|
||||
for (int d = 0; d < HD; d++) {
|
||||
s_val += sQ[d] * bf16_to_f32(kb[c * HD + d]);
|
||||
}
|
||||
for (int d=0; d<HD; d++) s_val += sQ[d] * bf16_to_f32(kb[c*HD+d]);
|
||||
s_val *= scale;
|
||||
if (swa_len>0 && 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<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);
|
||||
}
|
||||
tmem_fence();
|
||||
|
||||
row_sum *= rescale;
|
||||
row_max = new_max;
|
||||
row_sum *= rescale; row_max = new_max;
|
||||
}
|
||||
|
||||
float p_val = expf(s_val - row_max);
|
||||
row_sum += p_val;
|
||||
|
||||
// P@V: accumulate p_val * V[:, c] into TMEM O
|
||||
for (int col = 0; col < tmem_o_cols; col++) {
|
||||
int d0 = col * 4 + 0;
|
||||
int d1 = col * 4 + 1;
|
||||
int d2 = col * 4 + 2;
|
||||
int d3 = col * 4 + 3;
|
||||
// P@V: accumulate p_val * V[:,c] into TMEM O
|
||||
for (int col=0; col<tmem_o_cols; col++) {
|
||||
int d0=col*4+0, d1=col*4+1, d2=col*4+2, d3=col*4+3;
|
||||
float v0=(d0<HD)?bf16_to_f32(vb[d0*s_k+c]):0.0f;
|
||||
float v1=(d1<HD)?bf16_to_f32(vb[d1*s_k+c]):0.0f;
|
||||
float v2=(d2<HD)?bf16_to_f32(vb[d2*s_k+c]):0.0f;
|
||||
float v3=(d3<HD)?bf16_to_f32(vb[d3*s_k+c]):0.0f;
|
||||
|
||||
float v0 = (d0 < HD) ? bf16_to_f32(vb[d0 * s_k + c]) : 0.0f;
|
||||
float v1 = (d1 < HD) ? bf16_to_f32(vb[d1 * s_k + c]) : 0.0f;
|
||||
float v2 = (d2 < HD) ? bf16_to_f32(vb[d2 * s_k + c]) : 0.0f;
|
||||
float v3 = (d3 < HD) ? bf16_to_f32(vb[d3 * s_k + c]) : 0.0f;
|
||||
|
||||
// Load current O, add p*V, store back
|
||||
float r0, r1, r2, r3;
|
||||
tmem_load_col(to + col, 0, r0, r1, r2, r3);
|
||||
r0 += p_val * v0;
|
||||
r1 += p_val * v1;
|
||||
r2 += p_val * v2;
|
||||
r3 += p_val * v3;
|
||||
tmem_store_col(to + col, 0, r0, r1, r2, r3);
|
||||
float r0,r1,r2,r3;
|
||||
tmem_load_col(to+col, 0, r0,r1,r2,r3);
|
||||
r0+=p_val*v0; r1+=p_val*v1; r2+=p_val*v2; r3+=p_val*v3;
|
||||
tmem_store_col(to+col, 0, r0,r1,r2,r3);
|
||||
}
|
||||
tmem_fence();
|
||||
}
|
||||
sRowSums[0] = row_sum;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Store row_sum for the epilogue
|
||||
if (tid == 0) sRowSums[0] = row_sum;
|
||||
__syncthreads();
|
||||
|
||||
// =================================================================
|
||||
// One-way Correction Epilogue: TMEM → regs → normalize → BF16 → GMEM
|
||||
// =================================================================
|
||||
// This is the key pattern from the MoE kernel, adapted for FMHA:
|
||||
// 1. tcgen05.ld: Load O from TMEM to registers (FP32)
|
||||
// 2. Divide by row_sum in registers (normalize)
|
||||
// 3. cvt.rn.bf16.f32: Cast to BF16 in registers
|
||||
// 4. st.global: Write to GMEM
|
||||
// ================================================================
|
||||
// One-way Correction Epilogue
|
||||
// TMEM → regs → normalize → BF16 → GMEM
|
||||
//
|
||||
// Future (NVFP4-1.2): Step 2.5 adds FP4 amax + pack in registers
|
||||
// Future (D2 multi-CTA): Step 4 uses TMA store with flat_divide coordinates
|
||||
|
||||
// This is the MoE epilogue pattern adapted for FMHA:
|
||||
// 1. tcgen05.ld: Load O from TMEM to registers (FP32)
|
||||
// 2. O[i] /= row_sum (normalize in registers)
|
||||
// 3. cvt.rn.bf16.f32: Cast to BF16
|
||||
// 4. st.global: Write to GMEM
|
||||
// ================================================================
|
||||
if (tid == 0) {
|
||||
float inv_sum = 1.0f / sRowSums[0];
|
||||
for (int col = 0; col < tmem_o_cols; col++) {
|
||||
float r0, r1, r2, r3;
|
||||
tmem_load_col(to + col, 0, r0, r1, r2, r3);
|
||||
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*=inv_sum; r1*=inv_sum; r2*=inv_sum; r3*=inv_sum;
|
||||
|
||||
// Normalize
|
||||
r0 *= inv_sum; r1 *= inv_sum; r2 *= inv_sum; r3 *= inv_sum;
|
||||
|
||||
// Cast to BF16 and write to GMEM
|
||||
int d0 = col * 4 + 0;
|
||||
int d1 = col * 4 + 1;
|
||||
int d2 = col * 4 + 2;
|
||||
int d3 = col * 4 + 3;
|
||||
if (d0 < HD) oh[d0] = f32_to_bf16(r0);
|
||||
if (d1 < HD) oh[d1] = f32_to_bf16(r1);
|
||||
if (d2 < HD) oh[d2] = f32_to_bf16(r2);
|
||||
if (d3 < HD) oh[d3] = f32_to_bf16(r3);
|
||||
int d0=col*4+0, d1=col*4+1, d2=col*4+2, d3=col*4+3;
|
||||
if(d0<HD) oh[d0]=f32_to_bf16(r0);
|
||||
if(d1<HD) oh[d1]=f32_to_bf16(r1);
|
||||
if(d2<HD) oh[d2]=f32_to_bf16(r2);
|
||||
if(d3<HD) oh[d3]=f32_to_bf16(r3);
|
||||
}
|
||||
}
|
||||
|
||||
// LSE
|
||||
if (lse_out && tid == 0) {
|
||||
lse_out[batch * gridDim.y + head] = logf(row_sum) + row_max;
|
||||
}
|
||||
|
||||
// TMEM dealloc
|
||||
if (wid == 0 && lane == 0) tmem_dealloc(tb, tmem_n);
|
||||
if(lse_out && tid==0) lse_out[batch*gridDim.y+head] = logf(row_sum) + row_max;
|
||||
if(wid==0 && lane==0) tmem_dealloc(tb, tmem_n);
|
||||
}
|
||||
|
||||
} // namespace dsv4::kernels::attention
|
||||
} // namespace
|
||||
|
||||
@@ -1,161 +1,59 @@
|
||||
/**
|
||||
* DSV4 FMHA Decode Kernel — Phase 1 Reference
|
||||
*
|
||||
* Correct scalar implementation. Each CTA processes one (batch, head).
|
||||
* All 192 threads cooperate on the softmax and PV for T=1 decode.
|
||||
*
|
||||
* Strategy: Each thread independently computes S for a subset of KV positions,
|
||||
* does online softmax with O rescale, then a parallel reduction for P@V.
|
||||
* DSV4 FMHA Phase 1 Reference — scalar implementation.
|
||||
* Uses SMEM for Q and O. Single-thread for correctness.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
#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<int HD>
|
||||
__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<HD; d+=NTHREADS) sQ[d] = bf16_to_f32(qh[d]);
|
||||
for (int d=tid; d<HD; d+=NTHREADS) sO[d] = 0.0f;
|
||||
__syncthreads();
|
||||
|
||||
// Online softmax: process KV in blocks
|
||||
float row_max = -INFINITY;
|
||||
float row_sum = 0.0f;
|
||||
|
||||
// Each thread processes s_k/NTHREADS KV positions
|
||||
// For s_k=128, NTHREADS=192: most threads get 0-1 positions
|
||||
// Better: have each thread process a range, accumulate locally, then reduce
|
||||
|
||||
// Simpler: warp-level processing
|
||||
// Each warp processes s_k/NWARPS KV positions
|
||||
// Then warp-reduce for softmax state
|
||||
|
||||
int kv_per_warp = (s_k + NWARPS - 1) / NWARPS;
|
||||
int my_kv_start = wid * kv_per_warp;
|
||||
int my_kv_end = min(my_kv_start + kv_per_warp, s_k);
|
||||
|
||||
// Warp-local softmax state
|
||||
float warp_max = -INFINITY;
|
||||
float warp_sum = 0.0f;
|
||||
float warp_o[HD]; // Warp-local O accumulation
|
||||
for (int d = lane; d < HD; d += WARP) warp_o[d % (HD > 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<s_k; c++) {
|
||||
float s_val = 0.0f;
|
||||
for (int d = 0; d < HD; d++) {
|
||||
s_val += sQ[d] * bf16_to_f32(kb[c * HD + d]);
|
||||
}
|
||||
for (int d=0; d<HD; d++) s_val += sQ[d] * bf16_to_f32(kb[c*HD+d]);
|
||||
s_val *= scale;
|
||||
|
||||
// D3: SWA mask
|
||||
if (swa_len > 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<HD; d++) sO[d] *= rescale;
|
||||
row_sum *= rescale; row_max = new_max;
|
||||
}
|
||||
|
||||
float p_val = expf(s_val - row_max);
|
||||
row_sum += p_val;
|
||||
|
||||
for (int d = 0; d < HD; d++) {
|
||||
o_acc[d] += p_val * bf16_to_f32(vb[d * s_k + c]);
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize
|
||||
for (int d = 0; d < HD; d++) {
|
||||
sO[d] = o_acc[d] / row_sum;
|
||||
for (int d=0; d<HD; d++) sO[d] += p_val * bf16_to_f32(vb[d*s_k+c]);
|
||||
}
|
||||
for (int d=0; d<HD; d++) sO[d] /= row_sum;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Write output
|
||||
for (int d = tid; d < HD; d += NTHREADS) {
|
||||
oh[d] = f32_to_bf16(sO[d]);
|
||||
}
|
||||
|
||||
// LSE
|
||||
if (lse_out && tid == 0) {
|
||||
lse_out[batch * gridDim.y + head] = logf(row_sum) + row_max;
|
||||
}
|
||||
for (int d=tid; d<HD; d+=NTHREADS) oh[d] = f32_to_bf16(sO[d]);
|
||||
if (lse_out && tid==0) lse_out[batch*gridDim.y+head] = logf(row_sum) + row_max;
|
||||
}
|
||||
|
||||
} // namespace dsv4::kernels::attention
|
||||
} // namespace
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
* Standalone CUDA test for FMHA SM100 — Reference + TMEM kernels.
|
||||
* Tests both the Phase 1 reference and Phase 2 TMEM+epilogue kernels.
|
||||
*/
|
||||
#include "dsv4/kernels/attention/fmha_common.cuh"
|
||||
#include "dsv4/kernels/attention/fmha_sm100.cuh"
|
||||
#include "dsv4/kernels/attention/fmha_epilogue_sm100.cuh"
|
||||
#include <stdio.h>
|
||||
|
||||
Reference in New Issue
Block a user