From 93e3840578dbef7857649fc4300f0843ea2a4765 Mon Sep 17 00:00:00 2001 From: LuminolT Date: Mon, 6 Jul 2026 10:44:05 +0800 Subject: [PATCH] test(moe): add GLM MegaMoE validation coverage Add parser coverage for --moe-a2a-backend=megamoe and local unit coverage for MegaMoE gating, GLM fast-path routing, fallback routing, and sidecar weight preservation. Use monkeypatch-style tests so the forward and weight-layout checks do not require loading GLM 5.2 or importing a real deep_gemm runtime. Constraint: tests must not stage or depend on local skill/doc artifacts. Feature-flag: --moe-a2a-backend=megamoe. Conflict-hotspots: test/registered/unit/server_args/test_server_args.py, test/registered/unit/moe/test_glm_megamoe.py. Scope-risk: runtime GPU e2e still needs target Blackwell/DeepGEMM environment. Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile test/registered/unit/server_args/test_server_args.py test/registered/unit/moe/test_glm_megamoe.py. Tested: git diff --check. Not-tested: PYTHONPATH=python python3 -m pytest test/registered/unit/moe/test_glm_megamoe.py -q; local Python has no pytest. Not-tested: PYTHONPATH=python python3 -m unittest test.registered.unit.moe.test_glm_megamoe -v; local Python has no torch. Not-tested: GLM 5.2 MegaMoE GPU e2e; local environment lacks target runtime and hardware. --- test/registered/unit/moe/test_glm_megamoe.py | 140 ++++++++++++++++++ .../unit/server_args/test_server_args.py | 41 +++++ 2 files changed, 181 insertions(+) create mode 100644 test/registered/unit/moe/test_glm_megamoe.py diff --git a/test/registered/unit/moe/test_glm_megamoe.py b/test/registered/unit/moe/test_glm_megamoe.py new file mode 100644 index 000000000..a19bd45fc --- /dev/null +++ b/test/registered/unit/moe/test_glm_megamoe.py @@ -0,0 +1,140 @@ +import sys +import types +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.moe import mega_moe +from sglang.srt.models import glm4_moe +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + + +class _MegaBackend: + def is_megamoe(self): + return True + + +class _NonDeepEPBackend: + def is_deepep(self): + return False + + +class _Param: + def __init__(self, data): + self.data = data + + +class TestGLMMegaMoE(unittest.TestCase): + def test_should_use_mega_moe_respects_env_and_token_cap(self): + hidden_states = torch.empty((2, 32)) + moe = SimpleNamespace( + experts=SimpleNamespace(_mega_moe_weights_built=True), + ) + + with patch.object( + mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend() + ), patch.object( + mega_moe, "get_dp_global_num_tokens", return_value=None + ), patch.object( + mega_moe, "get_is_capture_mode", return_value=False + ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override(False): + self.assertFalse(mega_moe.should_use_mega_moe(moe, hidden_states)) + + with patch.object( + mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend() + ), patch.object( + mega_moe, "get_dp_global_num_tokens", return_value=None + ), patch.object( + mega_moe, "get_is_capture_mode", return_value=False + ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override( + True + ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override(1): + self.assertFalse(mega_moe.should_use_mega_moe(moe, hidden_states)) + + with patch.object( + mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend() + ), patch.object( + mega_moe, "get_dp_global_num_tokens", return_value=None + ), patch.object( + mega_moe, "get_is_capture_mode", return_value=False + ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override( + True + ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override(2): + self.assertTrue(mega_moe.should_use_mega_moe(moe, hidden_states)) + + def test_glm_forward_uses_megamoe_fast_path(self): + block = object.__new__(glm4_moe.Glm4MoeSparseMoeBlock) + hidden_states = torch.empty((1, 32)) + forward_batch = object() + expected = torch.empty_like(hidden_states) + + with patch.object( + mega_moe, "should_use_mega_moe", return_value=True + ) as should, patch.object( + mega_moe, "forward_mega_moe", return_value=expected + ) as forward: + output = block.forward(hidden_states, forward_batch=forward_batch) + + self.assertIs(output, expected) + should.assert_called_once_with(block, hidden_states) + forward.assert_called_once_with(block, hidden_states, forward_batch) + + def test_glm_forward_falls_back_to_normal_path(self): + block = object.__new__(glm4_moe.Glm4MoeSparseMoeBlock) + object.__setattr__(block, "alt_stream", None) + object.__setattr__(block, "num_fused_shared_experts", 0) + object.__setattr__(block, "forward_normal", Mock(return_value="normal-output")) + hidden_states = torch.empty((1, 32)) + + with patch.object( + mega_moe, "should_use_mega_moe", return_value=False + ), patch.object( + glm4_moe, "get_moe_a2a_backend", return_value=_NonDeepEPBackend() + ): + output = block.forward(hidden_states) + + self.assertEqual(output, "normal-output") + block.forward_normal.assert_called_once_with(hidden_states, False, False) + + 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) + + fake_deep_gemm = types.SimpleNamespace( + transform_sf_into_required_layout=fake_transform_sf, + ) + 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)), + ) + original_w13 = experts.w13_weight.data.clone() + + with patch.dict(sys.modules, {"deep_gemm": fake_deep_gemm}): + mega_moe.build_mega_moe_experts_weights( + experts, preserve_runner_layout=True + ) + + self.assertTrue(experts._mega_moe_weights_built) + self.assertTrue(experts._mega_moe_preserve_runner_layout) + self.assertTrue(torch.equal(experts.w13_weight.data, original_w13)) + self.assertNotEqual( + experts.mega_l1_weights[0].data_ptr(), + experts.w13_weight.data.data_ptr(), + ) + + experts.w13_weight.data.zero_() + self.assertFalse( + torch.equal(experts.mega_l1_weights[0], experts.w13_weight.data) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index c214369ee..79b85d49d 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1,4 +1,5 @@ import json +import os import tempfile import unittest from unittest.mock import MagicMock, patch @@ -195,6 +196,46 @@ def test_nsa_index_skip_topk_offset_zero_means_unset(): assert override_args == {"index_topk_freq": 4} +def test_megamoe_a2a_backend_parser_sets_ep_to_tp(): + import argparse + + from sglang.srt.environ import envs + + field = envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY + backup_present = field.name in os.environ + backup_value = os.environ.get(field.name) + backup_set_to_none = field._set_to_none + os.environ.pop(field.name, None) + field._set_to_none = False + field._cache_valid = False + + try: + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + raw_args = parser.parse_args( + [ + "--model-path", + "dummy", + "--tp-size", + "4", + "--moe-a2a-backend", + "megamoe", + ] + ) + args = ServerArgs.from_cli_args(raw_args) + assert args.moe_a2a_backend == "megamoe" + assert args.tp_size == 4 + assert args.ep_size == 4 + assert field.get() is True + finally: + if backup_present: + os.environ[field.name] = backup_value + else: + os.environ.pop(field.name, None) + field._set_to_none = backup_set_to_none + field._cache_valid = False + + class TestLoadBalanceMethod(unittest.TestCase): def test_non_pd_defaults_to_round_robin(self): server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")