ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-07 07:37:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

remove split num padd

Browse Source
This commit is contained in:
ltqin
2025-08-02 08:26:39 +00:00
parent 7ce108b5ab
commit 0708d47b34
2 changed files with 3 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp
View File

@@ -366,7 +366,7 @@ struct FmhaFwdSplitKVCombineKernel
o_acc_dram_naive,
make_tuple(
number<kNumWarps>{}, number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN1>{}),
sequence<true, kPadSeqLenQ, kPadHeadDimV>{});
sequence<false, kPadSeqLenQ, kPadHeadDimV>{});
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<FmhaPipeline::kMaxSplits>{}, number<FmhaPipeline::kM0>{}),
{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<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN1>{}),
{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]() {

4
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp
View File

@@ -282,12 +282,10 @@ struct BlockFmhaFwdSplitKVCombinePipeline
auto o_acc = make_static_distributed_tensor<OaccDataType>(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);
{