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

test: HD=64 QK with contiguous SMEM + offset descriptors

Browse Source
This commit is contained in:
biondizzle
2026-05-28 12:50:07 +00:00
parent 1c01e8e412
commit e8ac2120ad
1 changed files with 52 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

129
tests/unit/test_umma_qk_hd64.cu
View File

@@ -1,12 +1,16 @@
/**
* UMMA QK GEMM Test — HD=64 (4 K-tiles)
*
* Fixed from previous version:
* 1. Zero TMEM before accumulate loop (tcgen05.alloc does NOT zero)
* 2. tcgen05.fence::after_thread_sync after MMA loop (correct MMA→TMEM fence)
* 3. Single-thread MMA call (tid==0, matching gau-nernst's elect_one pattern)
* 4. Separate SMEM per K-tile (avoids offset descriptor issues)
* 5. Full 128×128 output validation
* Using gau-nernst's approach: single contiguous SMEM, offset descriptors.
* Q is (128, HD) and K is (128, HD) in canonical K-major layout.
* For each K-tile kt, the descriptor start address is:
* base + kt * BLOCK_MN * 32 (each 16-BF16 K-tile is BLOCK_MN * 32 bytes apart)
*
* Fixes from previous version:
* 1. Zero TMEM before accumulate (tcgen05.alloc does NOT zero)
* 2. tcgen05.fence::after_thread_sync after each MMA + before TMEM read
* 3. Single-thread MMA call (tid==0)
* 4. Contiguous SMEM with offset descriptors
*/
#include <cuda_runtime.h>
@@ -25,7 +29,7 @@ static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; m
constexpr int HD = 64;
constexpr int SK = 128;
constexpr int NKT = HD / MMA_K_BF16; // 4 K-tiles
constexpr int NKT = HD / MMA_K_BF16; // 4
constexpr int BLOCK_MN = 128;
__global__ void __launch_bounds__(128)
@@ -36,29 +40,29 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
const int wid = tid / WARP, lane = tid % WARP;
// SMEM layout:
// [0..3] tmem_base (written by tcgen05.alloc)
// [4..) K-tile SMEM regions (alternating Q, K per K-tile)
// Each (128, 16) = 128*16*2 = 4096 bytes in canonical layout
// [0..3] tmem_base
// [64..) sQ (128, HD) canonical = 128*64*2 = 16384 bytes
// sK (128, HD) canonical = 16384 bytes
// Align to 128 bytes for SMEM access patterns
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
// Align SMEM data to 16 bytes for descriptor start_address
char* data_base = (char*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
// Separate SMEM for each Q K-tile and K K-tile
bf16_t* sQ[NKT];
bf16_t* sK[NKT];
for (int kt = 0; kt < NKT; kt++) {
sQ[kt] = (bf16_t*)(data_base + (kt * 2 + 0) * BLOCK_MN * MMA_K_BF16 * sizeof(bf16_t));
sK[kt] = (bf16_t*)(data_base + (kt * 2 + 1) * BLOCK_MN * MMA_K_BF16 * sizeof(bf16_t));
}
bf16_t* sQ = (bf16_t*)(sbuf + 64); // 64-byte aligned
bf16_t* sK = sQ + BLOCK_MN * HD; // follows Q
// Step 1: TMEM alloc — 128 columns (Layout D for M=128, N=128)
// Load Q (1, HD) into (128, HD) canonical layout — only row 0 has data
write_q_to_smem<HD>(sQ, q);
// Load K (SK, HD) into (128, HD) canonical layout
write_k_to_smem<SK, HD>(sK, k);
__syncthreads();
// TMEM alloc — 128 columns for (128, 128) Layout D output
if (wid == 0) {
tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128);
}
__syncthreads();
uint32_t tb = *sTmemBase;
// Step 2: Zero TMEM (tcgen05.alloc does NOT zero)
// Zero TMEM (tcgen05.alloc does NOT zero)
if (wid == 0) {
for (int col = 0; col < 128; col++) {
tmem_store(tb + col, 0, 0, 0, 0);
@@ -67,53 +71,32 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
}
__syncthreads();
// Step 3: Load Q and K into per-K-tile SMEM regions
// Global Q is (1, HD). Global K is (SK, HD).
// K-tile kt: columns [16*kt, 16*kt+16) of the (128, HD) padded matrix.
for (int kt = 0; kt < NKT; kt++) {
bf16_t* sq = sQ[kt];
bf16_t* sk = sK[kt];
// Zero the (128, 16) tile
for (int i = tid; i < BLOCK_MN * MMA_K_BF16; i += 128) {
sq[i] = 0;
sk[i] = 0;
}
__syncthreads();
// Q: only row 0 has data. Write to canonical layout.
// For (128, 16): CORES_MN=16, CORES_K=2
// Row 0: core_mn=0, local_r=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];
}
// K: write (SK, 16) into (128, 16) canonical layout
for (int i = tid; i < SK * MMA_K_BF16; i += 128) {
int r = i / MMA_K_BF16;
int c = i % MMA_K_BF16;
int tmn = r / 8, ck = c / 8, lr = r % 8, lc = c % 8;
sk[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * MMA_K_BF16 + c];
}
}
__syncthreads();
// Step 4: Multi-K-tile QK GEMM: Q × K^T → S in TMEM
// Each MMA call: (128, 16) × (16, 128) → (128, 128) partial, accumulated in TMEM
// Multi-K-tile QK GEMM
// For K-major NONE layout, each 16-column K-tile starts at:
// base_smem + kt * BLOCK_MN * 32 (in bytes)
// Each K-tile spans 2 core-matrix columns (16 BF16).
// LBO = BLOCK_MN * 16 (bytes) = stride between the 2 core columns.
// The descriptor describes a (BLOCK_MN, 16) sub-matrix.
uint32_t sQ_smem = __cvta_generic_to_shared(sQ);
uint32_t sK_smem = __cvta_generic_to_shared(sK);
uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN);
for (int kt = 0; kt < NKT; kt++) {
// Construct descriptor for this K-tile
uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ[kt]), BLOCK_MN);
uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK[kt]), BLOCK_MN);
// Offset to the kt-th K-tile: kt * BLOCK_MN * 32 bytes
// (each K-tile is 16 BF16 = 2 core columns, BLOCK_MN * 16B apart, total 32B per row)
// Wait, that's not right. In canonical layout, the kt-th 16-column slice starts
// at the (2*kt)-th core column. Each core column is BLOCK_MN/8 * 128 bytes
// = BLOCK_MN * 16 bytes. So offset = 2*kt * BLOCK_MN * 16 bytes = kt * BLOCK_MN * 32 bytes.
uint32_t q_addr = sQ_smem + kt * BLOCK_MN * 32;
uint32_t k_addr = sK_smem + kt * BLOCK_MN * 32;
uint64_t dq = make_umma_desc_kmajor_none(q_addr, BLOCK_MN);
uint64_t dk = make_umma_desc_kmajor_none(k_addr, BLOCK_MN);
// Single-thread MMA call (matching gau-nernst's elect_one pattern)
// accumulate=true for all tiles (TMEM zeroed above)
if (tid == 0) {
umma_ss_f16(tb, dq, dk, idesc, true);
umma_ss_f16(tb, dq, dk, idesc, true); // accumulate across K-tiles
}
// Ensure MMA completes before next K-tile
// Fence after each K-tile MMA to ensure TMEM is updated
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
}
@@ -122,9 +105,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
// Step 5: Read S from TMEM (Layout D: 32x32b.x8)
// 128 TMEM columns, 128 rows. Each warp reads 32 rows × 8 columns.
// 4 warps × 16 iterations = 128 rows × 128 columns
// Read S from TMEM (Layout D: 32x32b.x8)
for (int n = 0; n < 128 / 8; n++) {
const int row = wid * 32;
const int col = n * 8;
@@ -145,7 +126,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
}
__syncthreads();
// Step 6: Scalar reference — S[0, j] = sum(Q[0,d]*K[j,d], d=0..63) * scale
// Scalar reference
if (tid == 0) {
for (int j = 0; j < SK; j++) {
float dot = 0.0f;
@@ -155,12 +136,11 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
}
}
// TMEM dealloc
if (wid == 0) tmem_dealloc(tb, 128);
}
int main() {
printf("=== UMMA QK GEMM HD=64 (4 K-tiles, fixed) ===\n");
printf("=== UMMA QK GEMM HD=64 (4 K-tiles, contiguous SMEM) ===\n");
const float SCALE = 1.0f / sqrtf((float)HD);
bf16_t* h_q = (bf16_t*)malloc(HD * sizeof(bf16_t));
@@ -178,8 +158,8 @@ int main() {
cudaMemcpy(d_q, h_q, HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
// SMEM: tmem_base(4) + alignment(16) + 8 K-tiles × 128×16×2
int smem = (4 + 16 + NKT * 2 * BLOCK_MN * MMA_K_BF16 * sizeof(bf16_t) + 256 + 127) & ~127;
// SMEM: 64 (tmem_base+pad) + 2 * 128*64*2 (Q+K) + 256 (extra)
int smem = (64 + 2 * BLOCK_MN * HD * sizeof(bf16_t) + 256 + 127) & ~127;
printf("SMEM: %d bytes (%d KB)\n", smem, smem / 1024);
test_umma_hd64<<<1, 128, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE);
@@ -190,13 +170,11 @@ int main() {
cudaMemcpy(h_s_out, d_s_out, SK*SK*sizeof(float), cudaMemcpyDeviceToHost);
cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost);
// Print row 0 (the only row with data since Q is 1×HD)
printf("S[0,0..7] MMA: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n");
printf("S[0,0..7] ref: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_scalar[c]); printf("\n");
printf("S[0,120..127] MMA: "); for(int c=120;c<128;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n");
printf("S[0,120..127] ref: "); for(int c=120;c<128;c++) printf("%.6f ",h_s_scalar[c]); printf("\n");
// Validate row 0 (full 128 columns)
float max_diff = 0.0f, max_val = 0.0f;
for (int c = 0; c < SK; c++) {
float diff = fabsf(h_s_out[0*SK+c] - h_s_scalar[c]);
@@ -206,13 +184,10 @@ int main() {
float rel_err = max_val > 0 ? max_diff / max_val : max_diff;
printf("Row 0 rel err (128 cols): %.8f\n", rel_err);
// Check rows 1-127 should be zero (Q has no data there)
float max_nonzero = 0.0f;
for (int r = 1; r < SK; r++) {
for (int c = 0; c < SK; c++) {
for (int r = 1; r < SK; r++)
for (int c = 0; c < SK; c++)
max_nonzero = fmaxf(max_nonzero, fabsf(h_s_out[r*SK+c]));
}
}
printf("Rows 1-127 max abs: %.8f\n", max_nonzero);
bool row0_ok = rel_err < 0.001f;