From a048b5688695e2b6683bf61a7c15054b132fe12a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 10:23:06 +0000 Subject: [PATCH] =?UTF-8?q?test:=20single-thread=20MMA=20+=200.25=20scalin?= =?UTF-8?q?g=20for=204=C3=97=20factor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit/test_umma_qk.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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) } } }