diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index 9e6e59be..239ff9fe 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -1,8 +1,7 @@ /** - * 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. + * UMMA QK GEMM Test (HD=16, SK=128) + * Uses Layout D TMEM read (32x32b.x8, 4 warps) to correctly read MMA output. + * Debugging the 4× scaling factor observed with 16x256b read. */ #include @@ -29,73 +28,57 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k, 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 + 4096; // 8KB padding after Q + bf16_t* sK = sQ + 128 * 16 + 4096; float* sQ_row = (float*)(sK + 128 * 16); for (int d = tid; d < 16; d += NTHREADS) sQ_row[d] = bf16_to_f32(q[d]); - // TMEM alloc (128 cols for MMA output N=128) + // TMEM alloc (128 cols) if (wid == 0) { - uint32_t sp = __cvta_generic_to_shared(sTmemBase); - tmem_alloc(sp, 128); + tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128); } __syncthreads(); - uint32_t tmem_base = *sTmemBase; + uint32_t tb = *sTmemBase; - // Load Q and K into SMEM + // Load Q and K 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, 64); // M=64 - uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 64); // M=64 - - // Test with different N values to understand scaling - // N=8 → n_dim=1 - uint32_t idesc = make_idesc(64, 128); // Try M=64 - - 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(); + 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); // MMA if (tid == 0) { - umma_ss_f16(tmem_base, desc_q, desc_k, idesc, false); + umma_ss_f16(tb, desc_q, desc_k, idesc, false); } __syncwarp(); if (wid == 0 && lane == 0) tmem_fence_store(); __syncthreads(); - // Read S[0,0] from TMEM column 0 - 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); + // Read from TMEM using Layout D: 32x32b.x8 format + // Each warp reads 32 rows × 8 columns + // Warp 0: rows 0-31, warp 1: rows 32-63, warp 2: rows 64-95, warp 3: rows 96-127 + if (wid < 4) { + const int row = wid * 32; + const int col = 0; + const int addr = tb + (row << 16) + col; + float tmp[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]), + "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7]) + : "r"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + // Lane 0 gets S[wid*32 + 0, 0..7] in tmp[0..7] + if (lane == 0) { + for (int c = 0; c < 8; c++) { + s_out[wid * 8 + c] = tmp[c]; + } + } } __syncthreads(); @@ -110,11 +93,11 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k, } __syncthreads(); - if (wid == 0) tmem_dealloc(tmem_base, 128); + if (wid == 0) tmem_dealloc(tb, 128); } int main() { - printf("=== UMMA QK GEMM Test (HD=16) ===\n"); + printf("=== UMMA QK GEMM Test (Layout D read) ===\n"); const int HD = 16, SK = 128; const float SCALE = 1.0f / sqrtf((float)HD); @@ -143,15 +126,18 @@ int main() { 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] 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("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"); + // Print S[0,0] from each warp's perspective + printf("S from TMEM (Layout D, 32x32b.x8):\n"); + for (int w = 0; w < 4; w++) { + printf(" Warp %d (rows %d-%d): ", w, w*32, w*32+31); + for (int c = 0; c < 8; c++) printf("%.4f ", h_s_out[w*8+c]); + printf("\n"); + } + printf("S[0,0..7] scalar: "); + for (int c = 0; c < 8; c++) printf("%.4f ", h_s_scalar[c]); + printf("\n"); + printf("Ratio warp0/col0 vs scalar[0]: %.4f\n", + h_s_scalar[0] != 0 ? h_s_out[0] / h_s_scalar[0] : 0); 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);