From 813bd6f85c3c3697ac868bc7e8471de072e617d3 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Tue, 28 Oct 2025 12:01:57 +0800 Subject: [PATCH] [2/2] Use moe_sum_reduce cuda kernel (#10654) Co-authored-by: luoyuan.luo Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> --- python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index 6d3fb53b0..54c06661c 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -36,7 +36,7 @@ _is_cpu = is_cpu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda: - from sgl_kernel import gelu_and_mul, silu_and_mul + from sgl_kernel import gelu_and_mul, moe_sum_reduce, silu_and_mul elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: @@ -569,11 +569,12 @@ def fused_experts_impl( routed_scaling_factor, ) else: - moe_sum_reduce_triton( + moe_sum_reduce( intermediate_cache3.view(*intermediate_cache3.shape), out_hidden_states[begin_chunk_idx:end_chunk_idx], routed_scaling_factor, ) + elif _is_hip: if _use_aiter: moe_sum(