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,