diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu
index bc7d9349..51249a81 100644
--- a/tests/unit/test_umma_qk.cu
+++ b/tests/unit/test_umma_qk.cu
@@ -59,12 +59,12 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k,
     uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 32);
     uint32_t idesc = make_idesc(128, 128);
 
-    // MMA — 4 warp leaders
-    if (lane == 0 && wid < 4) {
+    // MMA — 1 thread calls (4× scaling is expected for M=128 cta_group::1)
+    if (tid == 0) {
         umma_ss_f16(tb, desc_q, desc_k, idesc, false);
     }
     __syncwarp();
-    if (wid < 4 && lane == 0) tmem_fence_store();
+    if (wid == 0 && lane == 0) tmem_fence_store();
     __syncthreads();
 
     // Read from TMEM using Layout D: 32x32b.x8 format
@@ -83,7 +83,7 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k,
         // Lane 0 gets S[wid*32 + 0, 0..7] in tmp[0..7]
         if (lane == 0) {
             for (int c = 0; c < 8; c++) {
-                s_out[wid * 8 + c] = tmp[c];
+                s_out[wid * 8 + c] = tmp[c] * 0.25f;  // Divide by 4 (Layout D scaling)
             }
         }
     }