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

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.

Install

Install from a source checkout into an environment that already provides a supported NVIDIA driver and CUDA toolkit. With pip:

python -m pip install .

Or with uv:

uv venv --python 3.12
source .venv/bin/activate
uv pip install .

The initial release is validated with this runtime:

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:

python -m pip install -e .
python -m pip install pytest
python -m pytest -q tests/test_dist_moe.py

The equivalent uv workflow is:

uv venv --python 3.12
source .venv/bin/activate
uv pip install -e .
uv pip install pytest
python -m pytest -q tests/test_dist_moe.py

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 process per expert-parallel GPU:

These examples run on one host and use the loopback address so rendezvous does not depend on whether the host name is reachable from child processes.

torchrun --standalone --local-addr=127.0.0.1 --nproc-per-node=2 examples/bf16_training.py

import os

import torch
import torch.distributed as dist

import dist_moe


device = torch.device("cuda", int(os.environ["LOCAL_RANK"]))
torch.cuda.set_device(device)
dist.init_process_group("nccl")
ep_group = dist.group.WORLD

T, D, F, K, E_local = 16, 256, 256, 2, 4
rank = dist.get_rank(ep_group)
ep_size = dist.get_world_size(ep_group)
E = E_local * ep_size

config = dist_moe.Config(
    num_local_input_tokens=T,
    hidden_dim=D,
    intermediate_dim=F,
    top_k=K,
    num_experts=E,
    max_moe_layers_per_activation_slot=1,
    device_scratch_capacity_factor=1.0,
    num_activation_slots=1,
)
print(dist_moe.plan_memory(config, ep_size=ep_size, device=device).explain())

context = dist_moe.create_context(
    group=ep_group,
    config=config,
    device=device,
)

try:
    x_TD = torch.randn(
        T,
        D,
        dtype=torch.bfloat16,
        device=device,
        requires_grad=True,
    )
    global_token_ids_T = rank * T + torch.arange(T, device=device)
    topk_expert_ids_TK = (
        global_token_ids_T[:, None] * K
        + torch.arange(K, device=device)[None, :]
    ) % E
    topk_scores_TK = torch.full(
        (T, K),
        1.0 / K,
        dtype=torch.float32,
        device=device,
        requires_grad=True,
    )
    w13_E2FD = torch.randn(
        E_local,
        2,
        F,
        D,
        dtype=torch.bfloat16,
        device=device,
        requires_grad=True,
    )
    w2_EDF = torch.randn(
        E_local,
        D,
        F,
        dtype=torch.bfloat16,
        device=device,
        requires_grad=True,
    )

    output_TD = dist_moe.routed_experts(
        x_TD,
        topk_expert_ids_TK,
        topk_scores_TK,
        w13_E2FD,
        w2_EDF,
        context,
    )
    output_TD.float().sum().backward()
    assert output_TD.shape == x_TD.shape
    assert x_TD.grad is not None
finally:
    context.close()
    dist.destroy_process_group()

Here T is the number of local input tokens, D is the model hidden width, F is the per-expert intermediate width, K is top-k, and E_local is the number of experts owned by this rank. The global expert count is E_local * ep_size.

The W13 weight may have shape [E_local, 2, F, D] or [E_local, 2F, D]. The flattened form [E_local * 2F, D] is also accepted when W2 is [E_local * D, F]. Inputs and weights must be contiguous CUDA BF16 tensors. Expert IDs must be contiguous integers in [0, num_experts); scores must be contiguous BF16 or FP32 tensors.

Understand the execution model

One call has four inputs owned by the caller and one reusable context owned by Dist-MoE:

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

The context owns:

• communication buffers: symmetric device memory used to exchange routed rows and gradients with expert-parallel peers;
• activation buffer: one context-owned virtual address range containing device-backed activation slots and scratch storage;
• activation slots: logical regions that retain forward state until its matching backward;
• device scratch: temporary memory reused by one active Dist-MoE call; and
• optional host-backed scratch: VMM overflow for routing beyond the device scratch capacity.

Forward expert matrix multiplication is called FPROP. Backward computes input gradients with DGRAD and weight gradients with WGRAD. The detailed BF16 and MXFP8 guides show every kernel and saved value.

The context is collective to create, sequentially reusable, and not reentrant. All expert-parallel ranks must create matching contexts in the same order.

Config.num_local_input_tokens is one fixed physical shape, not a maximum. Every rank and every invocation using the context must pass exactly that many rows. If a rank has fewer logical tokens, pad its inputs, assign zero routing scores to padded rows, and slice the returned output. Unequal physical token counts across expert-parallel ranks trigger a device-side trap before routing metadata is generated. That failure terminates the distributed iteration and invalidates the CUDA execution context; it is not a recoverable Python input error.

Choose a precision and pipeline

Start with BF16 unless the model and hardware require block-scaled compute.

BF16

Leave block_scaled=None, which is the default. Activations and weights enter the expert GEMMs in BF16; tensor-core products accumulate in FP32. See BF16 execution.

MXFP8

MXFP8 E4M3 supports training and inference. Activations are quantized after routing, and weights may be prepared once and reused. Run the complete training or inference program with:

The block-scaled paths use the concatenated W13 shape [E_local, 2F, D], named w13_EFD. The BF16 quickstart keeps the gate/up axis explicit as w13_E2FD, with shape [E_local, 2, F, D].

torchrun --standalone --local-addr=127.0.0.1 --nproc-per-node=2 examples/mxfp8_training.py
torchrun --standalone --local-addr=127.0.0.1 --nproc-per-node=2 examples/block_scaled_inference.py

policy = dist_moe.BlockScaledConfig(
    format=dist_moe.BlockScaledFormat.MXFP8_E4M3,
    pipeline="staged",
)
config = dist_moe.Config(
    num_local_input_tokens=T,
    hidden_dim=D,
    intermediate_dim=F,
    top_k=K,
    num_experts=E,
    max_moe_layers_per_activation_slot=1,
    block_scaled=policy,
)

The staged pipeline uses separate distributed W13 and SwiGLU/W2 kernels. The Mega pipeline fuses both projections into a chunked pipeline. They use the same public tensor and memory contracts. See MXFP8 execution.

NVFP4

NVFP4 is inference-only. It requires Config(inference=True), zero activation slots, and prepared weights. The complete NVFP4 inference example runs with:

torchrun --standalone --local-addr=127.0.0.1 --nproc-per-node=2 examples/block_scaled_inference.py

policy = dist_moe.BlockScaledConfig(
    format=dist_moe.BlockScaledFormat.NVFP4,
    pipeline="staged",
)
config = dist_moe.Config(
    num_local_input_tokens=T,
    hidden_dim=D,
    intermediate_dim=F,
    top_k=K,
    num_experts=E,
    max_moe_layers_per_activation_slot=1,
    num_activation_slots=0,
    block_scaled=policy,
    inference=True,
)

See NVFP4 inference.

Plan memory

Plan before creating the context. The complete memory-planning example runs without a process group:

python examples/memory_planning.py

config = dist_moe.Config(
    num_local_input_tokens=4,
    hidden_dim=128,
    intermediate_dim=128,
    top_k=2,
    num_experts=4,
    max_moe_layers_per_activation_slot=2,
    activation_slot_capacity_factor=1.0,
    device_scratch_capacity_factor=1.0,
    num_activation_slots=1,
    vmm=dist_moe.VmmConfig(total_scratch_capacity_factor=2.0),
)
print(dist_moe.plan_memory(config, ep_size=2).explain())

The three capacity controls have different meanings:

Leave both activation-slot controls unset to retain only mandatory layer inputs and recompute all eligible intermediates. Set activation_slot_capacity_factor=1.0 to retain all eligible state when the slot's aggregate routing matches balanced routing. Expert users may instead set an exact activation_slot_bytes; the two controls are mutually exclusive.

An activation-slot overflow is safe and selects dynamic recomputation. A total scratch overflow is fatal: the first routing kernel that knows the final local receive-row count prints the required and available rows, then traps before it publishes unsafe routing metadata. Only the overflowing rank is guaranteed to print that diagnostic; peers may time out or be terminated by the job supervisor. The failed CUDA context cannot be reused. See the memory planner guide for formulas and a complete plan.explain() example.

Reading plan.explain()

Memory planning is pure: it does not allocate a context or initialize communication. The example above makes every part of the plan visible.

Read the result from routing capacity to physical storage:

• Receive capacity: 8 balanced -> 8 device -> 16 total. Balanced routing stays in HBM; VMM covers up to twice that routing load.
• Activation slot: 1 x 16 KiB. Capacity factor 1 retains every eligible intermediate at balanced routing.
• Scratch: 12 KiB device + 12 KiB host. Scratch grows into host backing only after the device budget is exhausted.
• Allocation: 28 KiB device / 40 KiB virtual. Device bytes contain the slot and hot scratch; the virtual range also includes host overflow.

The useful range for one slot is 2 KiB for mandatory inputs, 16 KiB for a balanced full save, and 28 KiB for the maximum executable routing load.

Distributed MoE memory plan:
  balanced/device/total rows: 8/8/16
  device/total scratch capacity factors: 1/2
  activation slots: 1 x 0.000 GiB (16,384 bytes) = 0.000 GiB (16,384 bytes) (balanced capacity factor 1)
  device scratch: 0.000 GiB (12,288 bytes)
  host overflow scratch: 0.000 GiB (12,288 bytes)
  total device buffer: 0.000 GiB (28,672 bytes)
  activation slot minimum/balanced full-save/maximum useful: 0.000 GiB (2,048 bytes) / 0.000 GiB (16,384 bytes) / 0.000 GiB (28,672 bytes)
  virtual address range: 0.000 GiB (40,960 bytes); balanced routing fits without recompute; larger saved-state usage may recompute
  activation slots are a soft limit (recompute on exhaustion); total scratch capacity is a hard limit

The activation slot is a soft budget: exhaustion selects recomputation. Total scratch is a hard bound. Exact byte alignment remains the planner's responsibility.

Manage context lifetime

A context's static mode and an individual call's autograd behavior are different decisions:

The MXFP8 training example demonstrates that evaluation can reuse a training context:

with torch.no_grad():
    dist_moe.routed_experts(
        x_TD,
        topk_expert_ids_TK,
        topk_scores_TK,
        w13_EFD,
        w2_EDF,
        context,
    )

torch.no_grad() does not turn a training context into an inference context. Use Config(inference=True) for an inference-only context. Training, training-context evaluation, and specialized inference have different CUDA graph topologies and must be captured separately.

Call context.reset() only after an aborted schedule. Call context.close() after every eager call, trace, and CUDA-graph replay using the context has finished.

Use graph and checkpoint features

The supported graph boundary includes:

• FakeTensor metadata propagation;
• non-strict make_fx with opaque registered Dist-MoE operations;
• activation checkpointing with effectful planner transitions; and
• CUDA graph capture and replay over stable context-owned addresses.

Full torch.compile, strict export, and general functionalization of mutable activation and communication buffers are not supported. A registered operation does not imply those broader compiler guarantees.

Dynamic activation recomputation remains CUDA-graph compatible because the host launches one fixed sequence. Device predicates decide whether the recompute kernels perform work; the host does not branch on the planner result.

Use advanced features

Prepared block-scaled weights

dist_moe.prepare_block_scaled_weight() creates caller-owned quantized data and scales. The complete block-scaled inference example prepares both expert weights before execution:

prepared_w13 = dist_moe.prepare_block_scaled_weight(
    w13_EFD,
    policy,
    inference=True,
)
prepared_w2 = dist_moe.prepare_block_scaled_weight(
    w2_EDF,
    policy,
    inference=True,
)

Training integrations may refill compatible prepared MXFP8 storage with the out= argument after a new weight unshard.

See MXFP8 execution and NVFP4 inference for format-specific fields.

In-place WGRAD accumulation

ExecutionOptions(inplace_wgrad_accum=True) lets serialized backward calls accumulate W13 and W2 gradients directly into existing standard parameter.grad storage. The first contribution is returned normally; later contributions use graph-visible mutable destinations and return no duplicate weight gradient to autograd. See the complete BF16 training example.

For direct accumulation, Config.wgrad_dtype and the parameter declaration resolve as follows. The examples assume a BF16 parameter.

Only BF16 and FP32 are supported, and W13 and W2 must resolve to the same dtype before either kernel writes. Outside direct accumulation, None uses the compute-weight dtype. A parameter's default grad_dtype is its own dtype; set it to None only when an explicitly selected or pre-existing gradient dtype should take precedence.

options = dist_moe.ExecutionOptions(inplace_wgrad_accum=True)
for _ in range(2):
    output_TD = dist_moe.routed_experts(
        x_TD,
        topk_expert_ids_TK,
        topk_scores_TK,
        w13_E2FD,
        w2_EDF,
        context,
        options=options,
    )
    output_TD.float().sum().backward()

This is an opt-in serialized-writer contract. Read WGRAD accumulation before integrating it with a distributed optimizer or pipeline schedule.

Post-expert processing

ExecutionOptions.experts_output_postprocess accepts either a direct eager callback or a typed RMSNormPostprocess. The typed policy fuses route-wise RMSNorm, router-score weighting, and top-k reduction. See the complete BF16 training example.

options = dist_moe.ExecutionOptions(
    experts_output_postprocess=dist_moe.RMSNormPostprocess(
        eps=1e-6,
        norm_output_dtype=torch.float32,
        output_dtype=torch.bfloat16,
        require_bitwise=False,
        recompute_rstd=True,
    )
)

Callback-free typed policies use the registered graph path. Arbitrary Python callbacks and observers use the eager extension path. See post-expert processing.

Clamped SwiGLU

Set activation="swiglu_clamped" with positive swiglu_alpha and swiglu_limit. The fused kernels use the same parameters in forward, backward, and dynamic recomputation. The complete BF16 training example configures the public policy:

config = dataclasses.replace(
    config,
    activation="swiglu_clamped",
    swiglu_alpha=1.702,
    swiglu_limit=7.0,
)

See the precision-specific execution guides for the exact formula.

Clip statistics

BF16 execution can rewrite caller-owned SwiGLU counters on every forward. The complete BF16 training example supplies the graph-visible destination:

clip_stats_out_3 = torch.empty(3, dtype=torch.float32, device=device)
options = dist_moe.ExecutionOptions(
    swiglu_clip_stats_out_3=clip_stats_out_3,
    swiglu_clip_limit=7.0,
)

Pipeline activation slots

A pipeline framework derives static activation-slot assignments from schedule liveness. The complete BF16 training example selects one resolved physical slot:

context.select_activation_slot(
    activation_slot=1,
    num_moe_layers_in_slot=1,
)

The context saves the selected device scalar for backward. See pipeline activation slots for microbatch and stage-microbatch coloring.

VMM host-backed scratch

VMM can extend device scratch with pinned host memory while preserving one stable virtual address range. The complete memory-planning example resolves both device and host-backed scratch before context creation:

config = dataclasses.replace(
    config,
    device_scratch_capacity_factor=1.0,
    vmm=dist_moe.VmmConfig(
        total_scratch_capacity_factor=2.0,
        prefetch=False,
    ),
)

Saved activations always remain in device memory. Only scratch beyond the device capacity may use host backing. See VMM host-backed scratch.

Public API

The API docstrings are authoritative for individual arguments and validation errors. The guides below explain how those APIs compose.

Read the detailed guides

Start with the execution guide for the selected precision:

• BF16 end-to-end execution
• Asynchronous MXFP8 end-to-end execution
• NVFP4 inference execution

Then use the focused guides as needed:

• Activation and scratch memory planner
• VMM host-backed scratch
• Pipeline activation slots
• Post-expert processing
• WGRAD accumulation
• BF16 kernel design
• MXFP8 kernel design
• NVFP4 kernel design

Acknowledgements

Shikai Li is the primary author of this work, for designing and building its core: the distributed MoE implementation, the SyncFree execution model, the activation planning and memory management layers, and the grouped-GEMM and low-precision kernels. Additional contributions came from Jongsoo Park, Summer Deng, Jianyu Huang, Jie Wang, Xiaozhu Meng, Vijay Thakkar, Jiecao Yu, Edward Yang, and others. We also acknowledge the earlier work on mxfp8 dropless CUDA-Graphable MoE by Daniel Haziza, Luca Wehrstedt, Simon Layton and Driss Guessous. This work is open-sourced by Sanket Purandare, advised by Tianyu Liu, Natalia Gimelshein, and Edward Yang.