diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index ffd491a3..8334d8e0 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -58,8 +58,9 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k, uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 128); uint32_t idesc = make_idesc(128, 128); - // MMA — 1 thread issues (following gau-nernst pattern) - if (tid == 0) { + // MMA — 4 warp leaders call simultaneously (Layout D requires all 4 warps) + // Each warp's lane 0 issues the MMA instruction + if (lane == 0) { umma_ss_f16(tb, desc_q, desc_k, idesc, false); } // tcgen05.fence::after_thread_sync (CRITICAL — correct MMA→TMEM load fence)