From eb7f44a8eea6376ac67b1ee0e703296aa3d06b10 Mon Sep 17 00:00:00 2001
From: LuminolT <lumischen01@gmail.com>
Date: Tue, 7 Jul 2026 14:31:03 +0800
Subject: [PATCH] fix(moe): pass megamoe fp4 weights as int8 to deepgemm

DeepGEMM's TVM FFI binding rejects PyTorch FP4x2 DLPack lanes, while the MegaMoE kernel expects the same packed bytes. Convert FP4 packed expert weights to int8 views immediately before fp8_fp4_mega_moe without changing storage or scale tensors.

Tested: python3 -m pytest -q test/registered/unit/moe/test_glm_megamoe.py on remote B300 container (5 passed).

Tested: 1x B300 synthetic forward_mega_moe smoke passed.

Tested: 8x B300 synthetic direct DeepGEMM MegaMoE and forward_mega_moe wrapper smoke passed.
---
 python/sglang/srt/layers/moe/mega_moe.py     | 15 +++++++++++++--
 test/registered/unit/moe/test_glm_megamoe.py | 19 +++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

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()