diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index ab40fdfcd..7673c49c9 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -15,6 +15,7 @@ """Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 model compatible with HuggingFace weights""" import logging +import re from typing import Any, Dict, Iterable, List, Optional, Tuple, Union import torch @@ -22,6 +23,7 @@ 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.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( get_moe_expert_parallel_world_size, @@ -63,6 +65,7 @@ from sglang.srt.layers.moe import ( ) 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.kt_ep_wrapper import KTEPWrapperMethod from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.moe.utils import ( RoutingMethodType, @@ -183,7 +186,7 @@ class Glm4MoeAttention(nn.Module): num_heads: int, num_kv_heads: int, layer_id: int = 0, - rope_theta: float = 10000, + rope_theta: float = 1000000, partial_rotary_factor: float = 0.5, rope_scaling: Optional[Dict[str, Any]] = None, max_position_embeddings: int = 8192, @@ -439,9 +442,12 @@ class Glm4MoeSparseMoeBlock(nn.Module): fused_shared_experts_scaling_factor=1, ) - # shared expert + 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 = Glm4MoeMLP( hidden_size=config.hidden_size, intermediate_size=intermediate_size, @@ -453,16 +459,54 @@ class Glm4MoeSparseMoeBlock(nn.Module): dict(tp_rank=0, tp_size=1) if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() + or get_moe_a2a_backend().is_nixl() + or get_moe_a2a_backend().is_mori() + or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() or should_use_flashinfer_cutlass_moe_fp4_allgather() else {} ), ) + is_packed_weight = hasattr( + self.shared_experts.gate_up_proj.quant_method, "quant_config" + ) and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() in { + "awq", + "awq_marlin", + "moe_wna16", + } + 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 + ) + if self.shared_experts_is_fp8: + if ( + _use_aiter + and config.quantization_config.get("quant_method") + == "compressed-tensors" + ): + # For compressed-tensors ptpc model, don't need to check the weight_block_size + pass + else: + assert ( + self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size + == self.shared_experts.down_proj.quant_method.quant_config.weight_block_size + ) + self.shared_experts_weight_block_size = ( + self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size + ) + + self.top_k = config.num_experts_per_tok if ( get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() + or get_moe_a2a_backend().is_mori() + or get_moe_a2a_backend().is_ascend_fuseep() ): # TODO: we will support tp < ep in the future self.ep_size = get_moe_expert_parallel_world_size() @@ -483,7 +527,11 @@ class Glm4MoeSparseMoeBlock(nn.Module): get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() + or get_moe_a2a_backend().is_mori() + or get_moe_a2a_backend().is_ascend_fuseep() + or get_moe_a2a_backend().is_flashinfer() ) + self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo() def get_moe_weights(self): return [ @@ -502,8 +550,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: - - if not get_moe_a2a_backend().is_deepep(): + if not self._enable_a2a_moe: if ( self.alt_stream is not None and self.num_fused_shared_experts == 0 @@ -511,11 +558,15 @@ class Glm4MoeSparseMoeBlock(nn.Module): and get_is_capture_mode() ): return self.forward_normal_dual_stream( - hidden_states, should_allreduce_fusion, use_reduce_scatter + hidden_states, + should_allreduce_fusion, + use_reduce_scatter, ) else: return self.forward_normal( - hidden_states, should_allreduce_fusion, use_reduce_scatter + hidden_states, + should_allreduce_fusion, + use_reduce_scatter, ) else: return self.forward_deepep(hidden_states, forward_batch) @@ -534,20 +585,12 @@ class Glm4MoeSparseMoeBlock(nn.Module): # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) - final_hidden_states = self.experts(hidden_states, topk_output) - if not _is_cuda and not _use_aiter: - # fused in biased_grouped_topk so we can skip here + if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod): final_hidden_states *= self.routed_scaling_factor current_stream.wait_stream(self.alt_stream) - - with use_symmetric_memory( - parallel_state.get_tp_group(), disabled=not is_allocation_symmetric() - ): - final_hidden_states_out = torch.empty_like(final_hidden_states) - torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) - final_hidden_states = final_hidden_states_out + final_hidden_states += shared_output if ( self.tp_size > 1 and not should_allreduce_fusion @@ -1090,27 +1133,32 @@ class Glm4MoeForCausalLM(nn.Module): # For EAGLE3 support self.capture_aux_hidden_states = False - def get_input_embeddings(self) -> nn.Embedding: - return self.model.embed_tokens - def determine_num_fused_shared_experts(self): if get_global_server_args().disable_shared_experts_fusion: return disable_reason = None - if not getattr(self.config, "n_shared_experts", None): - disable_reason = "No shared experts are defined in the config." - elif not _is_cuda: - disable_reason = "Shared experts fusion currently requires CUDA devices." - elif _is_cuda and (_device_sm is not None) and (_device_sm < 80): - disable_reason = "Shared experts fusion requires SM80 or newer GPUs." - elif get_moe_expert_parallel_world_size() > 1: - disable_reason = "Shared experts fusion is not supported together with expert parallelism yet." - elif get_moe_a2a_backend().is_deepep(): - disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled." + if (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( + not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) + ): + disable_reason = ( + "Only GLM-4.5 on NV-platform with capability >= 80 " + "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization." + ) + elif get_moe_expert_parallel_world_size() > 1 and ( + not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) + ): + disable_reason = "Only GLM-4.5 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism." + elif disable_reason is None and ( + get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori() + ): + disable_reason = "GLM-4.5 cannot use shared experts fusion optimization under deepep expert parallelism." + elif self.quant_config and self.quant_config.get_name() == "w4afp8": + disable_reason = "GLM-4.5 W4AFP8 model uses different quant method for routed experts and shared experts." 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.", @@ -1118,10 +1166,9 @@ class Glm4MoeForCausalLM(nn.Module): return self.num_fused_shared_experts = self.config.n_shared_experts - assert ( - self.num_fused_shared_experts == 1 - ), "Only 1 fused shared expert is supported for Glm4MoeForCausalLM" - log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens @torch.no_grad() def forward( @@ -1154,7 +1201,12 @@ class Glm4MoeForCausalLM(nn.Module): def end_layer(self): return self.model.end_layer - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): + def load_weights( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + is_nextn=False, + params_dict=None, + ): if is_nextn: if hasattr(self.config, "num_nextn_predict_layers"): num_nextn_layers = self.config.num_nextn_predict_layers @@ -1177,6 +1229,28 @@ class Glm4MoeForCausalLM(nn.Module): ("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]]: + + 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", @@ -1192,20 +1266,17 @@ class Glm4MoeForCausalLM(nn.Module): "enorm", "hnorm", ] + else: + nextn_layer_prefix = None + nextn_spec_weight_names = [] + + if params_dict is None: + params_dict = dict(self.named_parameters()) - params_dict = dict(self.named_parameters()) weight_names = [] for name, loaded_weight in weights: weight_names.append(name) - if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: - # Map shared expert weights to the last expert slot - # Shared expert becomes expert ID = n_routed_experts - name = name.replace( - "mlp.shared_experts", - f"mlp.experts.{self.config.n_routed_experts}", - ) - if not is_nextn: if hasattr(self.config, "num_nextn_predict_layers"): num_nextn_layers = self.config.num_nextn_predict_layers @@ -1217,23 +1288,24 @@ class Glm4MoeForCausalLM(nn.Module): ): continue else: - if not name.startswith(nextn_layer_prefix): + if nextn_layer_prefix and not name.startswith(nextn_layer_prefix): continue - # Use shared head and embed weights from target model - if "shared_head.head" in name or "embed_tokens" in name: - 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") + 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 @@ -1295,6 +1367,7 @@ class Glm4MoeForCausalLM(nn.Module): # 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 diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index ffd9acc4e..03ff3d201 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -12,13 +12,15 @@ # limitations under the License. # ============================================================================== -"""Inference-only GLM-Lite model compatible with HuggingFace weights""" +"""Inference-only GLM-4.7-Flash model compatible with HuggingFace weights""" import logging +import re from typing import Iterable, Optional, Tuple import torch import torch.nn.functional as F +from sgl_kernel import dsv3_router_gemm from torch import nn from transformers import PretrainedConfig @@ -29,12 +31,14 @@ from sglang.srt.distributed import ( get_tensor_model_parallel_world_size, ) from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp from sglang.srt.layers.communicator import ( LayerCommunicator, LayerScatterModes, enable_moe_dense_fully_dp, ) from sglang.srt.layers.dp_attention import ( + get_attention_tp_rank, get_attention_tp_size, is_dp_attention_enabled, ) @@ -72,6 +76,7 @@ from sglang.srt.utils import ( log_info_on_rank0, make_layers, ) +from sglang.srt.utils.hf_transformers_utils import get_rope_config _is_cuda = is_cuda() _device_sm = get_device_sm() @@ -183,7 +188,6 @@ class Glm4MoeLiteGate(nn.Module): 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 @@ -335,12 +339,8 @@ class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer): 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 + rope_theta, rope_scaling = get_rope_config(config) max_position_embeddings = getattr(config, "max_position_embeddings", 202752) self.layer_id = layer_id @@ -429,8 +429,6 @@ class Glm4MoeLiteModel(DeepseekV2Model): 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 @@ -501,15 +499,8 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): ) 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: @@ -549,7 +540,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): 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"): @@ -579,7 +569,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): 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\.(.+)$" @@ -621,13 +610,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): 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()) @@ -635,7 +617,7 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): for name, loaded_weight in weights: weight_names.append(name) - if not is_nextn and not is_eagle: + if not is_nextn: 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"): @@ -725,8 +707,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): # 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 ( @@ -797,7 +777,7 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): # 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 + # GLM-4.7-Flash configs 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`