From 0e4d9ddbd6b5146e6664899c93255777e01d758e Mon Sep 17 00:00:00 2001 From: LHXuuu Date: Fri, 30 Jan 2026 13:11:47 +0800 Subject: [PATCH] Fix the scenario where eh_proj is quantized in the bailing moe nextn weights (#17808) Signed-off-by: LHXuuu --- python/sglang/srt/models/bailing_moe_nextn.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/bailing_moe_nextn.py b/python/sglang/srt/models/bailing_moe_nextn.py index 437ba098c..bcfb75149 100644 --- a/python/sglang/srt/models/bailing_moe_nextn.py +++ b/python/sglang/srt/models/bailing_moe_nextn.py @@ -28,6 +28,7 @@ from transformers import PretrainedConfig from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.vocab_parallel_embedding import ( @@ -69,7 +70,13 @@ class BailingMoEModelNextN(nn.Module): self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + self.eh_proj = ReplicatedLinear( + 2 * config.hidden_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("eh_proj", prefix), + ) self.decoder = BailingMoEBlock( config, @@ -96,7 +103,7 @@ class BailingMoEModelNextN(nn.Module): hidden_states = input_embeds if hidden_states.shape[0] > 0: - hidden_states = self.eh_proj( + hidden_states, _ = self.eh_proj( torch.cat( ( self.enorm(hidden_states),