diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index 2dd9c799..a75de0d5 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -69,9 +69,9 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k, uint64_t desc_q = make_umma_desc_kmajor_none(sQ_smem, 128); uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 128); - // Try different idesc values to understand the 4× factor - // Test 1: M=128, N=128 - uint32_t idesc = make_idesc(128, 128); + // Test with different N values to understand scaling + // N=8 → n_dim=1 + uint32_t idesc = make_idesc(128, 8); // Try N=8 if (tid == 0) { memcpy(&s_out[128], &desc_q, 8);