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,
                _mega_moe_weight_format="fp4",
            ),
        )

        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_should_use_mega_moe_fp8_does_not_require_fp4_act_env(self):
        hidden_states = torch.empty((2, 32))
        moe = SimpleNamespace(
            experts=SimpleNamespace(
                _mega_moe_weights_built=True,
                _mega_moe_weight_format="fp8",
            ),
        )

        with patch.object(
            mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend()
        ), patch.object(
            mega_moe, "_deep_gemm_supports_fp8_mega_moe", return_value=True
        ), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override(False):
            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.assertEqual(experts._mega_moe_weight_format, "fp4")
        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)
        )

    def test_fp4_weights_use_int8_view_for_deep_gemm_ffi(self):
        raw = torch.arange(16, dtype=torch.uint8).reshape(2, 8)
        fp4 = raw.view(torch.float4_e2m1fn_x2)
        scale = torch.ones((2, 8), dtype=torch.int32)

        data, out_scale = mega_moe._as_deep_gemm_packed_fp4_weights((fp4, scale))

        self.assertEqual(data.dtype, torch.int8)
        self.assertEqual(data.data_ptr(), fp4.data_ptr())
        self.assertEqual(data.shape, fp4.shape)
        self.assertIs(out_scale, scale)

        int8_data = raw.view(torch.int8)
        same_data, same_scale = mega_moe._as_deep_gemm_packed_fp4_weights(
            (int8_data, scale)
        )
        self.assertIs(same_data, int8_data)
        self.assertIs(same_scale, scale)

    def test_build_fp8_mega_moe_weights_preserves_runner_layout(self):
        experts = SimpleNamespace(
            w13_weight=_Param(torch.arange(2 * 8 * 4).reshape(2, 8, 4)),
            w2_weight=_Param(torch.arange(2 * 4 * 4).reshape(2, 4, 4)),
            w13_weight_scale=_Param(torch.arange(2, dtype=torch.float32) + 1),
            w2_weight_scale=_Param(torch.arange(2, dtype=torch.float32) + 3),
            quant_method=SimpleNamespace(
                quant_config=SimpleNamespace(weight_block_size=None)
            ),
        )
        original_w13 = experts.w13_weight.data.clone()
        fake_deep_gemm = types.SimpleNamespace(
            transform_sf_into_required_layout=(
                lambda sf, *, mn, k, recipe, num_groups, disable_ue8m0_cast: torch.zeros(
                    (num_groups, mn, max(k // recipe[1], 1)), dtype=torch.int32
                )
            ),
            transform_weights_for_mega_moe=lambda l1, l2: (l1, l2),
        )

        with patch.dict(sys.modules, {"deep_gemm": fake_deep_gemm}):
            mega_moe.build_mega_moe_fp8_experts_weights(
                experts, preserve_runner_layout=True, swap_w13_halves=True
            )

        self.assertTrue(experts._mega_moe_weights_built)
        self.assertEqual(experts._mega_moe_weight_format, "fp8")
        self.assertFalse(experts._mega_moe_fp8_block_quant)
        self.assertTrue(torch.equal(experts.w13_weight.data, original_w13))
        self.assertTrue(
            torch.equal(
                experts.mega_fp8_w13_weight,
                torch.cat((original_w13[:, 4:], original_w13[:, :4]), dim=1),
            )
        )
        self.assertNotEqual(
            experts.mega_fp8_w13_weight.data_ptr(),
            experts.w13_weight.data.data_ptr(),
        )

    def test_transform_sf_reports_group16_dependency_gap(self):
        def fake_transform_sf(*args, **kwargs):
            del args, kwargs
            raise RuntimeError("Unknown SF transformation")

        with self.assertRaisesRegex(RuntimeError, "group16"):
            mega_moe._transform_mega_moe_sf(
                fake_transform_sf,
                torch.ones((2, 256, 2), dtype=torch.float32),
                mn=256,
                k=32,
                group_size=16,
                num_groups=2,
                name="w13",
            )


if __name__ == "__main__":
    unittest.main()