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:
@@ -289,8 +289,8 @@ def _run_mega_routed(
|
|||||||
weight_group_size = getattr(moe.experts, "_mega_moe_weight_group_size", 32)
|
weight_group_size = getattr(moe.experts, "_mega_moe_weight_group_size", 32)
|
||||||
deep_gemm.fp8_fp4_mega_moe(
|
deep_gemm.fp8_fp4_mega_moe(
|
||||||
y,
|
y,
|
||||||
moe.experts.mega_l1_weights,
|
_as_deep_gemm_packed_fp4_weights(moe.experts.mega_l1_weights),
|
||||||
moe.experts.mega_l2_weights,
|
_as_deep_gemm_packed_fp4_weights(moe.experts.mega_l2_weights),
|
||||||
buf,
|
buf,
|
||||||
recipe=(1, 1, weight_group_size),
|
recipe=(1, 1, weight_group_size),
|
||||||
activation="swiglu",
|
activation="swiglu",
|
||||||
@@ -304,6 +304,17 @@ def _run_mega_routed(
|
|||||||
return y
|
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:
|
def _interleave_mega_moe_gate_up(t: torch.Tensor, gran: int = 8) -> torch.Tensor:
|
||||||
num_groups, n, *rest = t.shape
|
num_groups, n, *rest = t.shape
|
||||||
half = n // 2
|
half = n // 2
|
||||||
|
|||||||
@@ -136,6 +136,25 @@ class TestGLMMegaMoE(unittest.TestCase):
|
|||||||
torch.equal(experts.mega_l1_weights[0], experts.w13_weight.data)
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user