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

fix sorting bug

Browse Source
This commit is contained in:
coderfeli
2025-02-21 06:55:20 +00:00
parent 4cab3b5fb2
commit d5b2c900b9
1 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -574,11 +574,10 @@ struct MoeSortingKernel
int e_start = cumsum[tid];
int e_end = cumsum[tid + 1];
index_t *p_sorted_expert_cnts = p_total_tokens_post_pad + 1 + tid;
int e_size = unit_size_mdiv.div(e_end);
p_sorted_expert_cnts[1] = e_size;
if (tid == 0) {
p_sorted_expert_cnts[0] = 0;
}
p_sorted_expert_cnts[1] = unit_size_mdiv.div(e_end);
for(int i = e_start; i < e_end; i += unit_size_mdiv.divisor)
{
p_sorted_expert_ids[unit_size_mdiv.div(i)] = tid;
@@ -870,11 +869,19 @@ struct MoeSortingKernel
}
__syncthreads();
}
if (tid == 0) {
//temp hack ptr for expert tile cnt
p_total_tokens_post_pad[1] = 0;
}
for(int i_e = tid; i_e < num_experts; i_e += block_size)
{
int e_start = smem_cumsum(i_e);
int e_end = smem_cumsum(i_e + 1);
//temp hack ptr for expert tile cnt
index_t *p_sorted_expert_cnts = p_total_tokens_post_pad + 1 + i_e;
p_sorted_expert_cnts[1] = unit_size_mdiv.div(e_end);
int expert_id = [&]() {
if constexpr(Problem::LocalExpertMasking)
{
@@ -884,7 +891,6 @@ struct MoeSortingKernel
else
return i_e;
}();
smem_cumdup(i_e) = e_start; // duplicate cumsum for later use
if constexpr(Problem::SkipExpertsWithZeroTokens)
{
@@ -897,12 +903,6 @@ struct MoeSortingKernel
if(local_expert_mask[i_e] == 0)
continue;
}
index_t *p_sorted_expert_cnts = p_total_tokens_post_pad + 1 + tid;
int e_size = unit_size_mdiv.div(e_end);
p_sorted_expert_cnts[1] = e_size;
if (tid == 0) {
p_sorted_expert_cnts[0] = 0;
}
for(int i = e_start; i < e_end; i += unit_size_mdiv.divisor)
{
p_sorted_expert_ids[unit_size_mdiv.div(i)] = expert_id;