From db4f6618439a8068d7baaf1d0a1acf7726c44502 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 12:52:19 +0000 Subject: [PATCH] test: debug with (128,16) SMEM matching HD=16 exactly --- tests/unit/test_umma_qk_hd64.cu | 29 ++++++++--------------------- 1 file changed, 8 insertions(+), 21 deletions(-) diff --git a/tests/unit/test_umma_qk_hd64.cu b/tests/unit/test_umma_qk_hd64.cu index ff6704ec..a25ea42f 100644 --- a/tests/unit/test_umma_qk_hd64.cu +++ b/tests/unit/test_umma_qk_hd64.cu @@ -49,10 +49,9 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, bf16_t* sQ = (bf16_t*)(sbuf + 64); // 64-byte aligned bf16_t* sK = sQ + BLOCK_MN * HD; // follows Q - // 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); + // Load Q (1, 16) into (128, 16) canonical — only first K-tile for debug + write_q_to_smem<16>(sQ, q); + write_k_to_smem(sK, k); __syncthreads(); // TMEM alloc — 128 columns for (128, 128) Layout D output @@ -76,24 +75,12 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, uint32_t sK_smem = __cvta_generic_to_shared(sK); uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN); - for (int kt = 0; kt < 1; kt++) { // DEBUG: only first K-tile - // 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; + // Single K-tile: descriptors point to start of SMEM + uint64_t dq = make_umma_desc_kmajor_none(sQ_smem, BLOCK_MN); + uint64_t dk = make_umma_desc_kmajor_none(sK_smem, BLOCK_MN); - uint64_t dq = make_umma_desc_kmajor_none(q_addr, BLOCK_MN); - uint64_t dk = make_umma_desc_kmajor_none(k_addr, BLOCK_MN); - - if (tid == 0) { - umma_ss_f16(tb, dq, dk, idesc, false); // no accumulate, single tile - } - // Fence after each K-tile MMA to ensure TMEM is updated - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - __syncthreads(); + if (tid == 0) { + umma_ss_f16(tb, dq, dk, idesc, false); } // Final fence before TMEM read