test: read only 1 TMEM column after MMA

2026-05-28 09:35:02 +00:00
parent 41128122e3
commit 004046a6a8
1 changed files with 5 additions and 7 deletions
--- 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();