Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[CK_TILE] Shrink lse_accum distributed tensor size in fmha splitkv combine kernel #1577

Draft
wants to merge 14 commits into
base: develop
Choose a base branch
from
+69 −32
Draft

[CK_TILE] Shrink lse_accum distributed tensor size in fmha splitkv combine kernel #1577

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
d68ceca
Use smaller width for lse_accum dist tensor
poyenc Oct 16, 2024
f92661a
Update pipeline comment
poyenc Oct 16, 2024
253c9ba
Fix wrong distribution for lse_accum
poyenc Oct 16, 2024
7ad674e
Merge branch 'develop' into feature/optimize-splitkv-combine-kernel
poyenc Oct 16, 2024
3e15078
Remove duplicate dim in lse_accum dist encoding
poyenc Oct 16, 2024
81cb4c4
Decide fmha splitkv combine kernel kBlockSize by kM0
poyenc Oct 16, 2024
d878c46
Merge branch 'feature/optimize-splitkv-combine-kernel' of github.com:…
poyenc Oct 16, 2024
9df2983
Remove assumption of MPerThread=1
poyenc Oct 16, 2024
af50a11
Merge branch 'develop' into feature/optimize-splitkv-combine-kernel
poyenc Oct 16, 2024
f262bec
Add log<4> & log<8> specialization
poyenc Oct 17, 2024
f03ee03
Merge branch 'feature/optimize-splitkv-combine-kernel' of github.com:…
poyenc Oct 17, 2024
88e5de4
Enlarge occupancy array
poyenc Oct 17, 2024
598ee5c
Fix vector size for small tile
poyenc Oct 17, 2024
5776e16
Add support for kMaxSplits=8
poyenc Oct 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,9 +138,8 @@ struct BlockFmhaFwdSplitKVCombinePipeline
auto lse_accum = make_static_distributed_tensor<LSEDataType>(
Policy::template MakeLSEaccRegTileDistribution<Problem>());

// copy LDS (shape=[kM0, kMaxSplits]) to lse_accum (shape=[kM0, max(kMaxSplits, warp_size)])
// this will extend the distributed tensor width so that each thread in wave have data to
// reduce.
// copy LDS (shape=[kM0, kMaxSplits]) to lse_accum (shape=[kM0, kMaxSplits])
// and fill up -INF values outside the [kM0, num_splits] region.
{
constexpr auto spans = decltype(lse_accum)::get_distributed_spans();
sweep_tile_span(spans[number<0>{}], [&](auto idx0) {
Expand Down
20 changes: 12 additions & 8 deletions include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline_default_policy.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,26 +134,30 @@ struct BlockFmhaFwdSplitKVCombinePipelineDefaultPolicy
{
constexpr index_t kBlockSize = Problem::kBlockSize;

constexpr index_t kNPerBlock = max(Problem::kMaxSplits, get_warp_size());
constexpr index_t kNPerBlock = Problem::kMaxSplits;
constexpr index_t kMPerBlock = Problem::kM0;

constexpr index_t NThreads = get_warp_size();
constexpr index_t NThreads = 4;
constexpr index_t NPerThread = kNPerBlock / NThreads;

constexpr index_t MThreads = kBlockSize / NThreads;
constexpr index_t MPerThread = kMPerBlock / MThreads;
constexpr index_t MThreadPerWarp = get_warp_size() / NThreads;
constexpr index_t MWarps = kMPerBlock / MThreadPerWarp;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MWarps = kMPerBlock / MThreadPerWarp ? Confusing, why not MWarps = kBlockSize / get_warp_size()

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

because I'm using hard-coded kBlockSize=256, thus kBlockSize/ get_warp_size() (=4) will exceed the real # warp needed here if kMPerBlock is less than 64.

constexpr index_t MThreads = MThreadPerWarp * MWarps;
constexpr index_t MPerThread = kMPerBlock / MThreads;

static_assert(kBlockSize % MWarps == 0);
static_assert(NThreads * NPerThread == kNPerBlock);
static_assert(MThreads * MPerThread == kMPerBlock);

// duplicate MWarps if less than (kBlockSize / get_warp_size())
return make_static_tile_distribution(
tile_distribution_encoding<
sequence<1>,
tuple<sequence<MThreads, MPerThread>, sequence<NThreads, NPerThread>>,
tuple<sequence<1>, sequence<2>>,
tuple<sequence<0>, sequence<0>>,
tuple<sequence<MWarps, MThreadPerWarp, MPerThread>, sequence<NThreads, NPerThread>>,
tuple<sequence<0, 1>, sequence<2, 1>>,
tuple<sequence<0, 0>, sequence<0, 1>>,
sequence<1, 2>,
sequence<1, 1>>{});
sequence<2, 1>>{});
}

template <typename Problem>
Expand Down