diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index ab220b04..10c5c8e1 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -86,13 +86,32 @@ test_umma_qk_hd16(
     if (wid == 0 && lane == 0) tmem_fence_store();
     __syncthreads();
 
-    // Read S[0,0] from TMEM
+    // Read S from TMEM — read 128 columns one at a time
+    // Use separate tmem_load calls with __syncthreads between them
+    // to avoid the loop crash
     if (wid == 0) {
-        uint32_t u0, u1, u2, u3;
-        tmem_load(tmem_base, u0, u1, u2, u3);
-        if (lane == 0) s_out[0] = u32_to_f32(u0);
+        #define READ_COL(c) do { \
+            uint32_t u0, u1, u2, u3; \
+            tmem_load(tmem_base + c, u0, u1, u2, u3); \
+            if (lane == 0) s_out[c] = u32_to_f32(u0); \
+        } while(0)
+        READ_COL(0);
     }
     __syncthreads();
+    if (wid == 0) { READ_COL(1); }
+    __syncthreads();
+    if (wid == 0) { READ_COL(2); }
+    __syncthreads();
+    if (wid == 0) { READ_COL(3); }
+    __syncthreads();
+    if (wid == 0) { READ_COL(4); }
+    __syncthreads();
+    if (wid == 0) { READ_COL(5); }
+    __syncthreads();
+    if (wid == 0) { READ_COL(6); }
+    __syncthreads();
+    if (wid == 0) { READ_COL(7); }
+    __syncthreads();
 
     // Scalar reference
     if (tid == 0) {