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

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

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
+            del sf, disable_ue8m0_cast
-            return torch.zeros((num_groups, mn, 1), dtype=torch.int32)
+            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)),
+            w13_weight_scale=_Param(torch.ones((2, 256, 2), dtype=torch.float32)),
-            w2_weight_scale=_Param(torch.ones((2, 128, 1), 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(),