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.
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
+5 -4
View File
@@ -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(),