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:
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user