diff --git a/tests/unit/test_qk_minimal.cu b/tests/unit/test_qk_minimal.cu
index 646e26b7..09b84f94 100644
--- a/tests/unit/test_qk_minimal.cu
+++ b/tests/unit/test_qk_minimal.cu
@@ -91,7 +91,7 @@ test_qk_minimal_kernel(float* __restrict__ out_s,
     __syncthreads();
 
     // Read from TMEM
-    if (wid == 0) {
+    if (lane < T) {
         for (int n = 0; n < NUM_READS; n++) {
             float tmp[8];
             asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
@@ -108,7 +108,7 @@ test_qk_minimal_kernel(float* __restrict__ out_s,
         }
     }
     __syncthreads();
-    if (is_mma_warp) tmem_dealloc(tb, TMEM_N);
+    tmem_dealloc(tb, TMEM_N);
 }
 
 int main() {