From 2fd64c464d951d2a2a336130367fc93b751ee47c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 14:54:10 +0000 Subject: [PATCH] FMHA HD=64 with BLOCK_MN_B=64 for V, proper output dimensions --- tests/unit/test_fmha_hd64_smem_p.cu | 54 +++++++++++++++++------------ 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/tests/unit/test_fmha_hd64_smem_p.cu b/tests/unit/test_fmha_hd64_smem_p.cu index dd09dfc2..2ab0773e 100644 --- a/tests/unit/test_fmha_hd64_smem_p.cu +++ b/tests/unit/test_fmha_hd64_smem_p.cu @@ -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 @@ -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);