From 5adc2880f99601ecf3b3f3f62dfca9faaca7ed6d Mon Sep 17 00:00:00 2001 From: Zaili Wang <109502517+ZailiWang@users.noreply.github.com> Date: Thu, 27 Aug 2026 10:13:10 +0800 Subject: [PATCH] [CPU] Enable ERNIE models on CPU (#35222) --- python/sglang/srt/layers/moe/topk.py | 2 ++ python/sglang/srt/models/ernie4.py | 7 ++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 8bfbfcd8f..0eb9ce18b 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -782,6 +782,8 @@ def fused_topk_cpu( raise ValueError("packed_out is not supported for CPU fused topk") if num_token_non_padded is not None: raise ValueError("num_token_non_padded is not supported for CPU fused topk") + if correction_bias is not None and correction_bias.dtype != torch.float32: + correction_bias = correction_bias.to(torch.float32) if scoring_func == "softmax": topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu( diff --git a/python/sglang/srt/models/ernie4.py b/python/sglang/srt/models/ernie4.py index c3be1dde1..33cec5ea5 100644 --- a/python/sglang/srt/models/ernie4.py +++ b/python/sglang/srt/models/ernie4.py @@ -42,9 +42,10 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4MLP from sglang.srt.models.llama import LlamaAttention as Ernie4Attention from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix, is_npu, make_layers +from sglang.srt.utils import add_prefix, is_cpu, is_npu, make_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config +_is_cpu = is_cpu() _is_npu = is_npu() @@ -89,8 +90,8 @@ class Ernie4Moe(nn.Module): self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix)) correction_bias = self.gate.e_score_correction_bias - # npu only supports 1D, but current correction_bias is 2D - if _is_npu: + # npu/cpu only support 1D, but current correction_bias is 2D + if _is_npu or _is_cpu: correction_bias = correction_bias.squeeze(0) self.topk = TopK( top_k=config.moe_k,