diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index fef8818a9dba..c0b79f4eb4c4 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -170,10 +170,7 @@ def __init__( self.expert_bias = None def forward(self, hidden_states): - logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to( - hidden_states.dtype - ) - return logits + return F.linear(hidden_states.to(self.weight.dtype), self.weight, None) class BailingMoESparseMoeBlock(nn.Module): diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index ce8f8770589e..353d880ae55d 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -224,10 +224,7 @@ def __init__( self.expert_bias = None def forward(self, hidden_states): - logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to( - hidden_states.dtype - ) - return logits + return F.linear(hidden_states.to(self.weight.dtype), self.weight, None) class BailingMoE(nn.Module): diff --git a/python/sglang/srt/models/ernie4.py b/python/sglang/srt/models/ernie4.py index 5d30b1e77eba..6ddf450e5d2e 100644 --- a/python/sglang/srt/models/ernie4.py +++ b/python/sglang/srt/models/ernie4.py @@ -60,7 +60,7 @@ def __init__( torch.empty((config.moe_num_experts, config.hidden_size)) ) self.e_score_correction_bias = nn.Parameter( - torch.empty((1, config.moe_num_experts)) + torch.empty((1, config.moe_num_experts), dtype=torch.float32) ) def forward(self, hidden_states): diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index b0feaa7998d6..27b33df580f3 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -110,7 +110,9 @@ def __init__( prefix=f"{prefix}.gate", ) - self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(num_experts, dtype=torch.float32) + ) self.experts = get_moe_impl_class(quant_config)( num_experts=config.n_routed_experts, diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index f6f40c22f366..801576a70085 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -53,10 +53,7 @@ resolve_linear_parallel_group, ) from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.moe import ( - get_moe_a2a_backend, - get_moe_runner_backend, -) +from sglang.srt.layers.moe import get_moe_a2a_backend from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat from sglang.srt.layers.moe.utils import is_deepep_class_backend @@ -360,15 +357,8 @@ def __init__( torch.empty((config.n_routed_experts, config.hidden_size), dtype=self.dtype) ) if config.topk_method == "noaux_tc": - correction_bias_dtype = ( - torch.bfloat16 - if quant_config is not None - and quant_config.get_name() == "modelopt_fp4" - and get_moe_runner_backend().is_flashinfer_trtllm() - else torch.float32 - ) self.e_score_correction_bias = nn.Parameter( - torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) + torch.empty((config.n_routed_experts), dtype=torch.float32) ) else: self.e_score_correction_bias = None