meta-pytorch / dist_moe

# Dist-MoE
Dist-MoE provides speed-of-light distributed Mixture-of-Experts kernels for
NVIDIA SM100 and newer GPUs. Its fused communication-compute pipelines combine
token dispatch, grouped expert GEMMs, and token combine for high-throughput
training and inference.
The `dist_moe` package consumes caller-provided top-k expert IDs and scores,
and owns the communication buffers, activation planning, scratch memory, and
graph-stable PyTorch operations needed to execute routed experts.
## Contents
- Supported modes
- Install
- Run the BF16 quickstart
- Understand the execution model
- Choose a precision and pipeline
- Plan memory
- Manage context lifetime
- Use graph and checkpoint features
- Use advanced features
- Public API
- Read the detailed guides
- Acknowledgements
## Supported modes
| Precision | Pipeline | Training | Inference |
| --- | --- | :---: | :---: |
| BF16 | Standard | Yes | Yes |
| MXFP8 E4M3 | Staged or Mega | Yes | Yes |
| NVFP4 | Staged or Mega | No | Yes |
All supported modes work with non-strict `make_fx`, CUDA graphs, and optional
VMM host-backed scratch. Training modes include forward and backward.
The package also supports prepared MXFP8 and NVFP4 weights, device-selected
activation slots, plain and clamped SwiGLU, fused post-expert RMSNorm, clip
statistics, and optional in-place accumulation into standard
`parameter.grad` buffers.
 Not supported in this release 
Full `torch.compile`, strict export, TLX or hierarchical execution,
synchronous block-scaled backends, MXFP8 E5M2, MXFP4, A8W4, NVFP4 training,
expert borrowing, activation rings, and communication quantization.
## Install
Install from a source checkout into an environment that already provides a
supported NVIDIA driver and CUDA toolkit. With pip:
 [code] 
Or with uv:
 [code] 
The initial release is validated with this runtime:
| Layer | Validated release stack |
| --- | --- |
| Platform | Linux ARM64; NVIDIA Blackwell (SM100 class) |
| Framework | Python 3.12; PyTorch 2.14.0+cu130 |
| GPU runtime | CUDA toolkit 13.0.3.0; NVRTC 13.0.88; CUDA Python 13.4.1; NCCL 2.30.7 |
| Kernel toolchain | Triton 3.8 or 3.9; NVIDIA CUTLASS DSL 4.6.1 |
The PyTorch wheel runtime and the toolkit used for JIT compilation are separate
compatibility requirements. Importing `dist_moe` does not initialize CUDA or a
process group, but execution has no CPU fallback.
### Verify the installation
The complete runtime test file requires at least two supported Blackwell GPUs.
For an editable pip installation:
 [code] 
The equivalent uv workflow is:
 [code] 
This test file covers configuration, kernels, autograd, non-strict `make_fx`,
CUDA graphs, and real two-rank execution. Real VMM tests use a separate isolated
test target.
This project leverages third-party code made available under its own licenses.
See NOTICE and LICENSES.
## BF16 Quickstart
This quickstart shows the complete process lifecycle. The checked-in
BF16 training example extends the same flow with
advanced BF16 options. Run one pro