flashinfer.megamoe

SM90 W4A8 expert-parallel MegaMoE with NVFP4 fused/split and MXFP4 fused kernels. All paths use FP8 E4M3 activations and return BF16 outputs. Execution requires NVLink-connected Hopper GPUs and a symmetric buffer allocated collectively before CUDA graph capture. NVFP4 supports four or eight ranks; MXFP4 supports power-of-two rank counts from one to 32.

Buffer and execution

MegaMoESymmBuffer(group, num_experts, ...[, ...])	Allocate symmetric input, expert-token and output-combine storage.
nvfp4_mega_moe(y, l1_weights, l2_weights, ...)	Run NVFP4 MegaMoE using the fused or split family.
mxfp4_mega_moe(y, l1_weights, l2_weights, ...)	Run fused MXFP4 MegaMoE with chunk-major weights and FP8 activations.

Weight preparation

transform_nvfp4_weights_for_mega_moe(...)	Return braided NVFP4 weights for both fused and split SM90 kernels.
transform_mxfp4_weights_for_mega_moe(...)	Return chunk-major MXFP4 weights for the SM90 fused kernel.