diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index f1cc666f..fc629e7b 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -225,7 +225,7 @@ class FmhaV3StageCMulti: qp.tail() kvp.reset() kv_coord = n_kv_tiles - n_kv_tiles # SSA runtime zero - for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1): + for kt in range(n_kv_tiles): kvh = kvp.acquire_and_advance() cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)