diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index f8337ec8..f32ed781 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -71,23 +71,22 @@ test_umma_qk_hd16(
     __syncthreads();
     uint32_t tmem_base = *sTmemBase;
 
-    // Debug: write tmem_base to output
+    // Debug: write tmem_base to output BEFORE tmem_store
     if (tid == 0) {
-        s_out[138] = 999.0f;  // sentinel before tmem ops
+        s_out[138] = (float)tmem_base;  // should be 0.0f if first alloc
+        s_out[139] = 1.0f;  // sentinel: reached this point
     }
     __syncthreads();
 
-    // Zero TMEM — re-enable with debug
+    // Zero TMEM — only zero 1 column to test
     if (wid == 0) {
-        for (int col = 0; col < 128; col++) {
-            tmem_store(tmem_base + col, 0, 0, 0, 0);
-        }
+        tmem_store(tmem_base, 0, 0, 0, 0);
         tmem_fence_store();
     }
     __syncthreads();
 
     if (tid == 0) {
-        s_out[138] = (float)tmem_base;  // write tmem_base after tmem ops
+        s_out[140] = 2.0f;  // sentinel: survived tmem_store
     }
     __syncthreads();