From 41128122e3dc7874b328d3d49cccb852bbe112be Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 09:33:45 +0000 Subject: [PATCH] test: clean rewrite, 32 TMEM cols, MMA N=32, tmem_load loop --- tests/unit/test_umma_qk.cu | 135 +++++++------------------------------ 1 file changed, 24 insertions(+), 111 deletions(-) diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index 5af429da..b7dd55bb 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -1,11 +1,6 @@ /** * Standalone CUDA test for UMMA QK GEMM (tcgen05.mma SS, BF16). - * - * Tests that tcgen05.mma produces correct QK attention scores. - * Uses K-major, SWIZZLE_NONE, BF16 descriptors with canonical SMEM layout. - * Based on gau-nernst's tcgen05 tutorial. - * - * Test: HD=16, SK=128 (single K-tile, single MMA call) + * Test: HD=16, SK=128, single K-tile, MMA with N=32 (32 TMEM columns) */ #include @@ -19,128 +14,63 @@ using namespace dsv4::kernels::attention; -// Host-side BF16 conversion static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u, &f, 4); return (uint16_t)(u >> 16); } -// ================================================================== -// Test kernel: UMMA QK GEMM for HD=16, SK=128 -// ================================================================== -// A = Q: (128, 16) padded, K-major core-matrix layout -// B = K: (128, 16) in K-major core-matrix layout -// MMA computes S = A × B^T = Q × K^T → (128, 128) in TMEM - __global__ void __launch_bounds__(NTHREADS) test_umma_qk_hd16( const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, - float* __restrict__ s_out, // output: S[0, 0..127] from TMEM - float* __restrict__ s_scalar, // scalar reference + float* __restrict__ s_out, float* __restrict__ s_scalar, float scale ) { const int tid = threadIdx.x; const int wid = tid / WARP, lane = tid % WARP; - // ================================================================ // SMEM layout - // ================================================================ - // sQ: (128, 16) in canonical layout = 4096 bytes - // sK: (128, 16) in canonical layout = 4096 bytes - // sTmemBase: 4 bytes - // sQ_row: 16 floats for scalar ref extern __shared__ char sbuf[]; uint32_t* sTmemBase = (uint32_t*)sbuf; - // Align to 16 bytes for UMMA bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15); - bf16_t* sK = sQ + 128 * 16; // 4096 bytes after Q + bf16_t* sK = sQ + 128 * 16; float* sQ_row = (float*)(sK + 128 * 16); - // Load Q to sQ_row for scalar reference - for (int d = tid; d < 16; d += NTHREADS) { + // Load Q for scalar reference + for (int d = tid; d < 16; d += NTHREADS) sQ_row[d] = bf16_to_f32(q[d]); - } - // ================================================================ - // TMEM allocation - // ================================================================ + // TMEM allocation — 32 columns if (wid == 0) { uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); - tmem_alloc(smem_ptr, 32); // 32 columns like minimal test + tmem_alloc(smem_ptr, 32); } __syncthreads(); uint32_t tmem_base = *sTmemBase; - // Debug: write tmem_base to output BEFORE tmem_store - if (tid == 0) { - s_out[138] = (float)tmem_base; // should be 0.0f if first alloc - s_out[139] = 1.0f; // sentinel: reached this point - } - __syncthreads(); - - // Zero TMEM — skip (minimal test didn't zero) - // if (wid == 0) { - // tmem_store(tmem_base, 0, 0, 0, 0); - // __syncwarp(); - // tmem_store(tmem_base + 1, 0, 0, 0, 0); - // __syncwarp(); - // } - // __syncthreads(); - - if (tid == 0) { - s_out[140] = 2.0f; // sentinel: survived tmem_store - } - __syncthreads(); - - // ================================================================ // Load Q and K into SMEM in canonical layout - // ================================================================ write_q_to_smem<16>(sQ, q); write_k_to_smem<128, 16>(sK, k); __syncthreads(); - // ================================================================ - // Construct descriptors and instruction descriptor - // ================================================================ + // Construct descriptors uint32_t sQ_smem = __cvta_generic_to_shared(sQ); uint32_t sK_smem = __cvta_generic_to_shared(sK); - - // K-major NONE: LBO = BLOCK_MN * 16, SBO = 128 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); + uint32_t idesc = make_idesc(128, 32); // M=128, N=32 - // Debug: write descriptor values to output (positions 128-135) + // Debug output 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)__cvta_generic_to_shared(sQ); - s_out[136] = (float)(sQ_smem >> 4); // start_address for desc - s_out[137] = (float)(sQ_smem & 0xF); // alignment check + s_out[135] = (float)tmem_base; } __syncthreads(); - // ================================================================ - // SKIP MMA for now — just test descriptor construction - // ================================================================ - // if (tid == 0) { - // umma_ss_f16(tmem_base, desc_q, desc_k, idesc, /*accumulate=*/false); - // } - // __syncwarp(); - // if (wid == 0 && lane == 0) { - // tmem_fence_store(); - // } - // __syncthreads(); - - // ================================================================ - // Call tcgen05.mma SS — test with 32 columns (N=32) - // ================================================================ - // Rebuild idesc for N=32 - uint32_t idesc = make_idesc(128, 32); - + // Call tcgen05.mma SS if (tid == 0) { umma_ss_f16(tmem_base, desc_q, desc_k, idesc, /*accumulate=*/false); } @@ -150,29 +80,24 @@ test_umma_qk_hd16( } __syncthreads(); - // ================================================================ // Read S from TMEM - // ================================================================ if (wid == 0) { for (int col = 0; col < 32; col++) { uint32_t u0, u1, u2, u3; tmem_load(tmem_base + col, u0, u1, u2, u3); if (lane == 0) { - s_out[col] = u32_to_f32(u0); // S[0, col] + s_out[col] = u32_to_f32(u0); } } } __syncthreads(); - // ================================================================ // Scalar reference - // ================================================================ if (tid == 0) { for (int c = 0; c < 128; c++) { float dot = 0.0f; - for (int d = 0; d < 16; d++) { + for (int d = 0; d < 16; d++) dot += sQ_row[d] * bf16_to_f32(k[c * 16 + d]); - } s_scalar[c] = dot * scale; } } @@ -184,12 +109,8 @@ test_umma_qk_hd16( } } -// ================================================================== -// Host -// ================================================================== - int main() { - printf("=== UMMA QK GEMM Test (HD=16, SK=128) ===\n"); + 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); @@ -240,31 +161,23 @@ int main() { printf("\n"); float max_diff = 0.0f, max_val = 0.0f; - for (int c = 0; c < SK; c++) { + 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); - // Debug: print descriptor values - uint64_t dbg_desc_q, dbg_desc_k; - uint32_t dbg_idesc; - memcpy(&dbg_desc_q, &h_s_out[128], 8); - memcpy(&dbg_desc_k, &h_s_out[130], 8); - memcpy(&dbg_idesc, &h_s_out[132], 4); - printf("desc_q = 0x%016lx\n", dbg_desc_q); - printf("desc_k = 0x%016lx\n", dbg_desc_k); - printf("idesc = 0x%08x\n", dbg_idesc); - printf("sQ_smem = %.0f\n", h_s_out[133]); - printf("sK_smem = %.0f\n", h_s_out[134]); - printf(" desc_q start_addr = %lu, LBO = %lu, SBO = %lu\n", - dbg_desc_q & 0x3FFF, (dbg_desc_q >> 16) & 0x3FFF, (dbg_desc_q >> 32) & 0x3FFF); - printf(" desc_k start_addr = %lu, LBO = %lu, SBO = %lu\n", - dbg_desc_k & 0x3FFF, (dbg_desc_k >> 16) & 0x3FFF, (dbg_desc_k >> 32) & 0x3FFF); + // 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("Test %s\n", rel_err < 0.01f ? "PASSED" : "FAILED"); - 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;