test: 4 warp leaders (lane==0) call MMA simultaneously

2026-05-28 11:30:19 +00:00
parent c01d6fddf4
commit f5a0966afc
1 changed files with 3 additions and 2 deletions
--- 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)