From 97eb531f26278cdea774bcec9a4b5bdd0943bef6 Mon Sep 17 00:00:00 2001 From: charlieguo1106 Date: Tue, 13 Oct 2026 05:49:07 +0000 Subject: [PATCH 1/4] bug fix duplicated expert IDs --- aiter/fused_moe.py | 9 +-------- csrc/py_itfs_cu/moe_sorting_opus_kernels.cu | 11 +++++++++++ 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index e835868bdca..f710087d4c9 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -489,14 +489,7 @@ def _moe_sorting_impl( aux_m_indices = torch.empty( max_num_tokens_padded, dtype=dtypes.i32, device=device ) - # The sort writes local routes only; EP leaves remote routes at -1, - # which the scatter reduce skips. - if expert_mask is not None: - aux_reverse_sorted = torch.full( - (M * topk,), -1, dtype=dtypes.i32, device=device - ) - else: - aux_reverse_sorted = torch.empty(M * topk, dtype=dtypes.i32, device=device) + aux_reverse_sorted = torch.empty(M * topk, dtype=dtypes.i32, device=device) if use_opus: ws_size = aiter.moe_sorting_opus_get_workspace_size( diff --git a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu index 0e91f180070..a1e6453e477 100644 --- a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu +++ b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu @@ -51,6 +51,17 @@ void moe_sorting_opus_fwd(aiter_tensor_t& topk_ids, HipDeviceGuard device_guard(topk_ids.device_id); const hipStream_t stream = aiter::getCurrentHIPStream(); + if(reverse_sorted.has_value()) + { + auto& reverse = reverse_sorted.value(); + AITER_CHECK(reverse.dtype() == AITER_DTYPE_i32 && reverse.is_contiguous() && + reverse.numel() >= static_cast(num_tokens) * topk, + "reverse_sorted must be a contiguous int32 buffer for tokens * topk routes"); + HIP_CALL(hipMemsetAsync(reverse.data_ptr(), 0xff, + static_cast(num_tokens) * topk * sizeof(int32_t), + stream)); + } + void* ws_ptr = workspace.has_value() ? workspace.value().data_ptr() : nullptr; moe_sorting_opus( From ffce82f5e016619e6eaee8114a70d1adf42b2376 Mon Sep 17 00:00:00 2001 From: charlieguo1106 Date: Tue, 13 Oct 2026 07:13:04 +0000 Subject: [PATCH 2/4] [MoE] Fuse reverse route index initialization into Opus sorting kernels --- csrc/include/moe_sorting_opus.h | 34 +++++++++++++++++++++ csrc/py_itfs_cu/moe_sorting_opus_kernels.cu | 3 -- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/csrc/include/moe_sorting_opus.h b/csrc/include/moe_sorting_opus.h index ea2e7d99848..159c60b433b 100644 --- a/csrc/include/moe_sorting_opus.h +++ b/csrc/include/moe_sorting_opus.h @@ -1133,6 +1133,14 @@ struct MoeSortingKernel return; } + if(kargs.p_reverse_sorted) + { + opus::index_t* p_reverse_sorted = static_cast(kargs.p_reverse_sorted); + const opus::index_t routes = kargs.tokens * kargs.topk_mdiv.divisor; + for(opus::index_t i = threadIdx.x; i < routes; i += kBlockSize) + p_reverse_sorted[i] = -1; + } + extern __shared__ char smem[]; return moe_align_block_size_kernel_ex( @@ -1371,6 +1379,22 @@ OPUS_D void moe_buf_set_zero_kernel_2d(void* buf, } } +template +OPUS_D void moe_sorting_reverse_sorted_set_invalid(void* p_reverse_sorted, + opus::index_t tokens, + opus::index_t topk) +{ + if(p_reverse_sorted == nullptr) + return; + opus::index_t* p_buf = reinterpret_cast(p_reverse_sorted); + const opus::index_t routes = tokens * topk; + for(opus::index_t i = blockIdx.x * kBlockSize + threadIdx.x; i < routes; + i += gridDim.x * kBlockSize) + { + p_buf[i] = -1; + } +} + } // namespace impl // TODO: tokens could be from @@ -1629,6 +1653,7 @@ struct MoeSortingMultiPhaseKernel_P0_v1 const void* p_topk_ids; // [tokens, topk] const void* p_local_tokens; // [1], if not nullptr, use this as actual tokens void* p_expert_mesh; // [expert, tokens] + void* p_reverse_sorted; opus::index_t tokens; // if p_local_tokens is not nullptr, this indicate the max possible // tokens used for ws/LDS calculation opus::index_t num_experts; @@ -1654,6 +1679,7 @@ struct MoeSortingMultiPhaseKernel_P0_v1 k.p_topk_ids = h.p_topk_ids; k.p_local_tokens = h.p_local_tokens; k.p_expert_mesh = h.p_ws; + k.p_reverse_sorted = h.p_reverse_sorted; k.tokens = h.tokens; k.num_experts = h.num_experts; k.mesh_stride = impl::moe_sorting_mp_mesh_stride(h.tokens); @@ -1670,6 +1696,9 @@ struct MoeSortingMultiPhaseKernel_P0_v1 OPUS_D void operator()(Kargs kargs) const { + impl::moe_sorting_reverse_sorted_set_invalid( + kargs.p_reverse_sorted, kargs.tokens, kargs.topk_mdiv.divisor); + using topk_id_t = opus::vector_t; const topk_id_t* p_topk_ids = reinterpret_cast(kargs.p_topk_ids); @@ -1758,6 +1787,7 @@ struct MoeSortingMultiPhaseKernel_P0_v2 const void* p_local_expert_mask; // [expert] void* p_expert_cumsum; // [expert] + void* p_reverse_sorted; opus::index_t num_experts; }; @@ -1786,6 +1816,7 @@ struct MoeSortingMultiPhaseKernel_P0_v2 k.mesh_stride = impl::moe_sorting_mp_mesh_stride(h.tokens); k.topk_mdiv = opus::mdiv{static_cast(h.topk)}; k.p_local_expert_mask = h.p_local_expert_mask; + k.p_reverse_sorted = h.p_reverse_sorted; k.num_experts = h.num_experts; return k; } @@ -1802,6 +1833,9 @@ struct MoeSortingMultiPhaseKernel_P0_v2 OPUS_D void operator()(Kargs kargs) const { + impl::moe_sorting_reverse_sorted_set_invalid( + kargs.p_reverse_sorted, kargs.tokens, kargs.topk_mdiv.divisor); + constexpr opus::index_t index_pack = Problem::SubTokenTile; // always packed __shared__ char smem[GetSmemSize()]; using topk_id_t = opus::vector_t; diff --git a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu index a1e6453e477..c7335b2b90e 100644 --- a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu +++ b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu @@ -57,9 +57,6 @@ void moe_sorting_opus_fwd(aiter_tensor_t& topk_ids, AITER_CHECK(reverse.dtype() == AITER_DTYPE_i32 && reverse.is_contiguous() && reverse.numel() >= static_cast(num_tokens) * topk, "reverse_sorted must be a contiguous int32 buffer for tokens * topk routes"); - HIP_CALL(hipMemsetAsync(reverse.data_ptr(), 0xff, - static_cast(num_tokens) * topk * sizeof(int32_t), - stream)); } void* ws_ptr = workspace.has_value() ? workspace.value().data_ptr() : nullptr; From 7b8411ced2a167ba50cb666e2e023c64aafd291f Mon Sep 17 00:00:00 2001 From: charlieguo1106 Date: Tue, 13 Oct 2026 07:19:12 +0000 Subject: [PATCH 3/4] =?UTF-8?q?revert=20moe=5Fsorting=5Fopus=5Fkernels.cu?= =?UTF-8?q?=E2=80=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- csrc/py_itfs_cu/moe_sorting_opus_kernels.cu | 8 -------- 1 file changed, 8 deletions(-) diff --git a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu index c7335b2b90e..0e91f180070 100644 --- a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu +++ b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu @@ -51,14 +51,6 @@ void moe_sorting_opus_fwd(aiter_tensor_t& topk_ids, HipDeviceGuard device_guard(topk_ids.device_id); const hipStream_t stream = aiter::getCurrentHIPStream(); - if(reverse_sorted.has_value()) - { - auto& reverse = reverse_sorted.value(); - AITER_CHECK(reverse.dtype() == AITER_DTYPE_i32 && reverse.is_contiguous() && - reverse.numel() >= static_cast(num_tokens) * topk, - "reverse_sorted must be a contiguous int32 buffer for tokens * topk routes"); - } - void* ws_ptr = workspace.has_value() ? workspace.value().data_ptr() : nullptr; moe_sorting_opus( From c775637dae5c9756597f189020257292b8505d7f Mon Sep 17 00:00:00 2001 From: charlieguo1106 Date: Tue, 13 Oct 2026 07:33:35 +0000 Subject: [PATCH 4/4] [MoE] Check reverse_sorted holds tokens * topk int32 before the Opus sort seeds every route with -1 --- csrc/py_itfs_cu/moe_sorting_opus_kernels.cu | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu index 0e91f180070..c7335b2b90e 100644 --- a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu +++ b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu @@ -51,6 +51,14 @@ void moe_sorting_opus_fwd(aiter_tensor_t& topk_ids, HipDeviceGuard device_guard(topk_ids.device_id); const hipStream_t stream = aiter::getCurrentHIPStream(); + if(reverse_sorted.has_value()) + { + auto& reverse = reverse_sorted.value(); + AITER_CHECK(reverse.dtype() == AITER_DTYPE_i32 && reverse.is_contiguous() && + reverse.numel() >= static_cast(num_tokens) * topk, + "reverse_sorted must be a contiguous int32 buffer for tokens * topk routes"); + } + void* ws_ptr = workspace.has_value() ? workspace.value().data_ptr() : nullptr; moe_sorting_opus(