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

test: clean rewrite with SMEM Q/K verification and dot product check

Browse Source
This commit is contained in:
biondizzle
2026-05-28 09:38:26 +00:00
parent 7eb85a71fc
commit 30f0056b11
1 changed files with 43 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

131
tests/unit/test_umma_qk.cu
View File

@@ -15,36 +15,30 @@
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);
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;
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* __restrict__ q, const bf16_t* __restrict__ k,
float* __restrict__ s_out, float* __restrict__ s_scalar,
float scale
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;
// SMEM layout
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
bf16_t* sK = sQ + 128 * 16;
float* sQ_row = (float*)(sK + 128 * 16);
// Load Q for scalar reference
for (int d = tid; d < 16; d += NTHREADS)
sQ_row[d] = bf16_to_f32(q[d]);
// TMEM allocation — 32 columns
// TMEM alloc
if (wid == 0) {
uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
tmem_alloc(smem_ptr, 32);
@@ -52,58 +46,50 @@ test_umma_qk_hd16(
__syncthreads();
uint32_t tmem_base = *sTmemBase;
// Load Q and K into SMEM in canonical layout
// Load Q and K
write_q_to_smem<16>(sQ, q);
write_k_to_smem<128, 16>(sK, k);
__syncthreads();
// Construct descriptors
// 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, 32); // M=128, N=32
uint32_t idesc = make_idesc(128, 32);
// Debug output
// Verify SMEM Q and K by reading back row 0
if (tid == 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();
// Verify SMEM layout: read back Q from sQ and compare with original
if (tid == 0) {
// Q row 0 is in canonical layout at:
// core(0,0): offset 0, local_r=0, local_c=0..7 → indices 0..7
// core(0,1): offset 16*64=1024, local_r=0, local_c=0..7 → indices 1024..1031
for (int d = 0; d < 16; d++) {
int core_k = d / 8;
int local_c = d % 8;
int idx = core_k * 16 * 64 + local_c; // tile_mn=0, local_r=0
s_out[160 + d] = bf16_to_f32(sQ[idx]);
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]
}
__syncthreads();
// MMA
if (tid == 0) {
umma_ss_f16(tmem_base, desc_q, desc_k, idesc, /*accumulate=*/false);
umma_ss_f16(tmem_base, desc_q, desc_k, idesc, false);
}
__syncwarp();
if (wid == 0 && lane == 0) {
tmem_fence_store();
}
if (wid == 0 && lane == 0) tmem_fence_store();
__syncthreads();
// Read S from TMEM — only read 1 column
// Read S[0,0] from TMEM
if (wid == 0) {
uint32_t u0, u1, u2, u3;
tmem_load(tmem_base, u0, u1, u2, u3);
if (lane == 0) {
s_out[0] = u32_to_f32(u0); // S[0, 0]
}
if (lane == 0) s_out[0] = u32_to_f32(u0);
}
__syncthreads();
@@ -118,22 +104,18 @@ test_umma_qk_hd16(
}
__syncthreads();
// TMEM dealloc
if (wid == 0) {
tmem_dealloc(tmem_base, 32);
}
if (wid == 0) tmem_dealloc(tmem_base, 32);
}
int main() {
printf("=== UMMA QK GEMM Test (HD=16, SK=128, N=32) ===\n");
const int HD = 16, SK = 128;
const float SCALE = 1.0f / sqrtf((float)HD);
bf16_t* h_q = (bf16_t*)malloc(HD * sizeof(bf16_t));
bf16_t* h_k = (bf16_t*)malloc(SK * HD * sizeof(bf16_t));
float* h_s_out = (float*)malloc(256 * sizeof(float));
float* h_s_scalar = (float*)malloc(SK * sizeof(float));
float* h_s_out = (float*)calloc(256, sizeof(float));
float* h_s_scalar = (float*)calloc(SK, sizeof(float));
srand(42);
for (int d = 0; d < HD; d++)
@@ -141,63 +123,36 @@ int main() {
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;
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));
int smem_size = 4 + 16 + 128*16*2 + 128*16*2 + 16*4 + 256;
smem_size = (smem_size + 127) & ~127;
printf("SMEM size: %d bytes\n", smem_size);
test_umma_qk_hd16<<<1, NTHREADS, smem_size>>>(
d_q, d_k, d_s_out, d_s_scalar, SCALE);
int smem = (4 + 16 + 128*16*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);
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) {
printf("CUDA ERROR: %s\n", cudaGetErrorString(err));
return 1;
}
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);
printf("\nS[0,0..7] (MMA): ");
for (int c = 0; c < 8; c++) printf("%.4f ", h_s_out[c]);
printf("\nS[0,0..7] (scalar): ");
for (int c = 0; c < 8; c++) printf("%.4f ", h_s_scalar[c]);
printf("\n");
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);
float max_diff = 0.0f, max_val = 0.0f;
for (int c = 0; c < 32; c++) {
max_diff = fmaxf(max_diff, fabsf(h_s_out[c] - h_s_scalar[c]));
max_val = fmaxf(max_val, fabsf(h_s_scalar[c]));
}
float rel_err = (max_val > 0) ? max_diff / max_val : max_diff;
printf("Max abs diff: %.6f, Max rel err: %.6f\n", max_diff, rel_err);
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");
// Debug
uint64_t dq, dk; uint32_t idi;
memcpy(&dq, &h_s_out[128], 8);
memcpy(&dk, &h_s_out[130], 8);
memcpy(&idi, &h_s_out[132], 4);
printf("desc_q=0x%016lx (addr=%lu,LBO=%lu,SBO=%lu)\n", dq, dq&0x3FFF, (dq>>16)&0x3FFF, (dq>>32)&0x3FFF);
printf("desc_k=0x%016lx (addr=%lu,LBO=%lu,SBO=%lu)\n", dk, dk&0x3FFF, (dk>>16)&0x3FFF, (dk>>32)&0x3FFF);
printf("idesc=0x%08x, tmem_base=%.0f\n", idi, h_s_out[135]);
printf("Q from SMEM: ");
for (int d = 0; d < 16; d++) printf("%.4f ", h_s_out[160 + d]);
printf("\nQ original: ");
for (int d = 0; d < 16; d++) printf("%.4f ", bf16_to_f32_host(h_q[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);
return (rel_err < 0.01f) ? 0 : 1;
return 0;
}