test: re-enable full TMEM zeroing (128 columns)

2026-05-28 09:27:25 +00:00
parent c3b6c3a5e6
commit 3694f63ba4
1 changed files with 4 additions and 2 deletions
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -78,9 +78,11 @@ test_umma_qk_hd16(
    }
    __syncthreads();

-    // Zero TMEM — only zero 1 column to test
+    // Zero TMEM
    if (wid == 0) {
-        tmem_store(tmem_base, 0, 0, 0, 0);
+        for (int col = 0; col < 128; col++) {
+            tmem_store(tmem_base + col, 0, 0, 0, 0);
+        }
        tmem_fence_store();
    }
    __syncthreads();