From 926ae5d7bf0acd050b45233275eaa7d5571585b0 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 12:54:03 +0000
Subject: [PATCH] test: fix K source stride mismatch in manual SMEM write

---
 tests/unit/test_umma_qk_hd64.cu | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/tests/unit/test_umma_qk_hd64.cu b/tests/unit/test_umma_qk_hd64.cu
index 1c003652..85a025c9 100644
--- a/tests/unit/test_umma_qk_hd64.cu
+++ b/tests/unit/test_umma_qk_hd64.cu
@@ -41,9 +41,25 @@ test_umma_hd64_debug(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
     __syncthreads();
     uint32_t tb = *sTmemBase;
 
-    // Load Q and K into SMEM in canonical layout — ONLY first 16 dims
-    write_q_to_smem<16>(sQ, q);
-    write_k_to_smem<128, 16>(sK, k);  // k is (128, 64), but we only read first 16 cols
+    // Load Q (first 16 dims) into (128, 16) canonical — manual write
+    // Zero sQ
+    for (int i = tid; i < 128 * 16; i += 128) sQ[i] = 0;
+    // Q row 0, first 16 dims
+    for (int c = tid; c < 16; c += 128) {
+        int ck = c / 8, lc = c % 8;
+        sQ[ck * 16 * 64 + lc] = q[c];
+    }
+
+    // Load K (first 16 cols of (128, 64)) into (128, 16) canonical
+    // K source stride is 64, but we only write 16 columns to SMEM
+    for (int i = tid; i < 128 * 16; i += 128) sK[i] = 0;
+    for (int r = 0; r < 128; r++) {
+        for (int c = tid; c < 16; c += 128) {
+            int ck = c / 8, lc = c % 8;
+            int tmn = r / 8, lr = r % 8;
+            sK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * 64 + c];
+        }
+    }
     bf16_t* sQ_pad = sQ + 128 * 16;
     for (int i = tid; i < 4096; i += 128) sQ_pad[i] = 0;
     __syncthreads();