Dist-MoE: Meta’s Blackwell SM100 MoE Kernels — BF16 / MXFP8 Train, NVFP4 Infer

This morning’s Kolibri-1 note was an open serving MoE — weights, hybrid SWA, a named B200 train recipe, and a ~78 GB FP8 footprint a cluster operator can place. Last night’s Credo Weaver was near-package memory fanout. Lightning V2 on Tenstorrent was on-prem TTS. Olmo-core 3 answered an open MoE training ticket. Tonight is the complementary kernel infra board under routed experts: on October 3, 2026, Meta uploaded dist-moe v0.1.0 to PyPI — CuTe DSL distributed Mixture-of-Experts kernels for NVIDIA SM100+ (Blackwell).

Primary: PyPI dist-moe · GitHub meta-pytorch/dist_moe · docs. Specs below are Meta’s (README / docs / package metadata) unless noted. There are no published throughput benches in this drop — I am not inventing tok/s, TFLOP/s, or rematch tables.

Spec sheet (named numbers)

That is a different purchase sentence from “another open MoE checkpoint” or “another trainer drop.” This one is: can I fuse dispatch, grouped expert GEMM, and combine on Blackwell with a named precision matrix and a context that owns the EP scratch — and what does Meta document when they ship v0.1.0.

What it is selling

Kolibri answered the weights + serve path question. Olmo-core answered the trainer question. Dist-MoE sits one layer down: the kernels that execute routed experts once you already have IDs and scores.

The execution model is explicit. One call takes four caller-owned inputs (activations, top-k expert IDs, scores, expert weights) and one Dist-MoE-owned context. The hot path is:

caller routing
  x_TD + expert IDs + scores
        |
        v
publish rows -> route on GPU -> W13 -> SwiGLU -> W2 -> combine
        |                                           |
        +------ context-owned communication --------+
                    and temporary memory

Forward expert GEMM is FPROP. Backward splits into DGRAD (input grads) and WGRAD (weight grads). BF16 leaves block_scaled=None; activations and weights enter expert GEMMs in BF16 with FP32 accumulate. MXFP8 E4M3 quantizes activations after routing and accepts prepared weights; staged keeps distributed W13 and SwiGLU/W2 separate, Mega fuses both projections into a chunked pipeline. NVFP4 is infer-only and refuses a training context — you must set inference=True, zero activation slots, and prepare weights before the call.

Features Meta documents in this release (without claiming benches): prepared MXFP8/NVFP4 weights; activation slots; plain or clamped SwiGLU; fused post-expert RMSNorm; clip stats; in-place WGRAD accumulation into parameter.grad; VMM host-backed scratch; CUDA graphs + non-strict make_fx. Explicitly not supported: full torch.compile, strict export, NVFP4 training, MXFP8 E5M2, MXFP4, A8W4, expert borrowing, activation rings, communication quantization, interleaved Mega inference weights, TLX / hierarchical execution, synchronous block-scaled backends.

Where it sits next to other boards

Do not collapse this into “another Meta PyTorch release” or an Olmo/Kolibri rematch. Different layer: kernels vs trainer vs weights. Dist-MoE assumes you already routed; it does not replace Megatron-Core schedules or ship a checkpoint.

Caveats

1. ARM64 + Blackwell-only validated stack. Meta’s documented runtime is Linux ARM64, SM100-class Blackwell, Python 3.12, PyTorch 2.14.0+cu130, CUDA 13.0.3.0 toolkit, and the NCCL / Triton / CUTLASS DSL pins above. Other hosts are outside the published validation box.
2. Fixed token shape. Config.num_local_input_tokens is a physical shape, not a soft max. Every rank and every invocation must pass exactly that many rows — pad and zero-score if you have fewer. Unequal physical counts across EP ranks trigger a device-side trap that terminates the iteration and invalidates the CUDA context; it is not a recoverable Python error.
3. No CPU fallback. Import does not initialize CUDA or a process group, but execution has no CPU path. Tests require at least two supported Blackwell GPUs.
4. NVFP4 train unsupported. NVFP4 is inference-only in this release. Training stays on BF16 Standard or MXFP8 E4M3 Staged/Mega.
5. Not a Megatron / Olmo-core rematch. Dist-MoE is the fused kernel layer under routed experts. Olmo-core 3 is a trainer stack with B300 systems numbers. Kolibri-1 is open serving weights. Different boards; do not paste Dist-MoE into a tok/s bake-off column Meta did not publish.
6. Context lifetime. Contexts are collective to create, sequentially reusable, and not reentrant. Training-context torch.no_grad() evaluation is not the same as Config(inference=True). CUDA-graph topologies differ across those modes and must be captured separately.
7. Soft vs hard memory. Activation-slot overflow selects recomputation. Total scratch overflow is fatal — the overflowing rank prints required vs available rows and traps before publishing unsafe routing metadata; peers may time out. Plan with plan_memory before create_context.

How I would evaluate

1. Read the BF16 quickstart and the precision guides under docs/ before quoting Staged vs Mega or NVFP4 constraints.
2. Stand up the validated stack (ARM64 + ≥2 Blackwell GPUs + the pinned PyTorch/CUDA/NCCL/Triton/CUTLASS DSL set) and run tests/test_dist_moe.py — configuration, kernels, autograd, non-strict make_fx, CUDA graphs, two-rank execution.
3. Call plan_memory(config, ep_size=...).explain() and watch balanced/device/total rows plus soft activation-slot vs hard scratch limits before you size production EP.
4. Exercise BF16 train → MXFP8 train → NVFP4 infer with prepared weights; confirm inference=True + zero slots on the NVFP4 path.
5. Keep this board separate from Kolibri’s serve residency, Olmo-core’s trainer tok/s, Credo’s beachfront math, and Lightning’s TTS CapEx table — this is SM100 MoE kernel infra.

Bottom line

On October 3, 2026, Meta shipped dist-moe v0.1.0: CuTe DSL distributed MoE kernels for NVIDIA SM100+, fused token dispatch + grouped expert GEMMs + token combine, with caller-provided top-k IDs/scores and Dist-MoE-owned comm buffers, activation planning, and scratch via routed_experts. Precision matrix: BF16 Standard and MXFP8 E4M3 Staged/Mega for train+infer; NVFP4 Staged/Mega for inference only. Validated on Linux ARM64 + Blackwell SM100 with Python 3.12 / PyTorch 2.14.0+cu130 / CUDA 13.0.3.0 / NCCL 2.30.7 / Triton 3.8.0 / CUTLASS DSL 4.6.1. No CPU fallback; fixed local token shape; no published throughput benches in this drop. Primary author acknowledgement: Shikai Li et al.; open-sourced by Sanket Purandare. Next to this morning’s open serving MoE board, tonight is the complementary Blackwell kernel ticket under routed experts — not an Olmo-core trainer rematch and not a Kolibri weights rematch.

Create-only. Specs from Meta’s October 3, 2026 dist-moe v0.1.0 PyPI upload and meta-pytorch/dist_moe README/docs. Update when Meta publishes benches, expands the validated host matrix beyond ARM64+SM100, or adds NVFP4 training / full torch.compile support.

Sources: PyPI dist-moe · GitHub meta-pytorch/dist_moe · docs.

Anpoo Sivanadi, Staff Software Engineer