diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 0c1e1d000..c7bbffe81 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -415,6 +415,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_flashinfer() + or get_moe_a2a_backend().is_megamoe() or should_use_flashinfer_cutlass_moe_fp4_allgather() else {} ), @@ -463,6 +464,13 @@ class Glm4MoeSparseMoeBlock(nn.Module): should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: + from sglang.srt.layers.moe.mega_moe import ( + forward_mega_moe, + should_use_mega_moe, + ) + + if should_use_mega_moe(self, hidden_states): + return forward_mega_moe(self, hidden_states, forward_batch) if not get_moe_a2a_backend().is_deepep(): if (