From 53f3634806262a4c0f3298b746d019290eeb168e Mon Sep 17 00:00:00 2001 From: carlushuang Date: Thu, 25 Sep 2025 16:34:44 +0800 Subject: [PATCH 1/2] hot fix check eid range --- .../ops/fused_moe/kernel/moe_sorting_kernel.hpp | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp index faeb5cf6b3b..b91a8e56412 100644 --- a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp +++ b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp @@ -13,7 +13,7 @@ namespace ck_tile { #define MOE_SORTING_MOCK_ID(token_id_, topk_id_) \ - static_cast(((token_id_) & 0x00ffffff) | (((topk_id_) & 0xff) << 24)) + static_cast(((token_id_)&0x00ffffff) | (((topk_id_)&0xff) << 24)) #ifndef MOE_SORTING_USE_EX_KERNEL #define MOE_SORTING_USE_EX_KERNEL 1 @@ -1574,6 +1574,7 @@ struct MoeSortingMultiPhaseKernel_P0 void* p_expert_mesh; // [expert, tokens] index_t tokens; // if p_local_tokens is not nullptr, this indicate the max possible tokens // used for ws/LDS calculation + index_t num_experts; index_t mesh_stride; // mesh_stride for p_expert_mesh mdiv topk_mdiv; }; @@ -1597,6 +1598,7 @@ struct MoeSortingMultiPhaseKernel_P0 k.p_local_tokens = h.p_local_tokens; k.p_expert_mesh = h.p_ws; k.tokens = h.tokens; + k.num_experts = h.num_experts; k.mesh_stride = impl::moe_sorting_mp_mesh_stride(h.tokens); k.topk_mdiv = mdiv{static_cast(h.topk)}; return k; @@ -1655,14 +1657,18 @@ struct MoeSortingMultiPhaseKernel_P0 IndexType eid = x[j.value]; // ext_vector_type must use int to [] uint32_t curr_token_id, curr_topk_id; kargs.topk_mdiv.divmod(i * Problem::SubTokenTile + j, curr_token_id, curr_topk_id); - if constexpr(Problem::LocalToken) + if(eid < kargs.num_experts) { - if(static_cast(curr_token_id) < tokens) + if constexpr(Problem::LocalToken) + { + if(static_cast(curr_token_id) < tokens) + p_expert_mesh[eid * mesh_stride + curr_token_id] = + (curr_topk_id + 1) & 0xffff; + } + else p_expert_mesh[eid * mesh_stride + curr_token_id] = (curr_topk_id + 1) & 0xffff; } - else - p_expert_mesh[eid * mesh_stride + curr_token_id] = (curr_topk_id + 1) & 0xffff; }); } } From 655ac1a851438815bc8e46378436151cf3becc0b Mon Sep 17 00:00:00 2001 From: illsilin_amdeng Date: Fri, 26 Sep 2025 13:36:21 -0700 Subject: [PATCH 2/2] fix clang format --- include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp index 7f40be844dd..42e2fad2360 100644 --- a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp +++ b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp @@ -13,7 +13,7 @@ namespace ck_tile { #define MOE_SORTING_MOCK_ID(token_id_, topk_id_) \ - static_cast(((token_id_)&0x00ffffff) | (((topk_id_)&0xff) << 24)) + static_cast(((token_id_) & 0x00ffffff) | (((topk_id_) & 0xff) << 24)) #ifndef MOE_SORTING_USE_EX_KERNEL #define MOE_SORTING_USE_EX_KERNEL 1