Meta has open-sourced Dist-MoE, a library of fused kernels for running Mixture-of-Experts (MoE) models across many GPUs. The repository on GitHub describes it as "speed-of-light" distributed MoE forward and backward kernels that combine token dispatch, grouped expert GEMMs and token combine into a single communication-compute pipeline for NVIDIA SM100 and newer GPUs (Blackwell).
This is infrastructure, not a model, so the audience is people training or serving large MoE systems. But it matters beyond that niche: MoE is the architecture behind many of the open models we cover, and kernel efficiency decides what those models cost to train and serve. This post summarizes what the library does, its limits, and which numbers are documented versus claimed on social media.
TL;DR: Dist-MoE at a glance
| Question | Answer (from the README) |
|---|---|
| What is it? | Fused distributed MoE kernels: dispatch, grouped expert GEMMs, combine |
| Hardware | NVIDIA SM100 and newer (Blackwell); no CPU fallback |
| Validated stack | Linux ARM64, Python 3.12, PyTorch 2.14.0+cu130, CUDA 13.0.3, NCCL 2.30.7, Triton 3.8 or 3.9, CUTLASS DSL 4.6.1 |
| Precisions | BF16 (train, infer); MXFP8 E4M3 (train, infer); NVFP4 (infer only) |
| Pipelines | Standard (BF16); Staged or Mega (MXFP8, NVFP4) |
| Graph support | Non-strict make_fx, CUDA graphs, activation checkpointing |
| Not supported | Full torch.compile, strict export, MXFP4, NVFP4 training and more |
| Tests | Full test file needs at least two Blackwell GPUs |
What problem does it solve?
In an MoE layer, each token is routed to a few experts that may live on other GPUs. The expensive parts are sending tokens to the experts (dispatch), running small matrix multiplies per expert, and sending results back (combine). Done naively, these are separate steps with GPU idle time between them, and the irregular number of tokens per expert makes it hard to keep memory shapes static.
That last point is the "dropless" problem. Many implementations cap how many tokens each expert accepts and drop the overflow, which changes model quality. Dropless implementations process every token but have dynamic shapes, which clash with CUDA graphs, the mechanism that cuts launch overhead by replaying a recorded sequence of GPU work. A Meta engineer, Edward Yang, said in a post relayed by news aggregators: "Last year we were extremely worried about how to get cudagraphable, dropless mixture of experts. Dist_moe has essentially solved this for us." That is his statement, not a documented guarantee.
Dist-MoE's README addresses it in design: the library owns its communication buffers, activation planning and scratch memory, keeps addresses stable for CUDA graph capture, and handles overflow with a memory planner (including optional host-backed overflow scratch using virtual memory management).
How the execution model works
The caller passes routing decisions (top-k expert IDs and scores) and tokens. Dist-MoE then runs: publish rows, route on the GPU, the first projection (W13), SwiGLU, the second projection (W2), and combine, with communication through context-owned symmetric device memory.
Key design points from the README:
- A context owns the buffers. Creating it is collective: all expert-parallel ranks must create matching contexts in the same order. It is sequentially reusable but not reentrant.
- Activation slots are a soft limit. If saved state exceeds the budget, the planner recomputes eligible intermediates in backward instead of failing.
- Scratch is a hard limit. Exceeding total scratch is fatal: the routing kernel prints required and available rows and traps, invalidating the CUDA context.
- Padding uses expert ID minus one. Inactive routes get ID -1 and score zero, and are excluded from routing, communication and gradients.
- Equal row counts across ranks. Every rank must use the same token count per invocation, or a device-side trap fires.
These hard failure modes are a reminder that this is a low-level library. The memory planner (plan_memory(...).explain()) runs without a process group and prints a readable breakdown, which is the right first step before launching a job.
Supported modes
| Precision | Pipeline | Training | Inference |
|---|---|---|---|
| BF16 | Standard | Yes | Yes |
| MXFP8 E4M3 | Staged or Mega | Yes | Yes |
| NVFP4 | Staged or Mega | No | Yes |
The staged pipeline uses separate distributed W13 and SwiGLU/W2 kernels. The Mega pipeline fuses both projections into a chunked pipeline. Block-scaled paths can use prepared weights: prepare_block_scaled_weight quantizes expert weights once, with 32x32 blocks for MXFP8 training-compatible values or 1x32 for inference, and an option to fuse SwiGLU for Mega inference. Additional features include plain and clamped SwiGLU, fused post-expert RMSNorm, clip statistics, and optional in-place accumulation of weight gradients into standard parameter gradient buffers.
For readers new to the number formats, MXFP8 and NVFP4 are low-precision block-scaled formats that Blackwell tensor cores support natively. Lower precision means more math per second and less memory traffic, at the cost of careful scaling.
The speed claims: sourced, not documented
An aggregator's summary of posts from Meta engineers says:
- 3.09 useful PFLOP/s per GPU (3,090 TFLOP/s) for NVFP4 inference across 16 GPUs.
- 1.86 PFLOP/s per GPU for MXFP8 training across 64 GPUs.
- A 1.40x speedup over Cursor's Mixture-of-Kittens for ten-layer MXFP8 training on 64 GPUs, on GB300 NVL72 hardware.
We did not find these figures in the README, and the benchmark scripts and configurations behind them were not part of what we read. Treat them as developer-reported results on a specific, expensive cluster. The comparison with Mixture-of-Kittens is notable: third-party coverage describes Cursor's earlier open-sourced MoE megakernel as targeting the same Blackwell racks, and the two projects now compete in the same niche. We have not verified either vendor's numbers independently.
What it does not do yet
The README is explicit about the release's gaps: no full torch.compile, no strict export, no TLX or hierarchical execution, no synchronous block-scaled backends, no MXFP8 E5M2, MXFP4 or A8W4, no NVFP4 training, no expert borrowing, activation rings or communication quantization. The initial validation is on Linux ARM64 with Blackwell, which suggests the Grace-Blackwell systems Meta uses. If you run x86 or Hopper, check before assuming it will work; the project says execution has no CPU fallback and targets SM100 or newer.
The tooling is also tied to a particular software stack: PyTorch 2.14.0 with CUDA 13 wheels, a recent NCCL, and NVIDIA's CUTLASS Python DSL. Building that stack yourself is a prerequisite.
Why kernel work decides what models cost
MoE models activate only a fraction of their parameters per token, which is how a model with hundreds of billions of parameters can be served for the cost of a much smaller dense one. That saving only materializes if the system keeps the GPUs busy. Communication between GPUs for dispatch and combine can eat a large share of step time, and uneven token counts per expert leave some GPUs waiting while others work.
Fusing communication with computation, as Dist-MoE does, overlaps these costs. Capturing the whole sequence in a CUDA graph removes much of the launch overhead, which matters most at small batch sizes during inference. Lower-precision formats such as MXFP8 and NVFP4 raise the arithmetic throughput and cut memory traffic. Each piece helps a little; together they are why labs invest engineer-years in kernels.
Layers stacked in sequence, illustrating the fused dispatch, expert and combine stages in Meta Dist-MoE kernels
Questions engineers will ask
Is this only for Meta's models? No. The package takes caller-provided top-k expert IDs and scores, so it plugs into any MoE router. You still have to integrate it with your training framework and handle weight layout and quantization preparation.
Does it replace NCCL or Triton? It builds on them. The validated stack lists NCCL, Triton and CUTLASS DSL as dependencies, and the communication uses symmetric device memory buffers owned by the context.
What about AMD or older NVIDIA GPUs? The README states SM100 and newer. We saw no mention of AMD or Hopper support in this release.
How should I benchmark it? Compare against your current MoE path on identical routing distributions, because performance depends strongly on how balanced the router is. Test balanced and skewed loads, measure both throughput and peak memory, and check the memory planner's output for your configuration.
Who should care
- Teams training MoE models on Blackwell clusters who want a maintained, graph-friendly dropless path, particularly with TorchTitan, which the aggregator summary says Dist-MoE integrates with.
- Inference engine developers interested in NVFP4 expert layers. Our coverage of SGLang's plugin system for vendor kernels and TokenSpeed shows how many kernel backends are competing.
- Everyone else: a signal that the frontier of MoE efficiency is moving to fused, low-precision, graph-captured pipelines. Models like Meta's open Muse Spark and Microsoft's MAI-Base-2 rely on MoE efficiency to be affordable. See also PyTorch 2.15's MoE decoding on Apple silicon for the other end of the hardware spectrum.
How to try it
- Provision a machine with at least two Blackwell GPUs and the validated stack above.
- Install with
pip install .oruv pip install .from a source checkout. - Run the BF16 quickstart with
torchrunon two processes, one per expert-parallel GPU, then the MXFP8 and block-scaled inference examples. - Run
python examples/memory_planning.pyto see how the planner reports capacity before committing to a configuration. - Run
python -m pytest -q tests/test_dist_moe.pyto verify the install on your hardware.
Related reading
- SGLang plugin system for vendor kernels
- TokenSpeed inference engine
- Meta Muse Spark Open
- Microsoft MAI-Base-2 MoE efficiency claim
- PyTorch 2.15 MoE decoding on Apple silicon
- NVIDIA Vera Rubin open designs
- Official: meta-pytorch/dist_moe on GitHub and PyTorch's MXFP8 MoE training write-up
Details are accurate as of October 22, 2026 and come from the project README; performance figures are developer-reported.
