minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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
co-authored by luoyuan.luo huangtingwei
parent 729f612dc6
commit 813bd6f85c
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
+3 -2
View File
@@ -36,7 +36,7 @@ _is_cpu = is_cpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _is_cuda: 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: elif _is_cpu and _is_cpu_amx_available:
pass pass
elif _is_hip: elif _is_hip:
@@ -569,11 +569,12 @@ def fused_experts_impl(
routed_scaling_factor, routed_scaling_factor,
) )
else: else:
moe_sum_reduce_triton( moe_sum_reduce(
intermediate_cache3.view(*intermediate_cache3.shape), intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states[begin_chunk_idx:end_chunk_idx], out_hidden_states[begin_chunk_idx:end_chunk_idx],
routed_scaling_factor, routed_scaling_factor,
) )
elif _is_hip: elif _is_hip:
if _use_aiter: if _use_aiter:
moe_sum( moe_sum(