From c755e6fddea7f4a7ecab46409a1ae08a32b4f0e7 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Fri, 29 May 2026 18:40:24 +0000
Subject: [PATCH] fix: TMEM read/dealloc for 128-thread kernel

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

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() {