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

FMHA HD=64 with BLOCK_MN_B=64 for V, proper output dimensions

Browse Source
This commit is contained in:
biondizzle
2026-05-28 14:54:10 +00:00
parent 15ecc1f616
commit 2fd64c464d
1 changed files with 31 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
tests/unit/test_fmha_hd64_smem_p.cu
View File

@@ -1,8 +1,8 @@
/**
* Full FMHA HD=64, SK=128 — PV via SS MMA (SMEM-P approach)
*
* Extends test_fmha_v5 (HD=16) to HD=64.
* 4 QK K-tiles, 8 PV K-tiles, per-K-tile P fill.
* Key difference from HD=16: V K-tiles are (64, 16) with BLOCK_MN=64,
* producing 64 output values per PV MMA call. The idesc uses MMA_N=4.
*/
#include <cuda_runtime.h>
@@ -22,8 +22,11 @@ static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; m
constexpr int HD = 64, SK = 128, BLOCK_MN = 128;
constexpr int NKT_QK = HD / MMA_K_BF16; // 4
constexpr int NKT_PV = SK / MMA_K_BF16; // 8
constexpr int TMEM_N = 128;
constexpr int TILE_SZ = BLOCK_MN * MMA_K_BF16;
constexpr int TILE_SZ = BLOCK_MN * MMA_K_BF16; // 2048 BF16
// V K-tile size: (HD, 16) canonical with BLOCK_MN=HD
// CORES_MN = HD/8 = 8, CORES_K = 2
// 8 * 2 * 64 = 1024 BF16 per V K-tile
constexpr int V_TILE_SZ = (HD / 8) * 2 * 64; // 1024 BF16
__global__ void __launch_bounds__(128)
test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
@@ -34,22 +37,22 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
bf16_t* sK0 = sQ0 + NKT_QK * TILE_SZ; // 4 Q K-tiles
bf16_t* sK0 = sQ0 + NKT_QK * TILE_SZ;
bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + NKT_QK * 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 + NKT_PV * 256);
float* s_p_vals = (float*)(sV + NKT_PV * V_TILE_SZ);
// Load Q K-tiles
// Load Q K-tiles (4 × (128,16) canonical)
for (int kt = 0; kt < NKT_QK; kt++) {
bf16_t* sq = sQ0 + kt * TILE_SZ;
for (int i = tid; i < TILE_SZ; i += 128) sq[i] = 0;
for (int d = tid; d < MMA_K_BF16; d += 128) {
int ck = d / 8, lc = d % 8;
sq[ck * 16 * 64 + lc] = q[kt * MMA_K_BF16 + d]; // Row 0 only
sq[ck * 16 * 64 + lc] = q[kt * MMA_K_BF16 + d];
}
}
// Load K K-tiles
// Load K K-tiles (4 × (128,16) canonical)
for (int kt = 0; kt < NKT_QK; kt++) {
bf16_t* sk = sK0 + kt * TILE_SZ;
for (int i = tid; i < TILE_SZ; i += 128) sk[i] = 0;
@@ -62,27 +65,30 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
}
}
// Load V K-tiles
// Load V K-tiles (8 × (64,16) canonical with BLOCK_MN=64)
// B[d, r] in canonical: g_mn=d/8, g_k=r/8, llr=d%8, lc=r%8
// (64,16): CORES_MN=8, CORES_K=2
// core(g_mn, g_k) at g_k * 8 * 64 + g_mn * 64 + llr * 8 + lc
for (int kt = 0; kt < NKT_PV; kt++) {
bf16_t* sv = sV + kt * 256;
for (int i = tid; i < 256; i += 128) sv[i] = 0;
bf16_t* sv = sV + kt * V_TILE_SZ;
for (int i = tid; i < V_TILE_SZ; i += 128) sv[i] = 0;
for (int d = tid; d < HD; d += 128) {
for (int lr = 0; lr < MMA_K_BF16; lr++) {
int r = kt * MMA_K_BF16 + lr;
int g_mn = d / 8, g_k = lr / 8;
int llr = d % 8, lc = lr % 8;
sv[g_k * 2 * 64 + g_mn * 64 + llr * 8 + lc] = v[d * SK + r];
sv[g_k * 8 * 64 + g_mn * 64 + llr * 8 + lc] = v[d * SK + r];
}
}
}
__syncthreads();
// TMEM alloc
if (wid == 1) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N);
// TMEM alloc: 128 columns (S is 128×128, O is 128×64 → first 64 cols)
if (wid == 1) tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128);
__syncthreads();
uint32_t tb = *sTmemBase;
// ===== QK GEMM (4 K-tiles, accumulate) =====
// ===== QK GEMM (4 K-tiles) =====
{
uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN);
for (int kt = 0; kt < NKT_QK; kt++) {
@@ -107,7 +113,7 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
: "r"(tb + n*8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) {
s_vals[n*8+c] = tmp[c] * scale * 2.0f; // ×2 to undo QK MMA 0.5 scale
s_vals[n*8+c] = tmp[c] * scale * 2.0f; // ×2 to undo QK 0.5 scale
row_max = fmaxf(row_max, tmp[c] * scale * 2.0f);
}
}
@@ -124,8 +130,9 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
__syncthreads();
// ===== PV GEMM (8 K-tiles, per-K-tile P fill) =====
// B = V(64, 16) with BLOCK_MN=64. idesc: MMA_M=8, MMA_N=4.
{
uint32_t idesc_pv = make_idesc(BLOCK_MN, HD);
uint32_t idesc_pv = make_idesc(BLOCK_MN, HD); // (128, 64) → MMA_M=8, MMA_N=4
for (int kt = 0; kt < NKT_PV; kt++) {
// Fill sPk from s_p_vals
for (int i = tid; i < TILE_SZ; i += 128) sPk[i] = 0;
@@ -136,9 +143,9 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
}
__syncthreads();
bf16_t* sv = sV + kt * 256;
bf16_t* sv = sV + kt * V_TILE_SZ;
uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), BLOCK_MN);
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sv), 16);
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sv), HD);
if (tid == 0) umma_ss_f16(tb, dp, dv, idesc_pv, kt > 0);
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
@@ -181,11 +188,11 @@ test_fmha_hd64_smem_p(const bf16_t* q, const bf16_t* k, const bf16_t* v,
}
}
if (wid == 0) tmem_dealloc(tb, TMEM_N);
if (wid == 0) tmem_dealloc(tb, 128);
}
int main() {
printf("=== Full FMHA HD=64 SMEM-P (PV via SS MMA) ===\n");
printf("=== Full FMHA HD=64 SMEM-P (PV via SS MMA, BLOCK_MN_B=64) ===\n");
const float SCALE = 1.0f / sqrtf((float)HD);
bf16_t* h_q = (bf16_t*)malloc(HD*sizeof(bf16_t));
@@ -209,7 +216,8 @@ int main() {
cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_v, h_v, HD*SK*sizeof(bf16_t), cudaMemcpyHostToDevice);
int smem = (4+16 + NKT_QK*TILE_SZ*2 + NKT_QK*TILE_SZ*2 + TILE_SZ*2 + NKT_PV*256*2 + SK*4 + 256 + 127) & ~127;
// SMEM: tmem(4+12) + sQ(4*4096) + sK(4*4096) + sPk(4096) + sV(8*2048) + s_p_vals(512) + align
int smem = (4+16 + NKT_QK*TILE_SZ*2 + NKT_QK*TILE_SZ*2 + TILE_SZ*2 + NKT_PV*V_TILE_SZ*2 + SK*4 + 256 + 127) & ~127;
printf("SMEM: %d bytes (%.1f KB, limit 232 KB)\n", smem, smem/1024.0f);
test_fmha_hd64_smem_p<<<1, 128, smem>>>(d_q, d_k, d_v, d_o, d_o_scalar, SCALE);