ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-29 11:16:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

mock sorting multikernel

Browse Source
This commit is contained in:
coderfeli
2025-02-25 06:03:38 +00:00
parent d148949abd
commit f0e703f07e
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp
View File

@@ -1682,6 +1682,8 @@ struct MoeSortingMultiPhaseKernel_P2
if(position < kargs.num_experts)
{
index_t *p_sorted_expert_cnts = p_total_tokens_post_pad + 1 + position;//temp mock for p_sorted_expert_cnts, fixme:felix
p_sorted_expert_cnts[0] = out_0;
p_expert_cumsum[position] = out_0 * kargs.unit_size_mdiv.divisor;
}
@@ -1710,6 +1712,7 @@ struct MoeSortingMultiPhaseKernel_P2
{
auto total_tokens_post_pad = prev_cumsum_a * kargs.unit_size_mdiv.divisor;
p_total_tokens_post_pad[0] = total_tokens_post_pad;
p_total_tokens_post_pad[kargs.num_experts+1] = prev_cumsum_a; //temp mock for p_sorted_expert_cnts, fixme:felix
p_expert_cumsum[kargs.num_experts] = total_tokens_post_pad;
}
}