From 4f7c9649fd919e451091fef415acf7ac734514bd Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 10:01:39 +0000 Subject: [PATCH] test: clean UMMA QK test, debug 4x factor, 8KB padding, 128 TMEM cols --- tests/unit/test_umma_qk.cu | 124 ++++++++++++++++++------------------- 1 file changed, 61 insertions(+), 63 deletions(-) diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index eb7aecbe..2dd9c799 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -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 @@ -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);