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

test: clean UMMA QK test, debug 4x factor, 8KB padding, 128 TMEM cols

Browse Source
This commit is contained in:
biondizzle
2026-05-28 10:01:39 +00:00
parent ac65ece33b
commit 4f7c9649fd
1 changed files with 61 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

124
tests/unit/test_umma_qk.cu
View File

@@ -1,6 +1,8 @@
/**
* Standalone CUDA test for UMMA QK GEMM (tcgen05.mma SS, BF16).
* Test: HD=16, SK=128, single K-tile, MMA with N=32 (32 TMEM columns)
* UMMA QK GEMM Test (HD=16, SK=128) — debugging 4× output factor.
* MMA produces S[0,0] = 4× scalar reference. Investigating why.
*
* Working: 1 tmem_load at column 0. Multi-column load crashes.
*/
#include <cuda_runtime.h>
@@ -14,67 +16,70 @@
using namespace dsv4::kernels::attention;
static bf16_t f32_to_bf16_host(float f) {
uint32_t u; memcpy(&u, &f, 4); return (uint16_t)(u >> 16);
}
static float bf16_to_f32_host(bf16_t h) {
uint32_t u = (uint32_t)h << 16; float f; memcpy(&f, &u, 4); return f;
}
static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; }
__global__ void __launch_bounds__(NTHREADS)
test_umma_qk_hd16(
const bf16_t* q, const bf16_t* k,
float* s_out, float* s_scalar, float scale
) {
test_umma_qk_hd16(const bf16_t* q, const bf16_t* k,
float* s_out, float* s_scalar, float scale)
{
const int tid = threadIdx.x;
const int wid = tid / WARP, lane = tid % WARP;
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
// Add 8KB padding after sQ to prevent MMA from reading into sK
bf16_t* sK = sQ + 128 * 16 + 4096; // 4096 BF16 = 8KB padding
bf16_t* sK = sQ + 128 * 16 + 4096; // 8KB padding after Q
float* sQ_row = (float*)(sK + 128 * 16);
for (int d = tid; d < 16; d += NTHREADS)
sQ_row[d] = bf16_to_f32(q[d]);
for (int d = tid; d < 16; d += NTHREADS) sQ_row[d] = bf16_to_f32(q[d]);
// TMEM alloc
// TMEM alloc (128 cols for MMA output N=128)
if (wid == 0) {
uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
tmem_alloc(smem_ptr, 128); // 128 columns for N=128
uint32_t sp = __cvta_generic_to_shared(sTmemBase);
tmem_alloc(sp, 128);
}
__syncthreads();
uint32_t tmem_base = *sTmemBase;
// Load Q and K
// Load Q and K into SMEM
write_q_to_smem<16>(sQ, q);
write_k_to_smem<128, 16>(sK, k);
__syncthreads();
// Verify SMEM data
if (tid == 0) {
for (int d = 0; d < 16; d++) {
int ck = d / 8, lc = d % 8;
s_out[160 + d] = bf16_to_f32(sQ[ck * 16 * 64 + lc]);
s_out[176 + d] = bf16_to_f32(sK[ck * 16 * 64 + lc]);
}
float dot = 0.0f;
for (int d = 0; d < 16; d++) {
int ck = d / 8, lc = d % 8;
dot += bf16_to_f32(sQ[ck * 16 * 64 + lc]) * bf16_to_f32(sK[ck * 16 * 64 + lc]);
}
s_out[192] = dot * scale;
}
__syncthreads();
// Descriptors
uint32_t sQ_smem = __cvta_generic_to_shared(sQ);
uint32_t sK_smem = __cvta_generic_to_shared(sK);
uint64_t desc_q = make_umma_desc_kmajor_none(sQ_smem, 128);
uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 128);
uint32_t idesc = make_idesc(128, 128); // M=128, N=128
// Verify SMEM Q and K by reading back row 0
// Try different idesc values to understand the 4× factor
// Test 1: M=128, N=128
uint32_t idesc = make_idesc(128, 128);
if (tid == 0) {
for (int d = 0; d < 16; d++) {
int core_k = d / 8, local_c = d % 8;
s_out[160 + d] = bf16_to_f32(sQ[core_k * 16 * 64 + local_c]);
s_out[176 + d] = bf16_to_f32(sK[core_k * 16 * 64 + local_c]);
}
// Dot product from SMEM
float dot = 0.0f;
for (int d = 0; d < 16; d++) {
int core_k = d / 8, local_c = d % 8;
float qv = bf16_to_f32(sQ[core_k * 16 * 64 + local_c]);
float kv = bf16_to_f32(sK[core_k * 16 * 64 + local_c]);
dot += qv * kv;
}
s_out[192] = dot * scale; // Expected S[0,0]
memcpy(&s_out[128], &desc_q, 8);
memcpy(&s_out[130], &desc_k, 8);
memcpy(&s_out[132], &idesc, 4);
s_out[133] = (float)sQ_smem;
s_out[134] = (float)sK_smem;
s_out[135] = (float)tmem_base;
}
__syncthreads();
@@ -86,12 +91,10 @@ test_umma_qk_hd16(
if (wid == 0 && lane == 0) tmem_fence_store();
__syncthreads();
// Read S from TMEM — read 128 columns one at a time
// Use separate tmem_load calls with __syncthreads between them
// to avoid the loop crash
// Read S[0,0] from TMEM column 0
if (wid == 0) {
uint32_t u0, u1, u2, u3;
tmem_load(tmem_base + 1, u0, u1, u2, u3);
tmem_load(tmem_base, u0, u1, u2, u3);
if (lane == 0) s_out[0] = u32_to_f32(u0);
}
__syncthreads();
@@ -111,7 +114,7 @@ test_umma_qk_hd16(
}
int main() {
printf("=== UMMA QK GEMM Test (HD=16, SK=128, N=32) ===\n");
printf("=== UMMA QK GEMM Test (HD=16) ===\n");
const int HD = 16, SK = 128;
const float SCALE = 1.0f / sqrtf((float)HD);
@@ -121,20 +124,15 @@ int main() {
float* h_s_scalar = (float*)calloc(SK, sizeof(float));
srand(42);
for (int d = 0; d < HD; d++)
h_q[d] = f32_to_bf16_host((float)(rand() % 100) / 100.0f - 0.5f);
for (int i = 0; i < SK * HD; i++)
h_k[i] = f32_to_bf16_host((float)(rand() % 100) / 100.0f - 0.5f);
for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
for (int i = 0; i < SK*HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
bf16_t *d_q, *d_k; float *d_s_out, *d_s_scalar;
cudaMalloc(&d_q, HD * sizeof(bf16_t));
cudaMalloc(&d_k, SK * HD * sizeof(bf16_t));
cudaMalloc(&d_s_out, 256 * sizeof(float));
cudaMalloc(&d_s_scalar, SK * sizeof(float));
cudaMemcpy(d_q, h_q, HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_k, h_k, SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemset(d_s_out, 0, 256 * sizeof(float));
cudaMemset(d_s_scalar, 0, SK * sizeof(float));
cudaMalloc(&d_q, HD*sizeof(bf16_t)); cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));
cudaMalloc(&d_s_out, 256*sizeof(float)); cudaMalloc(&d_s_scalar, SK*sizeof(float));
cudaMemcpy(d_q, h_q, HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemset(d_s_out, 0, 256*sizeof(float)); cudaMemset(d_s_scalar, 0, SK*sizeof(float));
int smem = (4 + 16 + 128*16*2 + 4096*2 + 128*16*2 + 16*4 + 256 + 127) & ~127;
test_umma_qk_hd16<<<1, NTHREADS, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE);
@@ -142,18 +140,18 @@ int main() {
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }
cudaMemcpy(h_s_out, d_s_out, 256 * sizeof(float), cudaMemcpyDeviceToHost);
cudaMemcpy(h_s_scalar, d_s_scalar, SK * sizeof(float), cudaMemcpyDeviceToHost);
cudaMemcpy(h_s_out, d_s_out, 256*sizeof(float), cudaMemcpyDeviceToHost);
cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost);
printf("S[0,0] from MMA: %.6f\n", h_s_out[0]);
printf("S[0,0] from SMEM: %.6f\n", h_s_out[192]);
printf("S[0,0] scalar ref: %.6f\n", h_s_scalar[0]);
printf("Ratio MMA/scalar: %.4f\n", h_s_scalar[0] != 0 ? h_s_out[0]/h_s_scalar[0] : 0);
printf("S[0,0] MMA: %.6f\n", h_s_out[0]);
printf("S[0,0] SMEM: %.6f\n", h_s_out[192]);
printf("S[0,0] scalar: %.6f\n", h_s_scalar[0]);
printf("Ratio: %.4f\n", h_s_scalar[0] != 0 ? h_s_out[0]/h_s_scalar[0] : 0);
printf("\nQ from SMEM: "); for (int d = 0; d < 8; d++) printf("%.4f ", h_s_out[160+d]); printf("...\n");
printf("K[0] from SMEM: "); for (int d = 0; d < 8; d++) printf("%.4f ", h_s_out[176+d]); printf("...\n");
printf("Q original: "); for (int d = 0; d < 8; d++) printf("%.4f ", bf16_to_f32_host(h_q[d])); printf("...\n");
printf("K[0] original: "); for (int d = 0; d < 8; d++) printf("%.4f ", bf16_to_f32_host(h_k[d])); printf("...\n");
printf("Q SMEM: "); for(int d=0;d<4;d++) printf("%.4f ",h_s_out[160+d]); printf("...\n");
printf("Q orig: "); for(int d=0;d<4;d++) printf("%.4f ",bf16_to_f32_host(h_q[d])); printf("...\n");
printf("K[0] SMEM: "); for(int d=0;d<4;d++) printf("%.4f ",h_s_out[176+d]); printf("...\n");
printf("K[0] orig: "); for(int d=0;d<4;d++) printf("%.4f ",bf16_to_f32_host(h_k[d])); printf("...\n");
cudaFree(d_q); cudaFree(d_k); cudaFree(d_s_out); cudaFree(d_s_scalar);
free(h_q); free(h_k); free(h_s_out); free(h_s_scalar);