From 57bf729af9e2cad0950160ed9e9c7cba30da06c9 Mon Sep 17 00:00:00 2001 From: LuminolT Date: Mon, 6 Jul 2026 10:39:23 +0800 Subject: [PATCH] feat(glm): route GLM MoE through megamoe backend Add the GLM MoE MegaMoE fast path before the existing DeepEP/normal branch selection. Keep the existing none+flashinfer_trtllm path as the miss path when MegaMoE weights, envs, or token caps do not satisfy should_use_mega_moe(). Match upstream MegaMoE shared-expert TP1 semantics for the separate shared expert path. Constraint: do not change the baseline none+flashinfer_trtllm dispatch path. Feature-flag: --moe-a2a-backend=megamoe. Conflict-hotspots: python/sglang/srt/models/glm4_moe.py. Scope-risk: shared-expert TP1 behavior matters when shared expert fusion is disabled. Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile python/sglang/srt/models/glm4_moe.py. Tested: git diff --check. Not-tested: GLM 5.2 MegaMoE GPU e2e; local environment lacks target runtime and hardware. --- python/sglang/srt/models/glm4_moe.py | 8 ++++++++ 1 file changed, 8 insertions(+) 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 (