[CPU] Enable ERNIE models on CPU (#35222)
This commit is contained in:
@@ -782,6 +782,8 @@ def fused_topk_cpu(
|
|||||||
raise ValueError("packed_out is not supported for CPU fused topk")
|
raise ValueError("packed_out is not supported for CPU fused topk")
|
||||||
if num_token_non_padded is not None:
|
if num_token_non_padded is not None:
|
||||||
raise ValueError("num_token_non_padded is not supported for CPU fused topk")
|
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":
|
if scoring_func == "softmax":
|
||||||
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
|
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
|
||||||
|
|||||||
@@ -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.deepseek_v2 import DeepseekV2MLP as Ernie4MLP
|
||||||
from sglang.srt.models.llama import LlamaAttention as Ernie4Attention
|
from sglang.srt.models.llama import LlamaAttention as Ernie4Attention
|
||||||
from sglang.srt.runtime_context import get_parallel
|
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
|
from sglang.srt.utils.hf_transformers_utils import get_rope_config
|
||||||
|
|
||||||
|
_is_cpu = is_cpu()
|
||||||
_is_npu = is_npu()
|
_is_npu = is_npu()
|
||||||
|
|
||||||
|
|
||||||
@@ -89,8 +90,8 @@ class Ernie4Moe(nn.Module):
|
|||||||
self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix))
|
self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix))
|
||||||
|
|
||||||
correction_bias = self.gate.e_score_correction_bias
|
correction_bias = self.gate.e_score_correction_bias
|
||||||
# npu only supports 1D, but current correction_bias is 2D
|
# npu/cpu only support 1D, but current correction_bias is 2D
|
||||||
if _is_npu:
|
if _is_npu or _is_cpu:
|
||||||
correction_bias = correction_bias.squeeze(0)
|
correction_bias = correction_bias.squeeze(0)
|
||||||
self.topk = TopK(
|
self.topk = TopK(
|
||||||
top_k=config.moe_k,
|
top_k=config.moe_k,
|
||||||
|
|||||||
Reference in New Issue
Block a user