fix: TMEM read/dealloc for 128-thread kernel

2026-05-29 18:40:24 +00:00
parent bd1309ba88
commit c755e6fdde
1 changed files with 2 additions and 2 deletions
--- 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() {