diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index b7dd55bb..cdd0175e 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -80,14 +80,12 @@ test_umma_qk_hd16(
     }
     __syncthreads();
 
-    // Read S from TMEM
+    // Read S from TMEM — only read 1 column
     if (wid == 0) {
-        for (int col = 0; col < 32; col++) {
-            uint32_t u0, u1, u2, u3;
-            tmem_load(tmem_base + col, u0, u1, u2, u3);
-            if (lane == 0) {
-                s_out[col] = u32_to_f32(u0);
-            }
+        uint32_t u0, u1, u2, u3;
+        tmem_load(tmem_base, u0, u1, u2, u3);
+        if (lane == 0) {
+            s_out[0] = u32_to_f32(u0);  // S[0, 0]
         }
     }
     __syncthreads();