From b703dc0a50238aa051f8bff6f1e1694570f28c16 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 12:51:33 +0000
Subject: [PATCH] test: debug single K-tile with offset descriptor

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

diff --git a/tests/unit/test_umma_qk_hd64.cu b/tests/unit/test_umma_qk_hd64.cu
index 32b15c46..ff6704ec 100644
--- a/tests/unit/test_umma_qk_hd64.cu
+++ b/tests/unit/test_umma_qk_hd64.cu
@@ -76,7 +76,7 @@ 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 < NKT; kt++) {
+    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
@@ -89,7 +89,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
         uint64_t dk = make_umma_desc_kmajor_none(k_addr, BLOCK_MN);
 
         if (tid == 0) {
-            umma_ss_f16(tb, dq, dk, idesc, kt > 0);  // first tile: no accumulate
+            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");
@@ -125,7 +125,7 @@ test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
     if (tid == 0) {
         for (int j = 0; j < SK; j++) {
             float dot = 0.0f;
-            for (int d = 0; d < HD; d++)
+            for (int d = 0; d < 16; d++)  // DEBUG: only first K-tile
                 dot += bf16_to_f32(q[d]) * bf16_to_f32(k[j * HD + d]);
             s_scalar[j] = dot * scale;
         }