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.
This commit is contained in:
LuminolT
2026-07-07 14:31:03 +08:00
parent f482a06a51
commit eb7f44a8ee
2 changed files with 32 additions and 2 deletions
@@ -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()