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:
@@ -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 (
|
||||
|
||||
Reference in New Issue
Block a user