diff --git a/tests/unit/test_fmha_ts_full.cu b/tests/unit/test_fmha_ts_full.cu index bd2530bc..77c9888a 100644 --- a/tests/unit/test_fmha_ts_full.cu +++ b/tests/unit/test_fmha_ts_full.cu @@ -134,7 +134,7 @@ test_fmha_ts_full(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, { uint32_t idesc_pv = make_idesc(BLOCK_MN, HD); - for (int kt = 0; kt < NKT_PV; kt++) { + for (int kt = 0; kt < 1; kt++) { // DEBUG: single PV K-tile uint32_t tmem_a = tb + kt * MMA_K_BF16; // A from P's kt-th 16 columns bf16_t* sv = sV + kt * 256; uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sv), 16);