diff --git a/tests/unit/test_umma_qk_hd64.cu b/tests/unit/test_umma_qk_hd64.cu index c848780b..22b333e5 100644 --- a/tests/unit/test_umma_qk_hd64.cu +++ b/tests/unit/test_umma_qk_hd64.cu @@ -1,12 +1,16 @@ /** * UMMA QK GEMM Test — HD=64 (4 K-tiles) * - * Fixed from previous version: - * 1. Zero TMEM before accumulate loop (tcgen05.alloc does NOT zero) - * 2. tcgen05.fence::after_thread_sync after MMA loop (correct MMA→TMEM fence) - * 3. Single-thread MMA call (tid==0, matching gau-nernst's elect_one pattern) - * 4. Separate SMEM per K-tile (avoids offset descriptor issues) - * 5. Full 128×128 output validation + * Using gau-nernst's approach: single contiguous SMEM, offset descriptors. + * Q is (128, HD) and K is (128, HD) in canonical K-major layout. + * For each K-tile kt, the descriptor start address is: + * base + kt * BLOCK_MN * 32 (each 16-BF16 K-tile is BLOCK_MN * 32 bytes apart) + * + * Fixes from previous version: + * 1. Zero TMEM before accumulate (tcgen05.alloc does NOT zero) + * 2. tcgen05.fence::after_thread_sync after each MMA + before TMEM read + * 3. Single-thread MMA call (tid==0) + * 4. Contiguous SMEM with offset descriptors */ #include @@ -25,7 +29,7 @@ static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; m constexpr int HD = 64; constexpr int SK = 128; -constexpr int NKT = HD / MMA_K_BF16; // 4 K-tiles +constexpr int NKT = HD / MMA_K_BF16; // 4 constexpr int BLOCK_MN = 128; __global__ void __launch_bounds__(128) @@ -36,29 +40,29 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, const int wid = tid / WARP, lane = tid % WARP; // SMEM layout: - // [0..3] tmem_base (written by tcgen05.alloc) - // [4..) K-tile SMEM regions (alternating Q, K per K-tile) - // Each (128, 16) = 128*16*2 = 4096 bytes in canonical layout + // [0..3] tmem_base + // [64..) sQ (128, HD) canonical = 128*64*2 = 16384 bytes + // sK (128, HD) canonical = 16384 bytes + // Align to 128 bytes for SMEM access patterns extern __shared__ char sbuf[]; uint32_t* sTmemBase = (uint32_t*)sbuf; - // Align SMEM data to 16 bytes for descriptor start_address - char* data_base = (char*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15); - // Separate SMEM for each Q K-tile and K K-tile - bf16_t* sQ[NKT]; - bf16_t* sK[NKT]; - for (int kt = 0; kt < NKT; kt++) { - sQ[kt] = (bf16_t*)(data_base + (kt * 2 + 0) * BLOCK_MN * MMA_K_BF16 * sizeof(bf16_t)); - sK[kt] = (bf16_t*)(data_base + (kt * 2 + 1) * BLOCK_MN * MMA_K_BF16 * sizeof(bf16_t)); - } + bf16_t* sQ = (bf16_t*)(sbuf + 64); // 64-byte aligned + bf16_t* sK = sQ + BLOCK_MN * HD; // follows Q - // Step 1: TMEM alloc — 128 columns (Layout D for M=128, N=128) + // Load Q (1, HD) into (128, HD) canonical layout — only row 0 has data + write_q_to_smem(sQ, q); + // Load K (SK, HD) into (128, HD) canonical layout + write_k_to_smem(sK, k); + __syncthreads(); + + // TMEM alloc — 128 columns for (128, 128) Layout D output if (wid == 0) { tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128); } __syncthreads(); uint32_t tb = *sTmemBase; - // Step 2: Zero TMEM (tcgen05.alloc does NOT zero) + // Zero TMEM (tcgen05.alloc does NOT zero) if (wid == 0) { for (int col = 0; col < 128; col++) { tmem_store(tb + col, 0, 0, 0, 0); @@ -67,53 +71,32 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, } __syncthreads(); - // Step 3: Load Q and K into per-K-tile SMEM regions - // Global Q is (1, HD). Global K is (SK, HD). - // K-tile kt: columns [16*kt, 16*kt+16) of the (128, HD) padded matrix. - for (int kt = 0; kt < NKT; kt++) { - bf16_t* sq = sQ[kt]; - bf16_t* sk = sK[kt]; - - // Zero the (128, 16) tile - for (int i = tid; i < BLOCK_MN * MMA_K_BF16; i += 128) { - sq[i] = 0; - sk[i] = 0; - } - __syncthreads(); - - // Q: only row 0 has data. Write to canonical layout. - // For (128, 16): CORES_MN=16, CORES_K=2 - // Row 0: core_mn=0, local_r=0 - for (int d = tid; d < MMA_K_BF16; d += 128) { - int ck = d / 8, lc = d % 8; - sq[ck * 16 * 64 + lc] = q[kt * MMA_K_BF16 + d]; - } - - // K: write (SK, 16) into (128, 16) canonical layout - for (int i = tid; i < SK * MMA_K_BF16; i += 128) { - int r = i / MMA_K_BF16; - int c = i % MMA_K_BF16; - int tmn = r / 8, ck = c / 8, lr = r % 8, lc = c % 8; - sk[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * MMA_K_BF16 + c]; - } - } - __syncthreads(); - - // Step 4: Multi-K-tile QK GEMM: Q × K^T → S in TMEM - // Each MMA call: (128, 16) × (16, 128) → (128, 128) partial, accumulated in TMEM + // Multi-K-tile QK GEMM + // For K-major NONE layout, each 16-column K-tile starts at: + // base_smem + kt * BLOCK_MN * 32 (in bytes) + // Each K-tile spans 2 core-matrix columns (16 BF16). + // LBO = BLOCK_MN * 16 (bytes) = stride between the 2 core columns. + // The descriptor describes a (BLOCK_MN, 16) sub-matrix. + uint32_t sQ_smem = __cvta_generic_to_shared(sQ); + uint32_t sK_smem = __cvta_generic_to_shared(sK); uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN); for (int kt = 0; kt < NKT; kt++) { - // Construct descriptor for this K-tile - uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ[kt]), BLOCK_MN); - uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK[kt]), BLOCK_MN); + // Offset to the kt-th K-tile: kt * BLOCK_MN * 32 bytes + // (each K-tile is 16 BF16 = 2 core columns, BLOCK_MN * 16B apart, total 32B per row) + // Wait, that's not right. In canonical layout, the kt-th 16-column slice starts + // at the (2*kt)-th core column. Each core column is BLOCK_MN/8 * 128 bytes + // = BLOCK_MN * 16 bytes. So offset = 2*kt * BLOCK_MN * 16 bytes = kt * BLOCK_MN * 32 bytes. + uint32_t q_addr = sQ_smem + kt * BLOCK_MN * 32; + uint32_t k_addr = sK_smem + kt * BLOCK_MN * 32; + + uint64_t dq = make_umma_desc_kmajor_none(q_addr, BLOCK_MN); + uint64_t dk = make_umma_desc_kmajor_none(k_addr, BLOCK_MN); - // Single-thread MMA call (matching gau-nernst's elect_one pattern) - // accumulate=true for all tiles (TMEM zeroed above) if (tid == 0) { - umma_ss_f16(tb, dq, dk, idesc, true); + umma_ss_f16(tb, dq, dk, idesc, true); // accumulate across K-tiles } - // Ensure MMA completes before next K-tile + // Fence after each K-tile MMA to ensure TMEM is updated asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); __syncthreads(); } @@ -122,9 +105,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); __syncthreads(); - // Step 5: Read S from TMEM (Layout D: 32x32b.x8) - // 128 TMEM columns, 128 rows. Each warp reads 32 rows × 8 columns. - // 4 warps × 16 iterations = 128 rows × 128 columns + // Read S from TMEM (Layout D: 32x32b.x8) for (int n = 0; n < 128 / 8; n++) { const int row = wid * 32; const int col = n * 8; @@ -145,7 +126,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, } __syncthreads(); - // Step 6: Scalar reference — S[0, j] = sum(Q[0,d]*K[j,d], d=0..63) * scale + // Scalar reference if (tid == 0) { for (int j = 0; j < SK; j++) { float dot = 0.0f; @@ -155,12 +136,11 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, } } - // TMEM dealloc if (wid == 0) tmem_dealloc(tb, 128); } int main() { - printf("=== UMMA QK GEMM HD=64 (4 K-tiles, fixed) ===\n"); + printf("=== UMMA QK GEMM HD=64 (4 K-tiles, contiguous SMEM) ===\n"); const float SCALE = 1.0f / sqrtf((float)HD); bf16_t* h_q = (bf16_t*)malloc(HD * sizeof(bf16_t)); @@ -178,8 +158,8 @@ int main() { cudaMemcpy(d_q, h_q, HD*sizeof(bf16_t), cudaMemcpyHostToDevice); cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice); - // SMEM: tmem_base(4) + alignment(16) + 8 K-tiles × 128×16×2 - int smem = (4 + 16 + NKT * 2 * BLOCK_MN * MMA_K_BF16 * sizeof(bf16_t) + 256 + 127) & ~127; + // SMEM: 64 (tmem_base+pad) + 2 * 128*64*2 (Q+K) + 256 (extra) + int smem = (64 + 2 * BLOCK_MN * HD * sizeof(bf16_t) + 256 + 127) & ~127; printf("SMEM: %d bytes (%d KB)\n", smem, smem / 1024); test_umma_hd64<<<1, 128, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE); @@ -190,13 +170,11 @@ int main() { cudaMemcpy(h_s_out, d_s_out, SK*SK*sizeof(float), cudaMemcpyDeviceToHost); cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost); - // Print row 0 (the only row with data since Q is 1×HD) printf("S[0,0..7] MMA: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n"); printf("S[0,0..7] ref: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_scalar[c]); printf("\n"); printf("S[0,120..127] MMA: "); for(int c=120;c<128;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n"); printf("S[0,120..127] ref: "); for(int c=120;c<128;c++) printf("%.6f ",h_s_scalar[c]); printf("\n"); - // Validate row 0 (full 128 columns) float max_diff = 0.0f, max_val = 0.0f; for (int c = 0; c < SK; c++) { float diff = fabsf(h_s_out[0*SK+c] - h_s_scalar[c]); @@ -206,13 +184,10 @@ int main() { float rel_err = max_val > 0 ? max_diff / max_val : max_diff; printf("Row 0 rel err (128 cols): %.8f\n", rel_err); - // Check rows 1-127 should be zero (Q has no data there) float max_nonzero = 0.0f; - for (int r = 1; r < SK; r++) { - for (int c = 0; c < SK; c++) { + for (int r = 1; r < SK; r++) + for (int c = 0; c < SK; c++) max_nonzero = fmaxf(max_nonzero, fabsf(h_s_out[r*SK+c])); - } - } printf("Rows 1-127 max abs: %.8f\n", max_nonzero); bool row0_ok = rel_err < 0.001f;