flashinfer.megamoe.transform_mxfp4_weights_for_mega_moe

flashinfer.megamoe.transform_mxfp4_weights_for_mega_moe(l1_weights: Tuple[Tensor, Tensor], l2_weights: Tuple[Tensor, Tensor]) → Tuple[Tensor, Tensor]

Return chunk-major MXFP4 weights for the SM90 fused kernel.

Parameters:

• l1_weights (tuple[torch.Tensor, torch.Tensor]) – Packed UINT8 weights of shape [E, 2 * intermediate, hidden // 2] in Marlin byte order and raw UINT8 E8M0 scale bytes per 32 values in [2**-5, 2**6] for groups with nonzero codes (see the module docs). Gate rows precede up rows.

• l2_weights (tuple[torch.Tensor, torch.Tensor]) – Packed UINT8 weights of shape [E, hidden, intermediate // 2] in Marlin byte order and raw UINT8 E8M0 scale bytes per 32 values in [2**-5, 2**6] for groups with nonzero codes.

Returns:

Contiguous L1 and L2 weight tensors in the kernel’s packed layout, on the input device. Weight preparation runs before graph capture.

Return type:

tuple[torch.Tensor, torch.Tensor]