sglang/test/registered/unit/moe/test_glm_megamoe.py

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, 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,
        )
        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, 2), dtype=torch.float32)),
            w2_weight_scale=_Param(torch.ones((2, 128, 2), 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.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(),
            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()