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

fix(moe): infer megamoe fp4 weight scale group

Browse Source 
Infer the FP4 weight scale group size from the loaded scale tensor instead of hard-coding K/32.

This keeps upstream-style FP4 expert layouts working while allowing GLM/ModelOpt NVFP4 layouts that use K/16 scale columns to build MegaMoE sidecar weights.

Constraint: preserve the existing runner layout and only change MegaMoE sidecar metadata/recipe.

Feature-flag: --moe-a2a-backend=megamoe.

Conflict-hotspots: python/sglang/srt/layers/moe/mega_moe.py.

Scope-risk: actual DeepGEMM recipe support still needs target GPU runtime validation.

Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile python/sglang/srt/layers/moe/mega_moe.py test/registered/unit/moe/test_glm_megamoe.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:46:18 +08:00
parent 93e3840578
commit c00630088c
2 changed files with 32 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
python/sglang/srt/layers/moe/mega_moe.py
View File

@@ -286,12 +286,13 @@ def _run_mega_routed(
device=hidden_states.device,
)
swiglu_limit = getattr(moe.config, "swiglu_limit", None)
weight_group_size = getattr(moe.experts, "_mega_moe_weight_group_size", 32)
deep_gemm.fp8_fp4_mega_moe(
y,
moe.experts.mega_l1_weights,
moe.experts.mega_l2_weights,
buf,
recipe=(1, 1, 32),
recipe=(1, 1, weight_group_size),
activation="swiglu",
activation_clamp=swiglu_limit,
fast_math=True,
@@ -344,6 +345,21 @@ def _get_mega_moe_scale_tensors(experts) -> tuple[torch.Tensor, torch.Tensor]:
)
def _infer_mega_moe_scale_group_size(
scale: torch.Tensor,
*,
k: int,
name: str,
) -> int:
num_scale_cols = scale.shape[-1]
if num_scale_cols <= 0 or k % num_scale_cols != 0:
raise ValueError(
f"MegaMoE cannot infer {name} scale group size from "
f"k={k}, scale_shape={tuple(scale.shape)}."
)
return k // num_scale_cols
def build_mega_moe_experts_weights(
experts,
*,
@@ -362,12 +378,19 @@ def build_mega_moe_experts_weights(
k1 = half_k1 * 2
_, n2, half_k2 = w2.shape
k2 = half_k2 * 2
group_k1 = _infer_mega_moe_scale_group_size(w13_sf_src, k=k1, name="w13")
group_k2 = _infer_mega_moe_scale_group_size(w2_sf_src, k=k2, name="w2")
if group_k1 != group_k2:
raise ValueError(
f"MegaMoE requires matching w13/w2 scale group sizes, got "
f"{group_k1=} and {group_k2=}."
)
w13_sf = transform_sf_into_required_layout(
w13_sf_src.to(torch.float32),
mn=n1,
k=k1,
recipe=(1, 32),
recipe=(1, group_k1),
num_groups=num_groups,
disable_ue8m0_cast=False,
)
@@ -375,7 +398,7 @@ def build_mega_moe_experts_weights(
w2_sf_src.to(torch.float32),
mn=n2,
k=k2,
recipe=(1, 32),
recipe=(1, group_k2),
num_groups=num_groups,
disable_ue8m0_cast=False,
)
@@ -420,4 +443,5 @@ def build_mega_moe_experts_weights(
experts._mega_moe_hidden_size = k1
experts._mega_moe_intermediate_size = k2
experts._mega_moe_weight_group_size = group_k1
experts._mega_moe_weights_built = True