diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp index 1543568eff..c5a1a19e2c 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp @@ -366,7 +366,7 @@ struct FmhaFwdSplitKVCombineKernel o_acc_dram_naive, make_tuple( number{}, number{}, number{}), - sequence{}); + sequence{}); const index_t padded_num_splits = o_acc_dram_view.get_tensor_descriptor().get_lengths()[number<0>{}]; @@ -404,13 +404,10 @@ struct FmhaFwdSplitKVCombineKernel make_tuple(number{}, number{}), {0, i_m0}); - const index_t padded_num_splits = - integer_divide_ceil(kargs.num_splits, kNumWarps) * kNumWarps; - auto o_acc_dram_window = make_tile_window(o_acc_dram, make_tuple(number{}, number{}), - {i_tile_m * padded_num_splits * FmhaPipeline::kM0, i_n1}); + {i_tile_m * kargs.num_splits * FmhaPipeline::kM0, i_n1}); // LSE DRAM window auto lse_dram_window = [&, i_nhead_ = i_nhead]() { diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp index d119a8c074..1527d04104 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp @@ -282,12 +282,10 @@ struct BlockFmhaFwdSplitKVCombinePipeline auto o_acc = make_static_distributed_tensor(o_acc_dist); clear_tile(o_acc); - const index_t padded_num_splits = integer_divide_ceil(num_splits, kNumWarps) * kNumWarps; - __builtin_amdgcn_sched_barrier(0); block_sync_lds(); // each warp handles a [KM0, kN1] tile - for(index_t split_start = 0; split_start < padded_num_splits; split_start++) + for(index_t split_start = 0; split_start < num_splits; split_start++) { auto o_tile = load_tile(o_acc_dram_window); {