diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 3cd9cd661..ddefacd32 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -269,7 +269,10 @@ class ModelConfig: ): self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" - if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM": + if is_draft_model and self.hf_config.architectures[0] in [ + "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", + ]: self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN" if ( @@ -375,6 +378,7 @@ class ModelConfig: or "DeepseekV32ForCausalLM" in self.hf_config.architectures or "DeepseekV3ForCausalLM" in self.hf_config.architectures or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures + or "Glm4MoeLiteForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures or "DotsVLMForCausalLM" in self.hf_config.architectures @@ -394,15 +398,21 @@ class ModelConfig: else None ) - # Handle rope scaling with yarn - self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - if self.hf_config.rope_scaling: - mscale_all_dim = self.hf_config.rope_scaling.get( - "mscale_all_dim", False + if "Glm4MoeLiteForCausalLM" in self.hf_config.architectures: + self.scaling = 1 + self.hf_config.rope_scaling = None + else: + # Handle rope scaling with yarn + self.scaling = 1 / math.sqrt( + self.qk_nope_head_dim + self.qk_rope_head_dim ) - scaling_factor = self.hf_config.rope_scaling["factor"] - mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) - self.scaling = self.scaling * mscale * mscale + if self.hf_config.rope_scaling: + mscale_all_dim = self.hf_config.rope_scaling.get( + "mscale_all_dim", False + ) + scaling_factor = self.hf_config.rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale elif "MiniCPM3ForCausalLM" in self.hf_config.architectures: self.head_dim = 128 diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 1c23e8384..f7d79b7f8 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -454,6 +454,7 @@ class OpenAIServingChat(OpenAIServingBase): if request.chat_template_kwargs else {} ), + return_dict=False, ) except Exception as e: # If the first attempt fails, try transforming the tools format @@ -476,6 +477,7 @@ class OpenAIServingChat(OpenAIServingBase): if request.chat_template_kwargs else {} ), + return_dict=False, ) except jinja2.TemplateError as template_error: # Template errors (e.g., from raise_exception in Jinja templates) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 60b6031e9..6eff360c4 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -17,7 +17,7 @@ from sglang.srt.models.deepseek_common.utils import ( _use_aiter_gfx95, ) from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import BumpAllocator +from sglang.srt.utils import BumpAllocator, next_power_of_2 if TYPE_CHECKING: from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA @@ -472,7 +472,12 @@ class DeepseekMHAForwardMixin: ): k = k_nope.new_empty(*k_shape) concat_mla_k(k=k, k_nope=k_nope, k_rope=k_pe) - elif _is_cuda: + elif ( + _is_cuda + and next_power_of_2(self.num_local_heads) == self.num_local_heads + and next_power_of_2(self.qk_nope_head_dim) == self.qk_nope_head_dim + and next_power_of_2(self.qk_rope_head_dim) == self.qk_rope_head_dim + ): # fa3 mha support fp8 inputs if ( self.current_attention_backend == "fa3" diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 5c09bc681..190f01dc9 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -685,6 +685,8 @@ class Glm4MoeDecoderLayer(nn.Module): attention_bias = config.attention_bias self.layer_id = layer_id + use_qk_norm = config.use_qk_norm if hasattr(config, "use_qk_norm") else False + self.self_attn = Glm4MoeAttention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, @@ -699,7 +701,7 @@ class Glm4MoeDecoderLayer(nn.Module): attention_bias=attention_bias, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), - use_qk_norm=config.use_qk_norm, + use_qk_norm=use_qk_norm, alt_stream=alt_stream, ) diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py new file mode 100644 index 000000000..0490ac47e --- /dev/null +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -0,0 +1,808 @@ +# Copyright 2025-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Inference-only GLM-Lite model compatible with HuggingFace weights""" + +import logging +from typing import Iterable, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.batch_overlap.single_batch_overlap import SboFlags +from sglang.srt.distributed import ( + get_moe_expert_parallel_world_size, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import ( + get_attention_tp_size, + is_dp_attention_enabled, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import MergedColumnParallelLinear, RowParallelLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import ( + get_moe_a2a_backend, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.deepseek_v2 import ( + DeepseekV2AttentionMLA, + DeepseekV2DecoderLayer, + DeepseekV2ForCausalLM, + DeepseekV2Model, + DeepseekV2MoE, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import ( + BumpAllocator, + LazyValue, + add_prefix, + get_device_sm, + is_cuda, + log_info_on_rank0, + make_layers, +) + +_is_cuda = is_cuda() +_device_sm = get_device_sm() + +logger = logging.getLogger(__name__) + + +class Glm4MoeLiteMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + self.tp_size = tp_size + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward( + self, + x, + forward_batch=None, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + gemm_output_zero_allocator: BumpAllocator = None, + ): + # Keep parity with DeepseekV2MLP.forward signature since DeepseekV2DecoderLayer + # invokes MLP modules with these extra arguments. + if (self.tp_size == 1) and x.shape[0] == 0: + return x + + # Some quantization wrappers store the underlying parameter as `weight_packed`. + if not hasattr(self.gate_up_proj, "weight"): + self.gate_up_proj.weight = getattr(self.gate_up_proj, "weight_packed") + if not hasattr(self.down_proj, "weight"): + self.down_proj.weight = getattr(self.down_proj, "weight_packed") + + if ( + gemm_output_zero_allocator is not None + and x.shape[0] <= 256 + and self.gate_up_proj.weight.dtype == torch.uint8 + ): + y = gemm_output_zero_allocator.allocate( + x.shape[0] * self.gate_up_proj.output_size_per_partition + ).view(x.shape[0], self.gate_up_proj.output_size_per_partition) + x = (x, None, y) + + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj( + x, + skip_all_reduce=should_allreduce_fusion or use_reduce_scatter, + ) + return x + + +class Glm4MoeLiteGate(nn.Module): + def __init__( + self, + config, + prefix: str = "", + is_nextn: bool = False, + ): + super().__init__() + self.is_nextn = is_nextn + self.weight = nn.Parameter( + torch.empty((config.n_routed_experts, config.hidden_size)) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty((config.n_routed_experts), dtype=torch.float32) + ) + + def forward(self, hidden_states, gemm_output_zero_allocator: BumpAllocator = None): + # NOTE: For some unknown reason, router_gemm seems degrade accept length. + if ( + _is_cuda + and not self.is_nextn + and hidden_states.shape[0] < 4 + and hidden_states.shape[1] == 7168 + and self.weight.shape[0] == 256 + and _device_sm >= 90 + ): + from sgl_kernel import dsv3_router_gemm + + logits = dsv3_router_gemm(hidden_states, self.weight).to( + hidden_states.dtype + ) + else: + logits = F.linear(hidden_states, self.weight, None) + + return logits + + +class Glm4MoeLiteSparseMoeBlock(DeepseekV2MoE): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, + ): + nn.Module.__init__(self) + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.n_shared_experts = config.n_shared_experts + self.num_fused_shared_experts = ( + 0 + if get_global_server_args().disable_shared_experts_fusion + else config.n_shared_experts + ) + self.config = config + self.layer_id = layer_id + self.alt_stream = alt_stream + self.is_nextn = is_nextn + + if self.tp_size > config.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.n_routed_experts}." + ) + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + self.gate = Glm4MoeLiteGate( + config=config, prefix=add_prefix("gate", prefix), is_nextn=is_nextn + ) + + self.experts = get_moe_impl_class(quant_config)( + num_experts=config.n_routed_experts + + self.num_fused_shared_experts + + get_global_server_args().ep_num_redundant_experts, + num_fused_shared_experts=self.num_fused_shared_experts, + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + layer_id=self.layer_id, + quant_config=quant_config, + routed_scaling_factor=self.routed_scaling_factor, + prefix=add_prefix("experts", prefix), + ) + + self.topk = TopK( + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, + layer_id=self.layer_id, + renormalize=config.norm_topk_prob, + use_grouped_topk=True, + num_expert_group=config.n_group, + num_fused_shared_experts=self.num_fused_shared_experts, + topk_group=config.topk_group, + correction_bias=self.gate.e_score_correction_bias, + quant_config=quant_config, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk, + # Some Fp4 MoE backends require the output format to be bypassed but the MTP layers are unquantized + # and requires the output format to be standard. We use quant_config to determine the output format. + output_format=TopKOutputFormat.STANDARD if quant_config is None else None, + ) + + self.shared_experts_is_int8 = False + self.shared_experts_is_fp8 = False + # self.shared_experts_weight_block_size = None + if config.n_shared_experts is not None and self.num_fused_shared_experts == 0: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + # disable tp for shared experts when enable deepep moe, or with fp4 allgather + self.shared_experts = Glm4MoeLiteMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + **( + dict(tp_rank=0, tp_size=1) + if get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mooncake() + or should_use_flashinfer_cutlass_moe_fp4_allgather() + else {} + ), + ) + is_packed_weight = hasattr( + self.shared_experts.gate_up_proj.quant_method, "quant_config" + ) + self.shared_experts_is_int8 = ( + not is_packed_weight + and self.shared_experts.gate_up_proj.weight.dtype == torch.int8 + ) + self.shared_experts_is_fp8 = ( + not is_packed_weight + and self.shared_experts.gate_up_proj.weight.dtype == torch.float8_e4m3fn + ) + + self.top_k = config.num_experts_per_tok + + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): + # TODO: we will support tp < ep in the future + self.ep_size = get_moe_expert_parallel_world_size() + self.num_experts = ( + config.n_routed_experts + + get_global_server_args().ep_num_redundant_experts + ) + self.renormalize = config.norm_topk_prob + self.topk_group = config.topk_group + self.num_expert_group = config.n_group + self.correction_bias = ( + self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None + ) + + self._enable_a2a_moe = ( + get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() + ) + self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo() + + +class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + is_nextn: bool = False, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + nn.Module.__init__(self) + self.hidden_size = config.hidden_size + self.config = config + + from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp + + self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + rope_theta = 1000000 + rope_scaling = None + max_position_embeddings = getattr(config, "max_position_embeddings", 202752) + self.layer_id = layer_id + + self.self_attn = DeepseekV2AttentionMLA( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + reduce_results=False, + layer_id=layer_id, + prefix=add_prefix("self_attn", prefix), + ) + + self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn) + is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False) + is_next_layer_sparse = self._is_layer_sparse(layer_id + 1, is_nextn=False) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=1 if is_nextn else config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + + if self.is_layer_sparse: + self.mlp = Glm4MoeLiteSparseMoeBlock( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + layer_id=self.layer_id, + alt_stream=alt_stream, + is_nextn=is_nextn, + ) + else: + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = Glm4MoeLiteMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=( + is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) + ), + qkv_latent_func=self.self_attn.prepare_qkv_latent, + ) + + +class Glm4MoeLiteModel(DeepseekV2Model): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + nn.Module.__init__(self) + self.padding_id = config.pad_token_id + self.vocab_size = config.vocab_size + self.first_k_dense_replace = config.first_k_dense_replace + self.pp_group = get_pp_group() + + # DeepseekV2Model.forward expects these attributes to exist. + from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp + + self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.cp_size = get_attention_tp_size() if self.nsa_enable_prefill_cp else None + self.gemm_output_zero_allocator_size = 0 + self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not is_dp_attention_enabled(), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: Glm4MoeLiteDecoderLayer( + config=config, + layer_id=idx, + quant_config=quant_config, + prefix=prefix, + alt_stream=self.alt_stream, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + self.layers_to_capture = [] + + +class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + config.moe_layer_freq = 1 + self.config = config + self.tp_size = get_tensor_model_parallel_world_size() + self.quant_config = quant_config + self.pp_group = get_pp_group() + self.determine_num_fused_shared_experts("Glm4MoeLiteForCausalLM") + self.model = Glm4MoeLiteModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config) + + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) + } + ) + self.capture_aux_hidden_states = False + + from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp + + self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + if self.nsa_enable_prefill_cp: + from sglang.srt.layers.dp_attention import ( + get_attention_tp_rank, + get_attention_tp_size, + ) + + self.cp_rank = get_attention_tp_rank() + self.cp_size = get_attention_tp_size() + else: + self.cp_rank = self.cp_size = None + + def determine_num_fused_shared_experts( + self, architecture: str = "Glm4MoeLiteForCausalLM" + ): + self.num_fused_shared_experts = 0 + if get_global_server_args().disable_shared_experts_fusion: + return + + disable_reason = None + if ( + not _is_cuda + or torch.cuda.get_device_capability("cuda") < (8, 0) + or self.config.architectures[0] != architecture + or self.config.n_shared_experts != 1 + ): + disable_reason = "Only GLM-4.5 or GLM-4.6 on NV-platform with capability >= 80 can use shared experts fusion optimization." + elif get_moe_expert_parallel_world_size() > 1: + disable_reason = "GLM-4.5 or GLM-4.6 can not use shared experts fusion optimization under expert parallelism." + + if disable_reason is not None: + get_global_server_args().disable_shared_experts_fusion = True + self.num_fused_shared_experts = 0 + log_info_on_rank0( + logger, + f"{disable_reason} Shared experts fusion optimization is disabled.", + ) + return + + self.num_fused_shared_experts = self.config.n_shared_experts + + def load_weights( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + is_nextn=False, + params_dict=None, + is_eagle=False, + ): + if is_nextn: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + assert num_nextn_layers == 1, "Only 1 nextn layer is supported" + # compatible with old design + nextn_layer_id = ( + 0 + if self.config.num_hidden_layers == 1 + else self.config.num_hidden_layers + ) + else: + raise ValueError("num_nextn_predict_layers is not in the config") + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + if self.num_fused_shared_experts > 0: + assert self.num_fused_shared_experts == 1 + + def iter_weights_with_fused_shared_experts( + weights: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[Tuple[str, torch.Tensor]]: + import re + + pattern = re.compile( + r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$" + ) + for name, weight in weights: + match = pattern.match(name) + if match: + layer_id = int(match.group(1)) + suffix = match.group(2) + name = f"model.layers.{layer_id}.mlp.experts.{self.config.n_routed_experts}.{suffix}" + yield name, weight + + weights = iter_weights_with_fused_shared_experts(weights) + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, + ) + + # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None + fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( + self.config.q_lora_rank is not None + ) + cached_a_proj = {} if fuse_qkv_a_proj else None + + if is_nextn: + nextn_layer_prefix = f"model.layers.{nextn_layer_id}" + nextn_spec_weight_names = [ + "shared_head.norm", + "eh_proj", + "enorm", + "hnorm", + ] + else: + nextn_layer_prefix = None + nextn_spec_weight_names = [] + + eagle_ignore_weight_names = [] + if is_eagle: + eagle_ignore_weight_names = [ + "eagle_draft_tokens_map", + "eagle_lm_head.weight", + ] + + if params_dict is None: + params_dict = dict(self.named_parameters()) + + weight_names = [] + for name, loaded_weight in weights: + weight_names.append(name) + + if not is_nextn and not is_eagle: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers > 0 and name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) >= self.config.num_hidden_layers + ): + continue + else: + if nextn_layer_prefix and not name.startswith(nextn_layer_prefix): + continue + + if nextn_layer_prefix is not None: # mtp + # Use shared head and embed weights from target model + if "shared_head.head" in name or "embed_tokens" in name: + continue + + is_decoder = True + # For nextn specific weights + for weight_name in nextn_spec_weight_names: + if weight_name in name: + name = name.replace(nextn_layer_prefix, "model") + is_decoder = False + break + # For decoder layer weights + if is_decoder: + name = name.replace(nextn_layer_prefix, "model.decoder") + + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Track if this is an expert weight to enable early skipping + is_expert_weight = False + + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + + # Mark as expert weight regardless of whether we can process it + is_expert_weight = True + + name = name.replace(weight_name, param_name) + if name not in params_dict: + # Expert weight not on this rank, will be skipped below + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + if is_expert_weight: + # This is an expert weight but not mapped to this rank, skip all remaining processing + continue + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name in eagle_ignore_weight_names: + continue + + # GLM NOTE: for MLA + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + + # When both q_a_proj and kv_a_proj_with_mqa has been cached, load the fused weight to parameter + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + fused_weight = torch.cat( + [q_a_proj_weight, kv_a_proj_weight], dim=0 + ) + param_name = ( + name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa") + if "q_a_proj" in name + else name.replace( + "kv_a_proj_with_mqa", "fused_qkv_a_proj_with_mqa" + ) + ) + if param_name not in params_dict: + continue + param = params_dict[param_name] + + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + else: + if ( + "k_scale" in name or "v_scale" in name + ) and name not in params_dict: + # modelopt attn kv scale is named differently + if any(scale in name for scale in ["k_scale", "v_scale"]): + name = name.replace("_proj", "attn_mqa") + else: + logger.warning( + f"Unknown scale found in checkpoint: {name}" + ) + + if name not in params_dict: + continue + + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning(f"Parameter {name} not found in params_dict") + + # DeepseekV2AttentionMLA.forward_* expects post_load_weights() to populate + # per-layer packed weights like `w_kc`/`w_vc` (used during CUDA graph capture). + # GLM-Lite configs may not set `config.mla`, but this model always uses + # DeepseekV2AttentionMLA, so we must run the post-load processing. + # Use weight_names=None to ensure we always process all layers. Some checkpoints / + # naming schemes may not include "kv_b_proj" in `weight_names`, but `w_kc`/`w_vc` + # are still required by DeepseekV2AttentionMLA at runtime. + self.post_load_weights(is_nextn=is_nextn, weight_names=None) + + +EntryClass = [Glm4MoeLiteForCausalLM] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 629066f96..193904ef1 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1586,6 +1586,7 @@ class ServerArgs: "DeepseekV3ForCausalLM", "GptOssForCausalLM", "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", "Qwen3MoeForCausalLM", ] and (is_sm90_supported() or is_sm100_supported()) @@ -2147,6 +2148,7 @@ class ServerArgs: "DeepseekV32ForCausalLM", "DeepseekV3ForCausalLM", "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", "BailingMoeForCausalLM", "BailingMoeV2ForCausalLM", "MistralLarge3ForCausalLM", @@ -5422,6 +5424,7 @@ def auto_choose_speculative_params(self: ServerArgs): "DeepseekV2ForCausalLM", "GptOssForCausalLM", "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", "BailingMoeForCausalLM", "BailingMoeV2ForCausalLM", "MistralLarge3ForCausalLM",