diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index d147ad39d..aa7f02274 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -38,6 +38,7 @@ 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 import Fp8Config from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -101,7 +102,18 @@ class DeepseekModelNextN(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) + if quant_config is not None and quant_config.get_name() == "quark": + self.eh_proj = ReplicatedLinear( + 2 * config.hidden_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("eh_proj", prefix), + ) + else: + self.eh_proj = nn.Linear( + 2 * config.hidden_size, config.hidden_size, bias=False + ) self.rot_weight = None if _is_npu: @@ -166,21 +178,23 @@ class DeepseekModelNextN(nn.Module): hidden_states = input_embeds if hidden_states.shape[0] > 0: - hidden_states = self.eh_proj( - torch.cat( - ( - self.enorm(hidden_states), - self.hnorm( - forward_batch.spec_info.hidden_states - if self.rot_weight is None - else torch.matmul( - forward_batch.spec_info.hidden_states, self.rot_weight - ) - ), + eh_input = torch.cat( + ( + self.enorm(hidden_states), + self.hnorm( + forward_batch.spec_info.hidden_states + if self.rot_weight is None + else torch.matmul( + forward_batch.spec_info.hidden_states, self.rot_weight + ) ), - dim=-1, - ) + ), + dim=-1, ) + if isinstance(self.eh_proj, ReplicatedLinear): + hidden_states, _ = self.eh_proj(eh_input) + else: + hidden_states = self.eh_proj(eh_input) if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp): hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) @@ -249,8 +263,20 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): self.cp_rank = None self.cp_size = None + nextn_quant_config = quant_config + # For quark, if the MTP layer is listed in exclude_layers, set quant_config to None. + if nextn_quant_config is not None and nextn_quant_config.get_name() == "quark": + from sglang.srt.layers.quantization.quark.utils import ( + should_ignore_layer, + ) + + ckpt_prefix = f"model.layers.{config.num_hidden_layers}" + mapped_prefix = self.hf_to_sglang_mapper._map_name(ckpt_prefix) + if should_ignore_layer(mapped_prefix, nextn_quant_config.exclude_layers): + nextn_quant_config = None + self.model = DeepseekModelNextN( - config, quant_config, prefix=add_prefix("model", prefix) + config, nextn_quant_config, prefix=add_prefix("model", prefix) ) self.lm_head = ParallelLMHead( config.vocab_size,