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(),