From a944f90040abfc211263f34f2fb2fd8682890202 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 13:44:36 +0000 Subject: [PATCH] test: match isolated TS test exactly (V=all-1, BLOCK_MN=16) --- tests/unit/test_fmha_ts_hd16.cu | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/tests/unit/test_fmha_ts_hd16.cu b/tests/unit/test_fmha_ts_hd16.cu index 4566a045..5dc8331c 100644 --- a/tests/unit/test_fmha_ts_hd16.cu +++ b/tests/unit/test_fmha_ts_hd16.cu @@ -57,16 +57,14 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v, bf16_t* sQ_pad = sQ + 128 * 16; for (int i = tid; i < 4096; i += 128) sQ_pad[i] = 0; - // Load V as 8 K-tiles of (16, 16) canonical - for (int i = tid; i < VKT * V_TILE_SZ; i += 128) sV_base[i] = 0; - for (int kt = 0; kt < VKT; kt++) { - bf16_t* sv = sV_base + kt * V_TILE_SZ; - for (int i = tid; i < MMA_K_BF16 * HD; i += 128) { - int r = i / HD, d = i % HD; - int ck = d / 8, lc = d % 8; - int tmn = r / 8, lr = r % 8; - sv[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = v[d * SK + kt * MMA_K_BF16 + r]; - } + // Load V = all 1.0 into (16, 16) canonical (same as isolated test) + for (int i = tid; i < 16 * 16; i += 128) sV_base[i] = 0; + __syncthreads(); + for (int i = tid; i < 16 * 16; i += 128) { + int r = i / 16, c = i % 16; + int ck = c / 8, lc = c % 8; + int tmn = r / 8, lr = r % 8; + sV_base[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = f32_to_bf16(1.0f); } __syncthreads(); @@ -138,9 +136,8 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v, uint32_t idesc_pv = make_idesc(BLOCK_MN, HD); for (int kt = 0; kt < 1; kt++) { // DEBUG: only first PV K-tile - bf16_t* sv = sV_base + kt * V_TILE_SZ; - uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sv), MMA_K_BF16); - uint32_t tmem_a = tb + kt * MMA_K_BF16; // P's K-tile columns [16*kt, 16*kt+15] + uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV_base), 16); + uint32_t tmem_a = tb; // P's first 16 columns if (tid == 0) umma_ts_f16(tb_o, tmem_a, dv, idesc_pv, false); // no accumulate for first tile asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");