We’re sharing Dist-MoE, open-source library from Meta SuperIntelligence Labs (MSL
@AIatMeta ). Dist-MoE provides fused communication and compute kernels for BF16/MXFP8 training and inference, plus NVFP4 inference.
✅ Fused communication and compute kernels
✅ Full CUDA graph support
✅ Dynamic activation packing across layers into a bounded buffer
✅ Dynamic, capacity-driven recompute
✅ Forward and backward kernels for BF16/MXFP8
✅ No shared-expert branch required
Some highlights from our GB300 NVL72 benchmarks:
• Up to 1.40× speedup over Mixture-of-Kittens (MoK) in our imbalanced, ten-layer MXFP8 training workload across 64 GPUs. Dist-MoE packed activation state across all ten layers with zero capacity-induced recompute.
• 3.09 useful PFLOP/s per GPU (3,090 TFLOP/s) for NVFP4 Mega inference across 16 GPUs.
• 1.86 useful PFLOP/s per GPU (1,856 TFLOP/s) for MXFP8 Staged training across 64 GPUs.
• 1.27 useful PFLOP/s per GPU (1,271 TFLOP/s) for BF16 training across 64 GPUs.
• Sub-linear scaling under imbalance: in NVFP4 inference across 16 GPUs, 8× the received rows on the busiest GPU required 4.82× the iteration time; about 40% lower iteration time per received row.
Full configurations, raw results, and measurement definitions are in the blog.
@levendlee and I also put together an interactive visual explainer covering token dispatch and combine, quantization, forward/backward kernels, activation packing, dynamic recompute, and memory planning.
Blog:
meta-pytorch.org/dist_moe/pr…
Repository:
github.com/meta-pytorch/dist…
TorchTitan (
@PyTorch ) integration:
github.com/pytorch/torchtita…
If you work on MoE training, inference, or GPU kernels, take a look. We welcome feedback and contributions.
#PyTorch #MixtureOfExperts #MoE #OpenSource #TorchTitan #DistMoE