From 526fafb8086396d07cc0b36e4621364654437c78 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 12:16:09 +0000 Subject: [PATCH] test: revert volatile, fix wid==0, full 4 K-tiles --- tests/unit/test_umma_qk_hd64.cu | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_umma_qk_hd64.cu b/tests/unit/test_umma_qk_hd64.cu index bd9cac73..7bba6c54 100644 --- a/tests/unit/test_umma_qk_hd64.cu +++ b/tests/unit/test_umma_qk_hd64.cu @@ -31,7 +31,7 @@ test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k, bf16_t* sK = sQ + 128 * hd; // TMEM alloc - if (wid == 1) { + if (wid == 0) { tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128); } __syncthreads(); @@ -42,18 +42,15 @@ test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k, __syncthreads(); // Write Q (1, hd) to sQ row 0 in canonical layout - // Use volatile to prevent compiler optimization eliminating writes - volatile bf16_t* vsQ = (volatile bf16_t*)sQ; - volatile bf16_t* vsK = (volatile bf16_t*)sK; for (int d = tid; d < hd; d += 128) { int ck = d / 8, lc = d % 8; - vsQ[ck * 16 * 64 + lc] = q[d]; + sQ[ck * 16 * 64 + lc] = q[d]; } // Write K (sk, hd) to sK in canonical layout for (int i = tid; i < sk * hd; i += 128) { int r = i / hd, c = i % hd; int tmn = r / 8, ck = c / 8, lr = r % 8, lc = c % 8; - vsK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[i]; + sK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[i]; } __syncthreads(); @@ -70,7 +67,7 @@ test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k, uint32_t idesc = make_idesc(128, 128); // K-tile loop with accumulate - for (int kt = 0; kt < 1; kt++) { // 1 K-tile only + for (int kt = 0; kt < hd / 16; kt++) { // Full K-tile loop // K-tile kt: columns [16*kt, 16*kt+16) // In canonical layout, columns start at core_k = 2*kt and 2*kt+1 // Offset = 2*kt * 2048 bytes from matrix base