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

More consideration in MakeOaccDramTileDistribution() in splitkv_combine pipeline policy

Browse Source
This commit is contained in:
Qianfeng Zhang
2026-05-07 06:32:01 -04:00
parent e48bcff488
commit f73d7d2f8a
1 changed files with 31 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_combine_pipeline_policy.hpp
View File

@@ -26,22 +26,40 @@ struct HstuAttentionFwdSplitKVCombinePipelinePolicy
constexpr index_t KPerThread = KVector;
// try to assign more consecutive threads on dim-K
constexpr index_t KThreads = OtherK;
if constexpr(OtherK < get_warp_size())
{
// try to assign more consecutive threads on dim-K
constexpr index_t KThreads = OtherK;
static_assert(KThreads <= get_warp_size(), "Check failed!");
constexpr index_t MThreadPerWarp = get_warp_size() / KThreads;
constexpr index_t MPerThread = kMPerBlock / (MThreadPerWarp * NumWarps);
constexpr index_t MThreadPerWarp = get_warp_size() / KThreads;
constexpr index_t MPerThread = kMPerBlock / (MThreadPerWarp * NumWarps);
return make_static_tile_distribution(
tile_distribution_encoding<sequence<1>,
tuple<sequence<MPerThread, NumWarps, MThreadPerWarp>,
sequence<KThreads, KPerThread>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<1>, sequence<2, 0>>,
sequence<1, 2>,
sequence<0, 1>>{});
}
else
{
// all threads in the warp are assigned on dim-K
constexpr index_t KThreads = get_warp_size();
constexpr index_t KRepPerThread = OtherK / KThreads;
return make_static_tile_distribution(
tile_distribution_encoding<sequence<1>,
tuple<sequence<MPerThread, NumWarps, MThreadPerWarp>,
sequence<KThreads, KPerThread>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<1>, sequence<2, 0>>,
sequence<1, 2>,
sequence<0, 1>>{});
constexpr index_t MPerThread = kMPerBlock / NumWarps;
return make_static_tile_distribution(
tile_distribution_encoding<sequence<1>,
tuple<sequence<MPerThread, NumWarps>,
sequence<KRepPerThread, KThreads, KPerThread>>,
tuple<sequence<1>, sequence<2>>,
tuple<sequence<1>, sequence<1>>,
sequence<1, 2, 2>,
sequence<0, 0, 2>>{});
};
}
else // for kKPerBlock=96,160
{