[2/2] Use moe_sum_reduce cuda kernel (#10654)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com>
This commit is contained in:
Yuan Luo
2025-10-28 12:01:57 +08:00
committed by GitHub
parent 729f612dc6
commit 813bd6f85c

View File

@@ -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(