From c00630088cd052f9f6687701d3c7ae36df8c6826 Mon Sep 17 00:00:00 2001 From: LuminolT Date: Mon, 6 Jul 2026 10:46:18 +0800 Subject: [PATCH] fix(moe): infer megamoe fp4 weight scale group 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. --- python/sglang/srt/layers/moe/mega_moe.py | 30 ++++++++++++++++++-- test/registered/unit/moe/test_glm_megamoe.py | 9 +++--- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 5a8ce8f56..5313e4002 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -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 diff --git a/test/registered/unit/moe/test_glm_megamoe.py b/test/registered/unit/moe/test_glm_megamoe.py index a19bd45fc..d9807c182 100644 --- a/test/registered/unit/moe/test_glm_megamoe.py +++ b/test/registered/unit/moe/test_glm_megamoe.py @@ -103,8 +103,8 @@ class TestGLMMegaMoE(unittest.TestCase): def test_build_mega_moe_weights_preserves_runner_layout(self): def fake_transform_sf(sf, *, mn, k, recipe, num_groups, disable_ue8m0_cast): - del sf, k, recipe, disable_ue8m0_cast - return torch.zeros((num_groups, mn, 1), dtype=torch.int32) + del sf, disable_ue8m0_cast + return torch.zeros((num_groups, mn, k // recipe[1]), dtype=torch.int32) fake_deep_gemm = types.SimpleNamespace( transform_sf_into_required_layout=fake_transform_sf, @@ -112,8 +112,8 @@ class TestGLMMegaMoE(unittest.TestCase): experts = SimpleNamespace( w13_weight=_Param(torch.arange(2 * 256 * 16).reshape(2, 256, 16)), w2_weight=_Param(torch.arange(2 * 128 * 16).reshape(2, 128, 16)), - w13_weight_scale=_Param(torch.ones((2, 256, 1), dtype=torch.float32)), - w2_weight_scale=_Param(torch.ones((2, 128, 1), dtype=torch.float32)), + w13_weight_scale=_Param(torch.ones((2, 256, 2), dtype=torch.float32)), + w2_weight_scale=_Param(torch.ones((2, 128, 2), dtype=torch.float32)), ) original_w13 = experts.w13_weight.data.clone() @@ -124,6 +124,7 @@ class TestGLMMegaMoE(unittest.TestCase): self.assertTrue(experts._mega_moe_weights_built) self.assertTrue(experts._mega_moe_preserve_runner_layout) + self.assertEqual(experts._mega_moe_weight_group_size, 16) self.assertTrue(torch.equal(experts.w13_weight.data, original_w13)) self.assertNotEqual( experts.mega_l1_weights[0].data_ptr(),