fmha_multirow: rewrite with 32x32b.x8 only, no s_p_vals, row_page addressing
- Kill 64KB s_p_vals buffer — P is streamed per K-tile through sPk - All TMEM ops use 32x32b.x8 exclusively (16x256b.x1 crashes on 2nd call) - T>32: 4 softmax warps use row_page offset in TMEM address (row<<16) - Lane l in warp w handles row w*32+l - Two-pass softmax: pass 1 row_max, pass 2 exp/sum interleaved with PV - PV: N=16 sub-tiles, SS MMA sPk(128,16) × sV(16,16) → TMEM - Epilogue: 32x32b.x8 TMEM read, normalize, BF16 → GMEM - SMEM budget: ~14KB (well within 232KB)
This commit is contained in:
@@ -1,11 +1,57 @@
|
||||
/**
|
||||
* DSV4 FMHA — 6-warp specialized kernel, multi-row softmax (prefill T>1).
|
||||
*
|
||||
* T <= 32: Warp 0 only, 32x32b.x8, lane l = row l.
|
||||
* T > 32: 4 softmax warps, 16x256b.x1, two-pass online softmax.
|
||||
* Warp w handles rows [w*32, min((w+1)*32, T)).
|
||||
* Lane j reads rows j*4+0..3 per TMEM column.
|
||||
* Active lanes for warp w: j in [w*8, w*8+8).
|
||||
* ==================================================================
|
||||
* DESIGN
|
||||
* ==================================================================
|
||||
*
|
||||
* 6-warp CTA: warps 0-3 = softmax, warp 4 = MMA, warp 5 = TMA load.
|
||||
* Grid: (1, n_h, batch) — each CTA processes one head of one batch item.
|
||||
*
|
||||
* TMEM operations ALL use 32x32b.x8 format exclusively.
|
||||
* The 16x256b.x1 format crashes on the 2nd call — verified on B200.
|
||||
*
|
||||
* TMEM ADDRESSING for 32x32b.x8:
|
||||
* addr = tmem_base + (row_page << 16) + col_group * 8
|
||||
* - col_group selects 8 consecutive TMEM columns
|
||||
* - row_page selects a 32-row "page" (row_page=0 → rows 0-31, =1 → rows 32-63, etc.)
|
||||
* - Lane l within the warp reads row (row_page*32 + l) across the 8 columns
|
||||
* - ALL 32 lanes must call with the same address (warp-collective)
|
||||
*
|
||||
* SOFTMAX (two regimes):
|
||||
* - T <= 32: warp 0 only. row_page=0. Lane l processes row l. 32x32b.x8 reads.
|
||||
* - T > 32: 4 softmax warps. Warp w uses row_page=w.
|
||||
* Lane l in warp w processes row w*32+l. Two-pass:
|
||||
* (1) row_max, (2) exp+sum, stream P to sPk per K-tile.
|
||||
*
|
||||
* P STAGING FOR PV (no s_p_vals buffer):
|
||||
* P is written one K-tile at a time to sPk (128, 16) SMEM in canonical
|
||||
* layout. Softmax warps compute P values in registers, write to sPk.
|
||||
* Then MMA warp does PV with sPk × sV.
|
||||
* The softmax and PV are interleaved K-tile-by-K-tile (required because
|
||||
* we can't hold the full 128×128 P matrix in SMEM).
|
||||
*
|
||||
* PV GEMM:
|
||||
* N=16 sub-tiles (avoids Layout D N=64 bug). SS MMA: sPk(128,16) × sV(16,16).
|
||||
* Accumulate into TMEM across K-tiles and N-sub-tiles.
|
||||
*
|
||||
* EPILOGUE:
|
||||
* 32x32b.x8 TMEM read → normalize (1/row_sum) → BF16 → GMEM.
|
||||
* LSE = ln(row_sum) + row_max for D5 KV merge.
|
||||
*
|
||||
* ==================================================================
|
||||
* SMEM BUDGET (HD=64)
|
||||
* ==================================================================
|
||||
* sTmemBase: 8 B
|
||||
* sRowMax: 128 × 4 = 512 B
|
||||
* sRowSum: 128 × 4 = 512 B
|
||||
* padding: 128 B (128B alignment for UMMA)
|
||||
* sQ0: 128 × 16 × 2 = 4096 B (one K-tile)
|
||||
* sK0: 128 × 16 × 2 = 4096 B (one K-tile)
|
||||
* sPk: 128 × 16 × 2 = 4096 B (one K-tile P buffer)
|
||||
* sV: 16 × 16 × 2 = 512 B (one N-sub-tile, one K-tile)
|
||||
* Total: ~14 KB — well within 232 KB budget
|
||||
* ==================================================================
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
@@ -34,13 +80,14 @@ struct FmhaMultiRowParams {
|
||||
template<int HD, int SK_TILE = 128>
|
||||
__global__ void __launch_bounds__(192)
|
||||
fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
|
||||
static constexpr int NKT_QK = HD / MMA_K_BF16;
|
||||
static constexpr int NKT_PV = SK_TILE / MMA_K_BF16;
|
||||
static constexpr int N_NSUB = HD / 16;
|
||||
static constexpr int TILE_SZ = 128 * MMA_K_BF16;
|
||||
static constexpr int V_SUB_SZ = 256;
|
||||
static constexpr int NKT_QK = HD / MMA_K_BF16; // K-tiles for QK GEMM
|
||||
static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; // K-tiles for PV GEMM (= 8)
|
||||
static constexpr int N_NSUB = HD / 16; // N-sub-tiles for PV (avoids N=64 bug)
|
||||
static constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16 elements
|
||||
static constexpr int V_SUB_SZ = 16 * MMA_K_BF16; // 256 BF16 elements (16, 16) canonical
|
||||
static constexpr int TMEM_N = (HD <= 128) ? 128 : 256;
|
||||
static constexpr int MAX_ROWS = 128;
|
||||
static constexpr int CORES_MN = 128 / 8; // 16 core matrices along M
|
||||
|
||||
const int head_idx = blockIdx.y;
|
||||
const int batch_idx = blockIdx.z;
|
||||
@@ -53,56 +100,81 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
|
||||
const int T = params.T;
|
||||
const int s_k = params.s_k;
|
||||
const float scale = params.scale;
|
||||
const bool full_tile = (T > 32);
|
||||
|
||||
// Per-head GMEM pointers
|
||||
const bf16_t* __restrict__ q_head = params.q + head_idx * params.q_head_stride + batch_idx * params.q_batch_stride;
|
||||
const bf16_t* __restrict__ k_head = params.k + head_idx * params.k_head_stride + batch_idx * params.k_batch_stride;
|
||||
const bf16_t* __restrict__ v_head = params.v + head_idx * params.v_head_stride + batch_idx * params.v_batch_stride;
|
||||
bf16_t* __restrict__ o_head = params.o + head_idx * params.o_head_stride + batch_idx * params.o_batch_stride;
|
||||
float* __restrict__ lse_head = params.lse ? params.lse + head_idx * params.lse_head_stride + batch_idx * params.lse_batch_stride : nullptr;
|
||||
|
||||
// SMEM
|
||||
// ================================================================
|
||||
// SMEM allocation — NO s_p_vals buffer
|
||||
// ================================================================
|
||||
extern __shared__ char sbuf[];
|
||||
uint32_t* sTmemBase = (uint32_t*)sbuf;
|
||||
float* sRowMax = (float*)(sbuf + 4); // [MAX_ROWS]
|
||||
float* sRowSum = sRowMax + MAX_ROWS; // [MAX_ROWS]
|
||||
bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 15) & ~(uintptr_t)15);
|
||||
bf16_t* sK0 = sQ0 + TILE_SZ;
|
||||
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); // [MAX_ROWS][SK_TILE]
|
||||
uint32_t* sTmemBase = (uint32_t*)sbuf; // 4 B (8B-aligned)
|
||||
float* sRowMax = (float*)(sbuf + 8); // [MAX_ROWS] = 512 B
|
||||
float* sRowSum = sRowMax + MAX_ROWS; // [MAX_ROWS] = 512 B
|
||||
// Align sQ0 to 128B for UMMA descriptors
|
||||
bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 127) & ~(uintptr_t)127);
|
||||
bf16_t* sK0 = sQ0 + TILE_SZ; // (128, 16) BF16 canonical
|
||||
bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + TILE_SZ) + 127) & ~(uintptr_t)127); // (128, 16) BF16 canonical
|
||||
bf16_t* sV = (bf16_t*)(((uintptr_t)(sPk + TILE_SZ) + 127) & ~(uintptr_t)127); // (16, 16) BF16 canonical
|
||||
|
||||
// TMEM alloc
|
||||
if (is_mma_warp) { uint32_t p = __cvta_generic_to_shared(sTmemBase); tmem_alloc(p, TMEM_N); }
|
||||
// ================================================================
|
||||
// TMEM allocation (warp 4)
|
||||
// ================================================================
|
||||
if (is_mma_warp) {
|
||||
uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
|
||||
tmem_alloc(smem_ptr, TMEM_N);
|
||||
}
|
||||
__syncthreads();
|
||||
uint32_t tb = *sTmemBase;
|
||||
|
||||
// ================================================================
|
||||
// QK GEMM
|
||||
// Row assignment for softmax
|
||||
// ================================================================
|
||||
// T <= 32: only warp 0, row_page = 0, lane l = row l
|
||||
// T > 32: 4 softmax warps, warp w uses row_page = w
|
||||
// lane l in warp w handles row w*32 + l
|
||||
const bool my_warp_active = (T <= 32) ? (wid == 0) : is_softmax_warp;
|
||||
const int my_row_page = (T <= 32) ? 0 : wid;
|
||||
const int my_row = my_row_page * 32 + lane;
|
||||
const bool my_row_active = my_warp_active && (my_row < T);
|
||||
const uint32_t my_tmem_row_offset = (uint32_t)(my_row_page * 32) << 16;
|
||||
|
||||
// ================================================================
|
||||
// QK GEMM: accumulate all K-tiles into TMEM
|
||||
// ================================================================
|
||||
for (int kt = 0; kt < NKT_QK; kt++) {
|
||||
if (is_load_warp) {
|
||||
constexpr int CORES_MN = 128 / 8;
|
||||
// Load Q K-tile: Q is (T, HD), columns [kt*16, kt*16+16)
|
||||
for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0;
|
||||
for (int r = 0; r < T; r++) {
|
||||
for (int d = lane; d < MMA_K_BF16; d += 32) {
|
||||
int full_d = kt * MMA_K_BF16 + d;
|
||||
int ck = d / 8, lc = d % 8;
|
||||
int core_mn = r / 8, local_r = r % 8;
|
||||
sQ0[ck * CORES_MN * 64 + core_mn * 64 + local_r * 8 + lc] = q_head[r * HD + full_d];
|
||||
if (full_d < HD) {
|
||||
int ck = d / 8, lc = d % 8;
|
||||
int core_mn = r / 8, local_r = r % 8;
|
||||
sQ0[ck * CORES_MN * 64 + core_mn * 64 + local_r * 8 + lc] = q_head[r * HD + full_d];
|
||||
}
|
||||
}
|
||||
}
|
||||
// Load K K-tile: K is (s_k, HD), columns [kt*16, kt*16+16)
|
||||
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 full_d = kt * MMA_K_BF16 + d;
|
||||
int ck = d / 8, lc = d % 8;
|
||||
int tmn = r / 8, lr = r % 8;
|
||||
sK0[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = k_head[r * HD + full_d];
|
||||
if (full_d < HD) {
|
||||
int ck = d / 8, lc = d % 8;
|
||||
int tmn = r / 8, lr = r % 8;
|
||||
sK0[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = k_head[r * HD + full_d];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
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);
|
||||
@@ -113,189 +185,198 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) {
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// TMEM now contains S[128, s_k] — the raw QK dot products.
|
||||
// S[row, col] needs scaling by `scale` before softmax.
|
||||
|
||||
// ================================================================
|
||||
// SOFTMAX
|
||||
// SOFTMAX PASS 1: compute row_max for each row
|
||||
// ================================================================
|
||||
if (!full_tile) {
|
||||
// T <= 32: warp 0, 32x32b.x8, lane l = row l
|
||||
if (wid == 0) {
|
||||
float s_vals[SK_TILE], row_max = -INFINITY;
|
||||
for (int n = 0; n < SK_TILE / 8; n++) {
|
||||
float tmp[8];
|
||||
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
|
||||
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
|
||||
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
|
||||
: "r"(tb + n * 8));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;");
|
||||
if (lane < T) {
|
||||
for (int c = 0; c < 8; c++) {
|
||||
s_vals[n * 8 + c] = tmp[c] * scale;
|
||||
row_max = fmaxf(row_max, s_vals[n * 8 + c]);
|
||||
// Read all SK_TILE columns (16 groups of 8 columns each).
|
||||
// 32x32b.x8: addr = tb + (row_page*32 << 16) + n*8
|
||||
// Each lane reads its assigned row across 8 columns per read.
|
||||
float my_row_max = -INFINITY;
|
||||
const int NUM_READS = SK_TILE / 8; // = 16
|
||||
|
||||
if (my_warp_active) {
|
||||
for (int n = 0; n < NUM_READS; n++) {
|
||||
float tmp[8];
|
||||
uint32_t addr = tb + my_tmem_row_offset + n * 8;
|
||||
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
|
||||
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
|
||||
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
|
||||
: "r"(addr));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;");
|
||||
|
||||
if (my_row_active) {
|
||||
for (int c = 0; c < 8; c++) {
|
||||
int col = n * 8 + c;
|
||||
if (col < s_k) {
|
||||
float val = tmp[c] * scale;
|
||||
my_row_max = fmaxf(my_row_max, val);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (lane < T) sRowMax[lane] = row_max;
|
||||
float row_sum = 0.0f;
|
||||
if (lane < T) {
|
||||
for (int j = 0; j < SK_TILE; j++) { s_vals[j] = expf(s_vals[j] - row_max); row_sum += s_vals[j]; }
|
||||
}
|
||||
if (lane < T) sRowSum[lane] = row_sum;
|
||||
if (lane < T) {
|
||||
float inv = 1.0f / row_sum;
|
||||
for (int j = 0; j < SK_TILE; j++) s_p_vals[lane * SK_TILE + j] = s_vals[j] * inv;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// T > 32: 4 warps, 16x256b.x1, two-pass softmax
|
||||
if (is_softmax_warp) {
|
||||
const int warp_row_start = wid * 32;
|
||||
const int warp_row_end = min((wid + 1) * 32, T);
|
||||
const int lane_row_base = lane * 4;
|
||||
const bool active = (lane_row_base >= warp_row_start) && (lane_row_base < warp_row_end);
|
||||
const int vrows = active ? min(4, warp_row_end - lane_row_base) : 0;
|
||||
|
||||
float rmax[4], rsum[4];
|
||||
for (int r = 0; r < 4; r++) rmax[r] = -INFINITY;
|
||||
|
||||
// Pass 1: row_max
|
||||
for (int col = 0; col < SK_TILE; col++) {
|
||||
float v[4];
|
||||
asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0,%1,%2,%3},[%4,%5];"
|
||||
: "=f"(v[0]), "=f"(v[1]), "=f"(v[2]), "=f"(v[3]) : "r"(tb), "r"(col));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;");
|
||||
if (active) for (int r = 0; r < vrows; r++) rmax[r] = fmaxf(rmax[r], v[r] * scale);
|
||||
}
|
||||
|
||||
// Pass 2: exp, sum, write P
|
||||
for (int r = 0; r < 4; r++) rsum[r] = 0.0f;
|
||||
for (int col = 0; col < SK_TILE; col++) {
|
||||
float v[4];
|
||||
asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0,%1,%2,%3},[%4,%5];"
|
||||
: "=f"(v[0]), "=f"(v[1]), "=f"(v[2]), "=f"(v[3]) : "r"(tb), "r"(col));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;");
|
||||
if (active) {
|
||||
for (int r = 0; r < vrows; r++) {
|
||||
int gr = lane_row_base + r;
|
||||
float p = expf(v[r] * scale - rmax[r]);
|
||||
rsum[r] += p;
|
||||
s_p_vals[gr * SK_TILE + col] = p;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Normalize and store max/sum
|
||||
if (active) {
|
||||
for (int r = 0; r < vrows; r++) {
|
||||
int gr = lane_row_base + r;
|
||||
sRowMax[gr] = rmax[r];
|
||||
sRowSum[gr] = rsum[r];
|
||||
float inv = 1.0f / rsum[r];
|
||||
for (int j = 0; j < SK_TILE; j++) s_p_vals[gr * SK_TILE + j] *= inv;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Store row_max to SMEM
|
||||
if (my_row_active) sRowMax[my_row] = my_row_max;
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// PV GEMM
|
||||
// SOFTMAX PASS 2 + PV GEMM (interleaved per PV K-tile)
|
||||
// ================================================================
|
||||
for (int n = 0; n < N_NSUB; n++) {
|
||||
int d_base = n * 16;
|
||||
for (int kt = 0; kt < NKT_PV; kt++) {
|
||||
if (is_load_warp) {
|
||||
constexpr int CORES_MN = 128 / 8;
|
||||
for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0;
|
||||
for (int r = 0; r < T; r++) {
|
||||
for (int c = lane; c < MMA_K_BF16; c += 32) {
|
||||
int gc = kt * MMA_K_BF16 + c;
|
||||
float pv = s_p_vals[r * SK_TILE + gc];
|
||||
int cmn = r / 8, lr = r % 8, ck = c / 8, lc = c % 8;
|
||||
sPk[ck * CORES_MN * 64 + cmn * 64 + lr * 8 + lc] = f32_to_bf16(pv);
|
||||
// For each PV K-tile (16 columns of S):
|
||||
// 1. Softmax warps re-read TMEM, compute exp(S*scale - row_max)
|
||||
// 2. Write P to sPk in canonical layout
|
||||
// 3. Load V sub-tile to sV
|
||||
// 4. MMA does PV: sPk(128,16) × sV(16,16) → TMEM
|
||||
//
|
||||
// We accumulate row_sum across K-tiles in registers.
|
||||
// P written to sPk is UN-normalized (exp(S*scale - row_max)).
|
||||
// Normalization happens in the epilogue (divide by row_sum).
|
||||
// ================================================================
|
||||
float my_row_sum = 0.0f;
|
||||
|
||||
for (int pv_kt = 0; pv_kt < NKT_PV; pv_kt++) {
|
||||
const int col_start = pv_kt * MMA_K_BF16;
|
||||
const int col_end = col_start + MMA_K_BF16;
|
||||
// TMEM column groups for this K-tile: [col_start/8, col_end/8)
|
||||
const int rg_start = col_start / 8;
|
||||
const int rg_end = col_end / 8;
|
||||
|
||||
// P values for this K-tile (un-normalized)
|
||||
float p_vals[16]; // MMA_K_BF16 values
|
||||
for (int c = 0; c < MMA_K_BF16; c++) p_vals[c] = 0.0f;
|
||||
|
||||
if (my_warp_active) {
|
||||
float rm = my_row_active ? sRowMax[my_row] : 0.0f;
|
||||
|
||||
// Read TMEM for this K-tile's columns
|
||||
for (int n = rg_start; n < rg_end; n++) {
|
||||
float tmp[8];
|
||||
uint32_t addr = tb + my_tmem_row_offset + n * 8;
|
||||
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
|
||||
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
|
||||
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
|
||||
: "r"(addr));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;");
|
||||
|
||||
if (my_row_active) {
|
||||
for (int c = 0; c < 8; c++) {
|
||||
int col = n * 8 + c;
|
||||
if (col < s_k && col >= col_start && col < col_end) {
|
||||
float val = expf(tmp[c] * scale - rm);
|
||||
p_vals[col - col_start] = val;
|
||||
my_row_sum += val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Zero sPk first (load warp does this for all rows)
|
||||
if (is_load_warp) {
|
||||
for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Write P to sPk in canonical layout (softmax warps)
|
||||
if (my_row_active) {
|
||||
for (int c = 0; c < MMA_K_BF16; c++) {
|
||||
int ck = c / 8, lc = c % 8;
|
||||
int core_mn = my_row / 8, local_r = my_row % 8;
|
||||
sPk[ck * CORES_MN * 64 + core_mn * 64 + local_r * 8 + lc] = f32_to_bf16(p_vals[c]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// PV GEMM: N=16 sub-tiles
|
||||
for (int n_sub = 0; n_sub < N_NSUB; n_sub++) {
|
||||
int d_base = n_sub * 16;
|
||||
|
||||
if (is_load_warp) {
|
||||
// Load V sub-tile: V is (HD, s_k) in GMEM
|
||||
// Need V[d_base..d_base+15, col_start..col_end-1]
|
||||
// Canonical layout for (16, 16):
|
||||
// g_mn = d/8, g_k = lr/8, llr = d%8, lc = lr%8
|
||||
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++) {
|
||||
int r = kt * MMA_K_BF16 + lr;
|
||||
int gmn = dd / 8, gk = lr / 8, llr = dd % 8, lc = lr % 8;
|
||||
sV[gk * 2 * 64 + gmn * 64 + llr * 8 + lc] = v_head[(d_base + dd) * s_k + r];
|
||||
int r = col_start + lr;
|
||||
if (r < s_k && (d_base + dd) < HD) {
|
||||
int g_mn = dd / 8, g_k = lr / 8;
|
||||
int llr = dd % 8, lc = lr % 8;
|
||||
sV[g_k * 2 * 64 + g_mn * 64 + llr * 8 + lc] = v_head[(d_base + dd) * s_k + r];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (is_mma_warp) {
|
||||
uint32_t idesc16 = make_idesc(128, 16);
|
||||
uint32_t idesc_pv = make_idesc(128, 16);
|
||||
uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), 128);
|
||||
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
|
||||
if (tid == 128) umma_ss_f16(tb + n * 16, dp, dv, idesc16, kt > 0);
|
||||
// accumulate = true for pv_kt > 0 (K-tile accumulation)
|
||||
// For n_sub > 0, the PV result goes to different TMEM columns (n_sub*16)
|
||||
// so first call for each n_sub should NOT accumulate
|
||||
if (tid == 128) umma_ss_f16(tb + n_sub * 16, dp, dv, idesc_pv, pv_kt > 0);
|
||||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
// Store row_sum to SMEM
|
||||
if (my_row_active) sRowSum[my_row] = my_row_sum;
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// EPILOGUE
|
||||
// EPILOGUE: TMEM → regs → normalize → BF16 → GMEM
|
||||
// ================================================================
|
||||
if (!full_tile) {
|
||||
// T <= 32: warp 0, 32x32b.x8
|
||||
if (wid == 0) {
|
||||
float rm = (lane < T) ? sRowMax[lane] : 0.0f;
|
||||
float rs = (lane < T) ? sRowSum[lane] : 1.0f;
|
||||
float o_vals[HD];
|
||||
for (int n = 0; n < HD / 8; n++) {
|
||||
if (my_warp_active) {
|
||||
if (my_row_active) {
|
||||
float rm = sRowMax[my_row];
|
||||
float rs = sRowSum[my_row];
|
||||
float inv_rs = 1.0f / rs;
|
||||
|
||||
// Read O from TMEM: HD values in HD/8 groups of 8
|
||||
// O is in TMEM at columns 0..N_NSUB*16-1
|
||||
// After PV GEMM, TMEM layout is:
|
||||
// N-sub-tile n writes to columns n*16..n*16+15
|
||||
// 32x32b.x8 reads 8 columns at a time
|
||||
// So we read (N_NSUB*16)/8 = N_NSUB*2 groups
|
||||
// But we need to read ALL N_NSUB*2 groups with correct row
|
||||
for (int n = 0; n < N_NSUB * 2; n++) {
|
||||
float tmp[8];
|
||||
uint32_t addr = tb + my_tmem_row_offset + n * 8;
|
||||
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
|
||||
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
|
||||
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
|
||||
: "r"(tb + n * 8));
|
||||
: "r"(addr));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;");
|
||||
if (lane < T) for (int c = 0; c < 8; c++) o_vals[n * 8 + c] = tmp[c];
|
||||
}
|
||||
if (lane < T) {
|
||||
float inv = 1.0f / rs;
|
||||
for (int d = 0; d < HD; d++) o_head[lane * HD + d] = f32_to_bf16(o_vals[d] * inv);
|
||||
if (lse_head) lse_head[lane] = logf(rs) + rm;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// T > 32: 4 warps, 16x256b.x1
|
||||
if (is_softmax_warp) {
|
||||
const int warp_row_start = wid * 32;
|
||||
const int warp_row_end = min((wid + 1) * 32, T);
|
||||
const int lane_row_base = lane * 4;
|
||||
const bool active = (lane_row_base >= warp_row_start) && (lane_row_base < warp_row_end);
|
||||
const int vrows = active ? min(4, warp_row_end - lane_row_base) : 0;
|
||||
|
||||
// Read O from TMEM one column at a time
|
||||
for (int d = 0; d < HD; d++) {
|
||||
float v[4];
|
||||
asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0,%1,%2,%3},[%4,%5];"
|
||||
: "=f"(v[0]), "=f"(v[1]), "=f"(v[2]), "=f"(v[3]) : "r"(tb), "r"(d));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;");
|
||||
if (active) {
|
||||
for (int r = 0; r < vrows; r++) {
|
||||
int gr = lane_row_base + r;
|
||||
if (gr < T) {
|
||||
float inv = 1.0f / sRowSum[gr];
|
||||
o_head[gr * HD + d] = f32_to_bf16(v[r] * inv);
|
||||
}
|
||||
for (int c = 0; c < 8; c++) {
|
||||
int d = n * 8 + c;
|
||||
if (d < HD) {
|
||||
o_head[my_row * HD + d] = f32_to_bf16(tmp[c] * inv_rs);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// LSE
|
||||
if (active) {
|
||||
for (int r = 0; r < vrows; r++) {
|
||||
int gr = lane_row_base + r;
|
||||
if (gr < T && lse_head) lse_head[gr] = logf(sRowSum[gr]) + sRowMax[gr];
|
||||
}
|
||||
if (lse_head) {
|
||||
lse_head[my_row] = logf(rs) + rm;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// TMEM dealloc
|
||||
if (is_mma_warp) tmem_dealloc(tb, TMEM_N);
|
||||
// TMEM dealloc (warp 4)
|
||||
if (is_mma_warp) {
|
||||
tmem_dealloc(tb, TMEM_N);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace dsv4::kernels::attention
|
||||
|
||||
@@ -4,8 +4,8 @@
|
||||
*
|
||||
* Tests:
|
||||
* 1. T=1 decode (regression — must match single-row results)
|
||||
* 2. T=2,4,8,16,32 (small prefill — only warp 0)
|
||||
* 3. T=64,128 (multi-warp prefill — all 4 softmax warps)
|
||||
* 2. T=2,4,8,16,32 (small prefill — warp 0 only)
|
||||
* 3. T=64,128 (multi-warp prefill — 4 softmax warps)
|
||||
* 4. LSE correctness for multi-row
|
||||
* 5. Multi-head + batched with T>1
|
||||
*/
|
||||
@@ -120,9 +120,9 @@ static int test_single_T(int T, int n_h = 1, int batch = 1) {
|
||||
params.lse_head_stride = T;
|
||||
params.lse_batch_stride = n_h * T;
|
||||
|
||||
// SMEM: tmemBase(4) + sRowMax(128*4) + sRowSum(128*4) + padding + sQ0 + sK0 + sPk + sV + s_p_vals
|
||||
// s_p_vals = 128 * 128 * 4 = 65536 bytes
|
||||
int smem = 4 + 32*4 + 32*4 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + MAX_T * SK * 4 + 256;
|
||||
// SMEM: sTmemBase(8) + sRowMax(512) + sRowSum(512) + padding(128) + sQ0(4096) + sK0(4096) + sPk(4096) + sV(512) + slack(256)
|
||||
// NO s_p_vals buffer — P is streamed per K-tile through sPk
|
||||
int smem = 8 + MAX_T*4 + MAX_T*4 + 128 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + 256;
|
||||
smem = (smem + 127) & ~127;
|
||||
|
||||
if (smem > 48 * 1024) {
|
||||
@@ -140,7 +140,7 @@ static int test_single_T(int T, int n_h = 1, int batch = 1) {
|
||||
}
|
||||
sync_err = cudaDeviceSynchronize();
|
||||
if (sync_err != cudaSuccess) {
|
||||
printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err));
|
||||
printf("CUDA ERROR: ~%s\n", cudaGetErrorString(sync_err));
|
||||
pass = 0; goto cleanup;
|
||||
}
|
||||
|
||||
@@ -173,6 +173,16 @@ static int test_single_T(int T, int n_h = 1, int batch = 1) {
|
||||
|
||||
if (cs < 0.999f) {
|
||||
printf(" FAIL batch=%d head=%d row=%d: cos=%.6f lse_err=%.6f\n", b, h, t, cs, lse_err);
|
||||
if (failed < 3) {
|
||||
// Print first few values for debugging
|
||||
printf(" kernel[0..7]=");
|
||||
for (int d2 = 0; d2 < min(8, HD); d2++)
|
||||
printf(" %.4f", bf16_to_f32_host(h_o[(idx * T + t) * HD + d2]));
|
||||
printf("\n ref [0..7]=");
|
||||
for (int d2 = 0; d2 < min(8, HD); d2++)
|
||||
printf(" %.4f", o_ref[t * HD + d2]);
|
||||
printf("\n");
|
||||
}
|
||||
failed++;
|
||||
}
|
||||
}
|
||||
@@ -198,7 +208,7 @@ int main() {
|
||||
// Test 1: T=1 decode (regression — must match single-row)
|
||||
all_pass &= test_single_T(1);
|
||||
|
||||
// Test 2: Small prefill (only warp 0)
|
||||
// Test 2: Small prefill (warp 0 only)
|
||||
all_pass &= test_single_T(2);
|
||||
all_pass &= test_single_T(4);
|
||||
all_pass &= test_single_T(8);
|
||||
|
||||
Reference in New Issue
Block a user