From ca419c52f352538e97d56fa2b9241f0be2b31823 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 09:23:10 +0000
Subject: [PATCH] test: re-enable TMEM alloc + zero

---
 tests/unit/test_umma_qk.cu | 32 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index 2a880916..f8e2c4f6 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -62,23 +62,23 @@ test_umma_qk_hd16(
     }
 
     // ================================================================
-    // TMEM allocation — SKIP for now
+    // TMEM allocation
     // ================================================================
-    // if (wid == 0) {
-    //     uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
-    //     tmem_alloc(smem_ptr, 128);
-    // }
-    // __syncthreads();
-    // uint32_t tmem_base = *sTmemBase;
-    //
-    // // Zero TMEM
-    // if (wid == 0) {
-    //     for (int col = 0; col < 128; col++) {
-    //         tmem_store(tmem_base + col, 0, 0, 0, 0);
-    //     }
-    //     tmem_fence_store();
-    // }
-    // __syncthreads();
+    if (wid == 0) {
+        uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
+        tmem_alloc(smem_ptr, 128);
+    }
+    __syncthreads();
+    uint32_t tmem_base = *sTmemBase;
+
+    // Zero TMEM
+    if (wid == 0) {
+        for (int col = 0; col < 128; col++) {
+            tmem_store(tmem_base + col, 0, 0, 0, 0);
+        }
+        tmem_fence_store();
+    }
+    __syncthreads();
 
     // ================================================================
     // Load Q and K into SMEM in canonical layout