diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index ac403717..556beb0e 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -71,14 +71,14 @@ test_umma_qk_hd16(
     __syncthreads();
     uint32_t tmem_base = *sTmemBase;
 
-    // SKIP zeroing TMEM — just test alloc
-    // if (wid == 0) {
-    //     for (int col = 0; col < 128; col++) {
-    //         tmem_store(tmem_base + col, 0, 0, 0, 0);
-    //     }
-    //     tmem_fence_store();
-    // }
-    // __syncthreads();
+    // 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
@@ -149,10 +149,10 @@ test_umma_qk_hd16(
     }
     __syncthreads();
 
-    // TMEM dealloc — skipped
-    // if (wid == 0) {
-    //     tmem_dealloc(tmem_base, 128);
-    // }
+    // TMEM dealloc
+    if (wid == 0) {
+        tmem_dealloc(tmem_base, 128);
+    }
 }
 
 // ==================================================================