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

fix(cpu): skip GPU JIT MoE top-k on CPU (#35677)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Xinguo Zhu
2026-08-28 02:31:26 -07:00
committed by GitHub
co-authored by Ma Mingfei
parent 9cee0a31d1
commit 74df026877
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/topk.py
+6 -1
View File
@@ -2264,7 +2264,12 @@ def select_experts(
if scoring_func not in ("sqrtsoftplus", "sigmoid"): if scoring_func not in ("sqrtsoftplus", "sigmoid"):
assert not apply_routed_scaling_factor_on_output, "Not implemented" assert not apply_routed_scaling_factor_on_output, "Not implemented"
if scoring_func == "sqrtsoftplus" or scoring_func == "sigmoid": # The JIT route depends on GPU-only topk_sigmoid/topk_softmax imports
_can_use_jit_kernel = not _is_cpu
if _can_use_jit_kernel and (
scoring_func == "sqrtsoftplus" or scoring_func == "sigmoid"
):
_biased_topk = biased_topk_xpu if _is_xpu else biased_topk_jit_kernel_impl _biased_topk = biased_topk_xpu if _is_xpu else biased_topk_jit_kernel_impl
topk_weights, topk_ids = _biased_topk( topk_weights, topk_ids = _biased_topk(
hidden_states=hidden_states, hidden_states=hidden_states,