PyTorch APIs for High Performance MoE Training and Inference
summary
PyTorch adds GPU-native all_to_all_2d dispatch/combine that returns device-side offsets to avoid host-device syncs and to place per-expert tokens contiguously; a 2D pull‑and‑push prototype overlaps intra‑ and inter‑node kernels (planned for 2.10). A new grouped GEMM with offsets API multiplies dynamic per‑expert token blocks using precompiled CUTLASS kernels and Triton autotuned variants. torchao exposes MXFP8 scaled grouped MM (differentiable) and low‑precision comms; microbenchmarks show 1.6–1.8× op speedups and end‑to‑end gains (Llama4 1.42×, DeepSeekV3 ~1.18×) on single‑node AB200s.
AI-written summary · about
chapters (6)
- 0:00Introduction
- 1:27PyTorch MoE goals
- 3:35Token shuffle with 2D all-to-all
- 9:36Grouped matmul for experts
- 17:05MXFP8 in TorchAo
- 25:08Summary
related talks
10mPyTorch Performance
06m⚡ Speedy but Sturdy PyTorch Optimizers
05mApex
38mMPI Meets Machine Learning: Unlocking PyTorch distributed for scaling AI workloads
22mOptimizing Large MoE Inference on NVIDIA Blackwell: NVFP4, ADP, and DualPipe Strat...
54mKeynote: PyTorch 2.1 Technical Deep Dive
24mInside MoE Optimization: A Profiler-Guided Tour of torch.compile and vLLM
25mPyTorch Symmetric Memory + NCCL Device APIs: A New Path Towards Multi-GP...