From 74df026877a490720739749c825b8de3a8423dd5 Mon Sep 17 00:00:00 2001 From: Xinguo Zhu Date: Fri, 28 Aug 2026 17:31:26 +0800 Subject: [PATCH] fix(cpu): skip GPU JIT MoE top-k on CPU (#35677) Co-authored-by: Ma Mingfei --- python/sglang/srt/layers/moe/topk.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 0eb9ce18b..6e047ac88 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -2264,7 +2264,12 @@ def select_experts( if scoring_func not in ("sqrtsoftplus", "sigmoid"): 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 topk_weights, topk_ids = _biased_topk( hidden_states=hidden_states,