flashinfer.megamoe.transform_nvfp4_weights_for_mega_moe

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

Return braided NVFP4 weights for both fused and split SM90 kernels.

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 E4M3 scale bytes per 16 values with 6 * scale <= 448 (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 E4M3 scale bytes per 16 values with 6 * scale <= 448.

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]