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.
This commit is contained in:
LuminolT
2026-07-06 10:39:23 +08:00
parent c13d4556ff
commit 57bf729af9

View File

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