diff --git a/benchmark/kernels/fused_moe_triton/common_utils.py b/benchmark/kernels/fused_moe_triton/common_utils.py index d87350f9f..40f8697f1 100644 --- a/benchmark/kernels/fused_moe_triton/common_utils.py +++ b/benchmark/kernels/fused_moe_triton/common_utils.py @@ -69,11 +69,15 @@ def get_model_config( E = config.num_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size - elif architecture in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]: + elif architecture in [ + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + ]: E = (config.n_routed_experts // ep_size) + ( 0 if disable_shared_experts_fusion - or architecture not in ["DeepseekV3ForCausalLM"] + or architecture not in ["DeepseekV3ForCausalLM", "Glm4MoeForCausalLM"] else 1 ) topk = config.num_experts_per_tok + ( @@ -104,7 +108,7 @@ def get_model_config( E = config.num_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size - elif architecture in ["Glm4MoeForCausalLM", "NemotronHForCausalLM"]: + elif architecture == "NemotronHForCausalLM": E = config.n_routed_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=161,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=161,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json new file mode 100644 index 000000000..4dd4e49ff --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=161,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index 230389a35..e7d5a67cc 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -434,6 +434,7 @@ def fused_experts_impl( topk_ids.shape[1], config_dtype, block_shape=block_shape, + per_channel_quant=per_channel_quant, return_down_config=True, ) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py index c1624007f..90a6a31c7 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py @@ -208,6 +208,7 @@ def try_get_optimal_moe_config( M: int, is_marlin: bool = False, block_shape: Optional[List[int]] = None, + per_channel_quant: bool = False, return_down_config: bool = False, ): from sglang.srt.layers.moe.fused_moe_triton import get_config @@ -222,7 +223,15 @@ def try_get_optimal_moe_config( E, _, N = w2_shape block_n = block_shape[0] if block_shape else 0 block_k = block_shape[1] if block_shape else 0 - configs = get_moe_configs(E, N, dtype, block_n, block_k, down_moe=False) + configs = get_moe_configs( + E, + N, + dtype, + block_n, + block_k, + per_channel_quant=per_channel_quant, + down_moe=False, + ) if configs: # If an optimal configuration map has been found, look up the @@ -234,7 +243,15 @@ def try_get_optimal_moe_config( M, E, N, w1_shape[2], top_k, dtype, is_marlin, block_shape ) if return_down_config: - down_configs = get_moe_configs(E, N, dtype, block_n, block_k, down_moe=True) + down_configs = get_moe_configs( + E, + N, + dtype, + block_n, + block_k, + per_channel_quant=per_channel_quant, + down_moe=True, + ) if down_configs: down_config = down_configs[ min(down_configs.keys(), key=lambda x: abs(x - M)) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index 7de28f482..ebb6ba1b4 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -417,6 +417,7 @@ def pre_permute_standard_to_triton( topk_output.topk_ids.shape[1], config_dtype, block_shape=quant_info.block_shape, + per_channel_quant=quant_info.per_channel_quant, ) config = get_config_func(num_tokens) diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 0267a9593..1fcbfb45a 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -85,6 +85,7 @@ from sglang.srt.utils import ( is_cuda, is_hip, is_non_idle_and_non_empty, + log_info_on_rank0, make_layers, ) @@ -352,8 +353,14 @@ class Glm4MoeSparseMoeBlock(nn.Module): nn.Module.__init__(self) self.top_k = config.num_experts_per_tok self.tp_size = get_tensor_model_parallel_world_size() + self.moe_ep_size = get_moe_expert_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 @@ -372,19 +379,10 @@ class Glm4MoeSparseMoeBlock(nn.Module): self.gate = Glm4MoeGate(config=config, prefix=add_prefix("gate", prefix)) - self.topk = TopK( - top_k=self.top_k, - renormalize=config.norm_topk_prob, - use_grouped_topk=True, - num_expert_group=config.n_group, - topk_group=config.topk_group, - correction_bias=self.gate.e_score_correction_bias, - routed_scaling_factor=self.routed_scaling_factor, - ) - self.experts = get_moe_impl_class(quant_config)( - num_experts=config.n_routed_experts, - top_k=self.top_k, + num_experts=config.n_routed_experts + self.num_fused_shared_experts, + num_fused_shared_experts=self.num_fused_shared_experts, + top_k=self.top_k + self.num_fused_shared_experts, layer_id=self.layer_id, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, @@ -393,8 +391,23 @@ class Glm4MoeSparseMoeBlock(nn.Module): prefix=add_prefix("experts", prefix), ) + self.topk = TopK( + top_k=self.top_k + self.num_fused_shared_experts, + renormalize=config.norm_topk_prob, + use_grouped_topk=True, + num_expert_group=config.n_group, + topk_group=config.topk_group, + correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + num_fused_shared_experts=self.num_fused_shared_experts, + apply_routed_scaling_factor_on_output=getattr( + self.experts, "should_fuse_routed_scaling_factor_in_topk", False + ), + fused_shared_experts_scaling_factor=1, + ) + # shared expert - if config.n_shared_experts is not 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 self.shared_experts = Glm4MoeMLP( hidden_size=config.hidden_size, @@ -450,6 +463,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): if not get_moe_a2a_backend().is_deepep(): if ( self.alt_stream is not None + and self.num_fused_shared_experts == 0 and hidden_states.shape[0] > 0 and get_is_capture_mode() ): @@ -472,6 +486,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) shared_output = self._forward_shared_experts(hidden_states) + with torch.cuda.stream(self.alt_stream): # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states) @@ -483,6 +498,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): 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() ): @@ -515,7 +531,6 @@ class Glm4MoeSparseMoeBlock(nn.Module): 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 final_hidden_states *= self.routed_scaling_factor if shared_output is not None: with use_symmetric_memory( @@ -570,10 +585,10 @@ class Glm4MoeSparseMoeBlock(nn.Module): return final_hidden_states def _forward_shared_experts(self, hidden_states: torch.Tensor): - shared_output = None - if hidden_states.shape[0] > 0: - shared_output = self.shared_experts(hidden_states) - return shared_output + if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0): + return self.shared_experts(hidden_states) + else: + return None def op_gate(self, state): if is_non_idle_and_non_empty( @@ -993,6 +1008,8 @@ class Glm4MoeForCausalLM(nn.Module): self.config = config self.tp_size = get_tensor_model_parallel_world_size() self.quant_config = quant_config + self.num_fused_shared_experts = 0 + self.determine_num_fused_shared_experts() self.model = Glm4MoeModel( config, quant_config, prefix=add_prefix("model", prefix) ) @@ -1011,6 +1028,36 @@ class Glm4MoeForCausalLM(nn.Module): 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 disable_reason is not None: + get_global_server_args().disable_shared_experts_fusion = True + 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 + 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.") + @torch.no_grad() def forward( self, @@ -1069,7 +1116,7 @@ class Glm4MoeForCausalLM(nn.Module): 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, + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, ) if is_nextn: @@ -1086,6 +1133,14 @@ class Glm4MoeForCausalLM(nn.Module): 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 diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py index cb44a58e6..5e0c3ac59 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -139,6 +139,10 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM): ) self.logits_processor = LogitsProcessor(config) + self.num_fused_shared_experts = ( + 0 if get_global_server_args().disable_shared_experts_fusion else 1 + ) + @torch.no_grad() def forward( self,