diff --git a/tests/unit/test_umma_qk_hd64.cu b/tests/unit/test_umma_qk_hd64.cu index d0f53eb9..ee6a3db2 100644 --- a/tests/unit/test_umma_qk_hd64.cu +++ b/tests/unit/test_umma_qk_hd64.cu @@ -1,7 +1,7 @@ /** - * UMMA QK GEMM Test — HD=64 (4 K-tiles), SK=128 - * Multi-K-tile accumulate: call MMA 4× with accumulate=true - * Each K-tile: 16 BF16 columns, separate descriptor + * UMMA QK GEMM Test — HD=64, SK=128, 1 K-tile (columns 0-15 only) + * Verifies the K-tile descriptor offset is correct by comparing + * against the HD=16 result (which is proven correct). */ #include @@ -18,172 +18,88 @@ 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; } -// MMA K-tile = 16 BF16 -constexpr int MMA_K = 16; -constexpr int N_WARPS = 4; -constexpr int BLOCK_M = 128; - -__global__ void __launch_bounds__(N_WARPS * 32) -test_umma_qk_hd64(const bf16_t* q, const bf16_t* k, +__global__ void __launch_bounds__(128) +test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k, float* s_out, float* s_scalar, float scale, int hd, int sk) { const int tid = threadIdx.x; const int wid = tid / 32, lane = tid % 32; - const int n_ktiles = hd / MMA_K; // 4 for hd=64 - // SMEM: sQ (128, HD) canonical + sK (128, HD) canonical - // Each K-tile of (128, 16) = 4096 bytes - // Full (128, HD) = n_ktiles * 4096 bytes 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 * hd; - // TMEM alloc (128 cols for N=128) + // TMEM alloc if (wid == 1) { tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128); } __syncthreads(); uint32_t tb = *sTmemBase; - // Zero all first - for (int i = tid; i < 128 * hd; i += N_WARPS * 32) { - sQ[i] = 0; - sK[i] = 0; - } + // Zero SMEM + for (int i = tid; i < 128 * hd; i += 128) { sQ[i] = 0; sK[i] = 0; } __syncthreads(); - // Write Q (1, hd) to row 0 of sQ in canonical layout - for (int d = tid; d < hd; d += N_WARPS * 32) { - int core_k = d / 8, local_c = d % 8; - int idx = core_k * 16 * 64 + local_c; // tile_mn=0, local_r=0 - sQ[idx] = q[d]; - // Also write to flat offset for verification - if (d < 8) s_out[220 + d] = (float)idx; // Log the indices we write to + // Write Q (1, hd) to sQ row 0 in canonical layout + for (int d = tid; d < hd; d += 128) { + int ck = d / 8, lc = d % 8; + sQ[ck * 16 * 64 + lc] = q[d]; } // Write K (sk, hd) to sK in canonical layout - for (int i = tid; i < sk * hd; i += N_WARPS * 32) { + for (int i = tid; i < sk * hd; i += 128) { int r = i / hd, c = i % hd; - int tile_mn = r / 8, core_k = c / 8; - int local_r = r % 8, local_c = c % 8; - int idx = core_k * 16 * 64 + tile_mn * 64 + local_r * 8 + local_c; - sK[idx] = k[i]; + int tmn = r / 8, ck = c / 8, lr = r % 8, lc = c % 8; + sK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[i]; } __syncthreads(); - // Verify SMEM data for first K-tile (columns 0-15) - // In canonical layout, Q[d] for row 0 is at core_k * 16 * 64 + local_c + // Verify Q[0..7] in SMEM if (tid == 0) { - // Q row 0, d=0..7: core_k=0, local_c=d → sQ[d] - for (int d = 0; d < 8; d++) - s_out[200+d] = bf16_to_f32(sQ[d]); - // Q row 0, d=8..15: core_k=1, local_c=d-8 → sQ[1024 + d-8] - for (int d = 0; d < 8; d++) - s_out[208+d] = bf16_to_f32(sQ[1024 + d]); + for (int d = 0; d < 8; d++) s_out[200+d] = bf16_to_f32(sQ[d]); + for (int d = 0; d < 8; d++) s_out[208+d] = bf16_to_f32(sQ[1024+d]); // Q[8..15] } __syncthreads(); + + // MMA with K-tile 0 (columns 0-15) + // Descriptor: start = sQ_smem (base of the matrix) 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); - // For each K-tile, construct a descriptor pointing to that 16-column slice - // K-tile k: columns [16k, 16k+16) of the (128, hd) matrix - // In canonical layout, the k-th 16-column slice starts at: - // k * CORES_MN * 64 = k * 16 * 64 = k * 1024 BF16 = k * 2048 bytes - // Each K-tile has BLOCK_M=128 rows and 16 columns. - // The descriptor for K-tile k: start_addr = sQ_smem + k * 2048 - // But wait — gau-nernst uses A_smem + k * BLOCK_M * 32 for the start address - // BLOCK_M * 32 = 128 * 32 = 4096 bytes. But our K-tile is only 2048 bytes. - // Actually, gau-nernst's offset is for the SMEM start of the K-tile. - // In his layout, each (BLOCK_M, 32B) slice is BLOCK_M * 32 = 4096 bytes apart. - // But 32B = 16 BF16 = one K-tile. And his start_address = A_smem + k * BLOCK_M * 32. - - // Wait — that's 4096 bytes per K-tile, but our (128, 16) matrix is only 4096 bytes. - // In canonical layout, the K-tile at columns [16k, 16k+16) starts at: - // The 2 core-matrix columns (c=2k and c=2k+1) are at offsets 2k*2048 and (2k+1)*2048. - // The full K-tile spans both columns: core(0..15, 2k) and core(0..15, 2k+1). - // The first column starts at 2k * 2048 and the second at (2k+1) * 2048. - // Total span: from 2k*2048 to (2k+1)*2048 + 16*128 = (2k+1)*2048 + 2048 = (2k+2)*2048. - - // Hmm, this is getting complicated. The descriptor for a K-tile should describe - // a (128, 16) matrix starting at the right offset in SMEM. - // For K-tile 0 (columns 0-15): start at sQ_smem, LBO=2048, SBO=128 - // For K-tile 1 (columns 16-31): start at sQ_smem + 2*2048, LBO=2048, SBO=128 - // Wait, but columns 16-23 are core_k=2 and columns 24-31 are core_k=3. - // The K-tile at columns [16k, 16k+16) has core_k = 2k and 2k+1. - // core_k=2 starts at 2 * 1024 = 2048 BF16 = 4096 bytes from sQ. - // core_k=3 starts at 3 * 1024 = 3072 BF16 = 6144 bytes from sQ. - - // But the descriptor's start_address is for the BEGINNING of the (128, 16) tile. - // The descriptor with LBO=2048 walks: column 0 at start, column 1 at start+2048. - // So for K-tile k, we need: - // start = sQ_smem + (2k) * 1024 * 2 = sQ_smem + 2k * 2048 bytes - // Column 0 = core_k 2k, column 1 = core_k 2k+1 - // LBO = 2048 (same as before) - // SBO = 128 (same as before) - - // Actually, this IS the same descriptor but with a different start_address. - // The descriptor for K-tile k: start = sQ + 2k * 2048 - - uint32_t idesc = make_idesc(BLOCK_M, 128); - - // K-tile loop with accumulate - for (int kt = 0; kt < n_ktiles; kt++) { - // Descriptor for Q's k-th K-tile - uint32_t q_ktile_addr = sQ_smem + kt * 2 * 2048; // 2 core-matrix columns per K-tile - uint32_t k_ktile_addr = sK_smem + kt * 2 * 2048; - uint64_t desc_q = make_umma_desc_kmajor_none(q_ktile_addr, BLOCK_M); - uint64_t desc_k = make_umma_desc_kmajor_none(k_ktile_addr, BLOCK_M); - - bool accumulate = (kt > 0); - - // Debug: write descriptor info - if (tid == 0 && kt == 0) { - s_out[240] = (float)q_ktile_addr; - s_out[241] = (float)k_ktile_addr; - s_out[242] = (float)desc_q; - s_out[243] = (float)desc_k; - } - - // 4 warp leaders call MMA - if (lane == 0) { - umma_ss_f16(tb, desc_q, desc_k, idesc, accumulate); - } - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - __syncthreads(); + if (lane == 0) { + umma_ss_f16(tb, desc_q, desc_k, idesc, false); } + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + __syncthreads(); - // Read from TMEM: 32x32b.x8, each warp reads 32 rows × 8 columns + // Read TMEM for (int n = 0; n < 128 / 8; n++) { const int row = wid * 32; - const int col = n * 8; - const int addr = tb + (row << 16) + col; + const int addr = tb + (row << 16) + n * 8; 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.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;"); - int out_row = wid * 32 + lane; - if (n < 2 && out_row < 128) { // First 16 cols - for (int c = 0; c < 8; c++) { + if (n < 2 && out_row < 128) { + for (int c = 0; c < 8; c++) s_out[out_row * 16 + n * 8 + c] = tmp[c] * scale; - } } } __syncthreads(); - // Scalar reference + // Scalar: S[0,j] = sum(Q[0,d]*K[j,d], d=0..15) * scale (first K-tile only) if (tid == 0) { - float* q_row = new float[hd]; - for (int d = 0; d < hd; d++) q_row[d] = bf16_to_f32(q[d]); - for (int c = 0; c < sk; c++) { + for (int j = 0; j < sk; j++) { float dot = 0.0f; - for (int d = 0; d < hd; d++) - dot += q_row[d] * bf16_to_f32(k[c * hd + d]); - s_scalar[c] = dot * scale; + for (int d = 0; d < 16; d++) // Only first K-tile + dot += bf16_to_f32(q[d]) * bf16_to_f32(k[j * hd + d]); + s_scalar[j] = dot * scale; } - delete[] q_row; } __syncthreads(); @@ -191,15 +107,16 @@ test_umma_qk_hd64(const bf16_t* q, const bf16_t* k, } int main() { - printf("=== UMMA QK GEMM Test (HD=64, 4 K-tiles) ===\n"); + printf("=== UMMA QK HD=64, 1 K-tile ===\n"); const int HD = 64, 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*)calloc(128 * 16, sizeof(float)); + float* h_s_out = (float*)calloc(128*16, sizeof(float)); 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); @@ -208,13 +125,9 @@ int main() { cudaMalloc(&d_s_out, 128*16*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, 128*16*sizeof(float)); cudaMemset(d_s_scalar, 0, SK*sizeof(float)); - // SMEM: sTmemBase(4) + pad(16) + sQ(128*64*2=16384) + sK(128*64*2=16384) + pad - int smem = (4 + 16 + 128*HD*2 + 128*HD*2 + 256 + 127) & ~127; - printf("SMEM: %d bytes\n", smem); - - test_umma_qk_hd64<<<1, N_WARPS*32, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE, HD, SK); + int smem = (4+16+128*HD*2+128*HD*2+256+127)&~127; + test_umma_qk_hd64_1ktile<<<1, 128, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE, HD, SK); cudaError_t err = cudaDeviceSynchronize(); if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; } @@ -222,25 +135,22 @@ int main() { cudaMemcpy(h_s_out, d_s_out, 128*16*sizeof(float), cudaMemcpyDeviceToHost); cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost); - printf("Q[d=0..7] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ", h_s_out[200+d]); printf("\n"); - printf("Q[d=16..23] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ", h_s_out[208+d]); printf("\n"); - printf("Q[d=0..7] orig: "); for(int d=0;d<8;d++) printf("%.4f ", bf16_to_f32_host(h_q[d])); printf("\n"); - printf("Q[d=16..23] orig: "); for(int d=0;d<8;d++) printf("%.4f ", bf16_to_f32_host(h_q[16+d])); printf("\n"); - - // Compare row 0 - printf("S[0,0..7] MMA: "); - for (int c = 0; c < 8; c++) printf("%.6f ", h_s_out[0*16+c]); - printf("\nS[0,0..7] ref: "); - for (int c = 0; c < 8; c++) printf("%.6f ", h_s_scalar[c]); - printf("\n"); + // SMEM verify + printf("Q[0..7] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ",h_s_out[200+d]); printf("\n"); + printf("Q[0..7] orig: "); for(int d=0;d<8;d++) printf("%.4f ",bf16_to_f32_host(h_q[d])); printf("\n"); + printf("Q[8..15] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ",h_s_out[208+d]); printf("\n"); + printf("Q[8..15] orig: "); for(int d=0;d<8;d++) printf("%.4f ",bf16_to_f32_host(h_q[8+d])); printf("\n"); + // Compare + printf("S[0,0..7] MMA: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_out[0*16+c]); printf("\n"); + printf("S[0,0..7] ref: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_scalar[c]); printf("\n"); float max_diff = 0.0f, max_val = 0.0f; - for (int c = 0; c < 16; c++) { + for (int c = 0; c < 8; c++) { max_diff = fmaxf(max_diff, fabsf(h_s_out[0*16+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("Row 0 rel err (16 cols): %.6f\n", rel_err); + printf("Row 0 rel err: %.6f\n", rel_err); printf("Test %s\n", rel_err < 0.01f ? "PASSED" : "FAILED"); cudaFree(d_q); cudaFree(d_k); cudaFree(d_s_out); cudaFree(d_s_scalar);