diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index 48ad1cba..d62f472b 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -78,9 +78,9 @@ test_umma_qk_hd16(
     }
     __syncthreads();
 
-    // Zero TMEM
+    // Zero TMEM — test with 32 columns first
     if (wid == 0) {
-        for (int col = 0; col < 128; col++) {
+        for (int col = 0; col < 32; col++) {
             tmem_store(tmem_base + col, 0, 0, 0, 0);
         }
         tmem_fence_store();