From 9fc2d549e4f14cfd40ff3f5511c51a0e90f8c2a1 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Fri, 29 May 2026 18:42:03 +0000
Subject: [PATCH] fix: warp-collective TMEM read/dealloc in minimal QK test

---
 tests/unit/test_qk_minimal.cu | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/unit/test_qk_minimal.cu b/tests/unit/test_qk_minimal.cu
index 09b84f94..9df0b81f 100644
--- a/tests/unit/test_qk_minimal.cu
+++ b/tests/unit/test_qk_minimal.cu
@@ -90,8 +90,9 @@ test_qk_minimal_kernel(float* __restrict__ out_s,
     asm volatile("fence.sc.gpu;" ::: "memory");
     __syncthreads();
 
-    // Read from TMEM
-    if (lane < T) {
+    // TMEM read — ALL 32 lanes must participate (warp-collective)
+    // For T=1, only lane 0's data is meaningful
+    {
         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 +109,8 @@ test_qk_minimal_kernel(float* __restrict__ out_s,
         }
     }
     __syncthreads();
-    tmem_dealloc(tb, TMEM_N);
+    // TMEM dealloc — warp 0 does it (must be warp-collective)
+    if (tid < 32) tmem_dealloc(tb, TMEM_N);
 }
 
 int main() {