diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 4510cf683..c6730a161 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -289,8 +289,8 @@ def _run_mega_routed( weight_group_size = getattr(moe.experts, "_mega_moe_weight_group_size", 32) deep_gemm.fp8_fp4_mega_moe( y, - moe.experts.mega_l1_weights, - moe.experts.mega_l2_weights, + _as_deep_gemm_packed_fp4_weights(moe.experts.mega_l1_weights), + _as_deep_gemm_packed_fp4_weights(moe.experts.mega_l2_weights), buf, recipe=(1, 1, weight_group_size), activation="swiglu", @@ -304,6 +304,17 @@ def _run_mega_routed( return y +def _as_deep_gemm_packed_fp4_weights( + weights: tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + data, scale = weights + if data.dtype is torch.float4_e2m1fn_x2: + # DeepGEMM's TVM FFI binding rejects PyTorch's FP4x2 DLPack lanes, while + # the MegaMoE kernel expects the same packed bytes exposed as int8. + data = data.view(torch.int8) + return data, scale + + def _interleave_mega_moe_gate_up(t: torch.Tensor, gran: int = 8) -> torch.Tensor: num_groups, n, *rest = t.shape half = n // 2 diff --git a/test/registered/unit/moe/test_glm_megamoe.py b/test/registered/unit/moe/test_glm_megamoe.py index d9807c182..5df6c0609 100644 --- a/test/registered/unit/moe/test_glm_megamoe.py +++ b/test/registered/unit/moe/test_glm_megamoe.py @@ -136,6 +136,25 @@ class TestGLMMegaMoE(unittest.TestCase): 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) + if __name__ == "__main__": unittest.main()