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

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
2026-08-28 02:31:26 -07:00
parent 9cee0a31d1
commit 74df026877
1 changed files with 6 additions and 1 deletions
@@ -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,