diff --git a/tests/unit/test_qk_direct.cu b/tests/unit/test_qk_direct.cu
index edc27718..b205e370 100644
--- a/tests/unit/test_qk_direct.cu
+++ b/tests/unit/test_qk_direct.cu
@@ -92,7 +92,7 @@ test_qk_direct_kernel(float* __restrict__ out_s,
     __syncthreads();
 
     const bool my_warp_active = (wid == 0);
-    const int my_row = lane;
+    const int my_row = threadIdx.x % 32;
     const bool my_row_active = my_row < T;
     constexpr int NUM_READS = SK / 8;