From 3694f63ba49910d0ba1a050b66465a027e54801f Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 09:27:25 +0000
Subject: [PATCH] test: re-enable full TMEM zeroing (128 columns)

---
 tests/unit/test_umma_qk.cu | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index f32ed781..48ad1cba 100644
--- 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();