diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index 10c5c8e1..eb7aecbe 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -90,28 +90,11 @@ test_umma_qk_hd16(
     // Use separate tmem_load calls with __syncthreads between them
     // to avoid the loop crash
     if (wid == 0) {
-        #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);
+        uint32_t u0, u1, u2, u3;
+        tmem_load(tmem_base + 1, u0, u1, u2, u3);
+        if (lane == 0) s_out[0] = u32_to_f32(u0);
     }
     __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) {