From 004046a6a8fe1bfead88bb3b914e5c78cda87035 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 09:35:02 +0000
Subject: [PATCH] test: read only 1 TMEM column after MMA

---
 tests/unit/test_umma_qk.cu | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

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();