diff --git a/tests/unit/test_fmha_ts_hd16.cu b/tests/unit/test_fmha_ts_hd16.cu index 070bd33c..5cdc6f5e 100644 --- a/tests/unit/test_fmha_ts_hd16.cu +++ b/tests/unit/test_fmha_ts_hd16.cu @@ -40,8 +40,9 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v, uint32_t* sTmemBase = (uint32_t*)sbuf; bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15); bf16_t* sK = sQ + 128 * 16 + 4096; - bf16_t* sV_base = sK + 128 * 16; - float* sQ_row = (float*)(sV_base + VKT * V_TILE_SZ); + // Fresh aligned V buffer (ensure 128-byte alignment for descriptor) + bf16_t* sV = (bf16_t*)(((uintptr_t)(sK + 128 * 16) + 127) & ~(uintptr_t)127); + float* sQ_row = (float*)(sV + 16 * 16); for (int d = tid; d < HD; d += 128) sQ_row[d] = bf16_to_f32(q[d]); @@ -58,13 +59,13 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v, for (int i = tid; i < 4096; i += 128) sQ_pad[i] = 0; // 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; + for (int i = tid; i < 16 * 16; i += 128) sV[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); + sV[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = f32_to_bf16(1.0f); } __syncthreads(); @@ -136,7 +137,7 @@ 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 - uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV_base), 16); + uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16); uint32_t tmem_a = tb; // P's first 16 columns if (tid == 0) umma_ts_f16(tb, tmem_a, dv, idesc_pv, false); // Write C to tb (overwriting P)