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