diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index 618796fa..bc7d9349 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -59,9 +59,7 @@ 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 call the instruction (Layout D requires 4 warps) - // elect_one_sync selects 1 leader per warp. With 4 warps, 4 leaders call MMA. - int elect_one = __ballot_sync(0xFFFFFFFF, lane == 0); + // MMA — 4 warp leaders if (lane == 0 && wid < 4) { umma_ss_f16(tb, desc_q, desc_k, idesc, false); }