Feat: GLM-4.6 supports shared experts fusion (#13873)

Signed-off-by: UranusSeven <109661872+UranusSeven@users.noreply.github.com>
Co-authored-by: Kevin-XiongC <kevin_xiong1997@outlook.com>
Co-authored-by: Mingyi Jin <jinmingyi1998@sina.cn>

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

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
+    }
+}

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,
     )
 

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))

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)

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,
+            num_experts=config.n_routed_experts + self.num_fused_shared_experts,
-            top_k=self.top_k,
+            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) and (self.num_fused_shared_experts == 0):
-        if hidden_states.shape[0] > 0:
+            return self.shared_experts(hidden_states)
-            shared_output = self.shared_experts(hidden_states)
+        else:
-        return shared_output
+            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

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,