chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Uranus
2025-12-01 11:33:18 +08:00
committed by GitHub
parent f5f3a5d98c
commit 982db4ebac
7 changed files with 252 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

146
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 Normal file
View File

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

1
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View File

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

21
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py
View File

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

1
python/sglang/srt/layers/moe/moe_runner/triton.py
View File

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

93
python/sglang/srt/models/glm4_moe.py
View File

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

4
python/sglang/srt/models/glm4_moe_nextn.py
View File

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