diff --git a/tests/unit/test_tmem_layout_pv64.cu b/tests/unit/test_tmem_layout_pv64.cu
index e91de980..75b0a4c7 100644
--- a/tests/unit/test_tmem_layout_pv64.cu
+++ b/tests/unit/test_tmem_layout_pv64.cu
@@ -16,10 +16,11 @@
 using namespace dsv4::kernels::attention;
 
 static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
+static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; }
 
 constexpr int HD = 64, SK = 128, BLOCK_MN = 128;
-constexpr int MMA_K_BF16 = 16;
-constexpr int TILE_SZ = BLOCK_MN * MMA_K_BF16;
+constexpr int LOCAL_MMA_K = 16;
+constexpr int TILE_SZ = BLOCK_MN * LOCAL_MMA_K;
 constexpr int V_TILE_SZ = (HD / 8) * 2 * 64;  // 1024
 
 __global__ void __launch_bounds__(128)
@@ -40,19 +41,19 @@ test_tmem_layout_pv(const bf16_t* q, const bf16_t* k, const bf16_t* v,
     for (int kt = 0; kt < 4; kt++) {
         bf16_t* sq = sQ0 + kt * TILE_SZ;
         for (int i = tid; i < TILE_SZ; i += 128) sq[i] = 0;
-        for (int d = tid; d < MMA_K_BF16; d += 128) {
+        for (int d = tid; d < LOCAL_MMA_K; d += 128) {
             int ck = d / 8, lc = d % 8;
-            sq[ck * 16 * 64 + lc] = q[kt * MMA_K_BF16 + d];
+            sq[ck * 16 * 64 + lc] = q[kt * LOCAL_MMA_K + d];
         }
     }
     for (int kt = 0; kt < 4; kt++) {
         bf16_t* sk = sK0 + kt * TILE_SZ;
         for (int i = tid; i < TILE_SZ; i += 128) sk[i] = 0;
         for (int r = 0; r < SK; r++) {
-            for (int d = tid; d < MMA_K_BF16; d += 128) {
+            for (int d = tid; d < LOCAL_MMA_K; d += 128) {
                 int ck = d / 8, lc = d % 8;
                 int tmn = r / 8, lr = r % 8;
-                sk[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * MMA_K_BF16 + d];
+                sk[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * LOCAL_MMA_K + d];
             }
         }
     }
@@ -60,8 +61,8 @@ test_tmem_layout_pv(const bf16_t* q, const bf16_t* k, const bf16_t* v,
         bf16_t* sv = sV + kt * V_TILE_SZ;
         for (int i = tid; i < V_TILE_SZ; i += 128) sv[i] = 0;
         for (int d = tid; d < HD; d += 128) {
-            for (int lr = 0; lr < MMA_K_BF16; lr++) {
-                int r = kt * MMA_K_BF16 + lr;
+            for (int lr = 0; lr < LOCAL_MMA_K; lr++) {
+                int r = kt * LOCAL_MMA_K + lr;
                 int g_mn = d / 8, g_k = lr / 8;
                 int llr = d % 8, lc = lr % 8;
                 sv[g_k * 8 * 64 + g_mn * 64 + llr * 8 + lc] = v[d * SK + r];
@@ -131,7 +132,7 @@ test_tmem_layout_pv(const bf16_t* q, const bf16_t* k, const bf16_t* v,
             if (tid < 16) {
                 int c = tid;
                 int ck = c / 8, lc = c % 8;
-                sPk[ck * 16 * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[kt * MMA_K_BF16 + c]);
+                sPk[ck * 16 * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[kt * LOCAL_MMA_K + c]);
             }
             __syncthreads();