Mixture of Experts

Warning

MoE feature availability depends on the GPU, framework, Transformer Engine and dependency versions, and precision recipe. Not every operation or optimization works in every configuration. This guide explains the building blocks; consult the PyTorch API reference and JAX API reference for compatibility and input requirements.

The integrated JAX MoE layer currently requires multiple GPUs (at least two). A single-GPU JAX MoE example is not available through the public API.

Introduction

A Mixture of Experts (MoE) layer replaces a dense feed-forward network with a set of expert networks and a router that sends each token to one or more experts. A token passes through the layer in the following stages:

1. The router scores the experts for each token and selects the top-k of them.

2. Token dispatch gathers the tokens into expert-contiguous order.

3. The grouped MLP (the experts) runs a single batched computation over all expert blocks.

4. Token combine scatters the expert outputs back into the original token order, merging the contributions when a token was sent to more than one expert.

With expert parallelism the experts are sharded across ranks, and an all-to-all dispatch and all-to-all combine take the place of token dispatch and token combine: the dispatch takes the router output directly and delivers each rank’s tokens already grouped by local expert, and the combine returns the outputs to the source rank in the original token order.

Figure 1. The stages of an MoE layer on a single device and with expert parallelism.

Transformer Engine provides a building block for each stage. They are exposed as standalone functions, so they can be assembled into a complete MoE layer or dropped into an existing implementation one piece at a time:

• Router: fused score function and top-k selection, and a fused load-balancing loss.

• Token permutation: token dispatch and combine kernels that move tokens between their original order and the expert-contiguous layout.

• Grouped GEMM: the expert linear layers as one call over expert-contiguous blocks; the grouped MLP fuses activation and supported quantization steps into the expert GEMMs.

• Expert parallelism: all-to-all dispatch and combine for experts sharded across devices.

The single-device example wires the blocks into a complete MoE layer; Expert parallelism covers the sharded case.

Router

The router decides which experts each token is sent to. It applies a score function to the gating logits, selects the top-k experts per token, and produces the two tensors that drive the rest of the layer:

• routing_map - by default, a [num_tokens, num_experts] mask marking the selected experts. Token dispatch uses it to lay the tokens out by expert.

• probs - the routing weight of each selected expert. Token combine uses these as merging weights when a token was routed to more than one expert.

In PyTorch, passing a preallocated topk_indices buffer of shape [num_tokens, topk] to fused_topk_with_score_function selects a compact list of expert indices instead. The buffer must be a contiguous CUDA tensor on the same device as logits, with dtype torch.int16, torch.int32 or torch.int64. For example:

topk_indices = torch.empty(
    (num_tokens, topk), dtype=torch.int32, device=logits.device
)

Pass this buffer as topk_indices=topk_indices and leave routing_map_format at its default value. The function fills the buffer and returns it as its second output without materializing the routing mask. The probs output still has shape [num_tokens, num_experts]. The JAX router does not currently expose this index-output option. Both frontends also support a bit-packed mask through routing_map_format="bitmap_u8". The examples and figure below use the default mask format.

fused_topk_with_score_function runs the score function and the top-k selection in a single differentiable kernel. All internal math runs in FP32, regardless of the logits dtype.

Figure 2. The router scores the experts for each token and keeps the top-k. The selected entries populate routing_map (a 0/1 mask) and probs (the routing weights); all other entries are zero.

Options:

• Score function: softmax, sigmoid or sqrtsoftplus. With softmax, use_pre_softmax selects whether the softmax is applied before or after the top-k.

• Grouped routing: the experts are split into num_groups equal groups. Each group is scored by the sum of its best expert scores, the top group_topk groups are kept, and the top-k experts are chosen only from those groups.

• Expert bias: expert_bias is added to the scores before the top-k selection (see Load balancing).

• Scaling: scaling_factor rescales the returned probabilities.

PyTorchJAX

from transformer_engine.pytorch.router import fused_topk_with_score_function

# logits: [num_tokens, num_experts], produced by the gating (router) projection.
#
# Select the top-k experts for each token and return their routing weights. The
# score function and the top-k selection run in a single fused kernel (all math
# is done in fp32 internally for numerical stability).
score_function = "softmax"  # "softmax", "sigmoid" or "sqrtsoftplus"
probs, routing_map = fused_topk_with_score_function(
    logits,
    topk=2,
    use_pre_softmax=False,  # softmax after top-k; True selects softmax-then-top-k
    num_groups=None,  # set with group_topk to enable grouped (device-limited) routing
    group_topk=None,
    scaling_factor=None,  # optional scalar multiplied into the returned probs
    score_function=score_function,
    expert_bias=None,  # [num_experts] selection bias for sigmoid or sqrtsoftplus
)

# probs:       [num_tokens, num_experts], non-zero only at the selected experts.
#              Pass directly to moe_unpermute as merging_probs.
# routing_map: [num_tokens, num_experts] bool mask, True at the selected experts.
#              Cast to int32 and pass to moe_permute.

Load balancing

fused_moe_aux_loss computes the auxiliary load-balancing loss that penalizes uneven token counts across experts. It takes the per-expert token counts and the dense routing scores (one value per expert, not only the selected top-k), so the loss has a gradient with respect to every expert’s logit. The dense scores are returned by the router functions shown below; add the scaled loss to the training loss.

Count tokens from the original router’s routing_map. The auxiliary-score helpers ignore expert bias and grouped routing, so their returned maps can select different experts. Use the same score function for routing and the auxiliary scores.

expert_bias balances the load without an extra loss term:

• with sigmoid or sqrtsoftplus it is added to the scores only for the top-k selection, so it changes which experts are picked but not the returned routing weights;

• update it between steps: lower it for overloaded experts, raise it for under-used ones.

PyTorchJAX

from transformer_engine.pytorch.router import (
    fused_compute_score_for_moe_aux_loss,
    fused_moe_aux_loss,
)

# Use dense scores, but keep the actual router's map for counts: this helper
# does not apply expert bias or grouped routing.
_, scores = fused_compute_score_for_moe_aux_loss(
    logits,
    topk=2,
    score_function=score_function,
)
tokens_per_expert = routing_map.sum(dim=0)  # [num_experts]

aux_loss = fused_moe_aux_loss(
    scores,
    tokens_per_expert,
    total_num_tokens=logits.shape[0],
    num_experts=logits.shape[1],
    topk=2,
    coeff=1e-2,  # loss weight; add aux_loss to the training loss
)

Token permutation

Token dispatch moves the tokens into the expert-contiguous layout expected by the grouped GEMM, and token combine moves the expert outputs back.

• The PyTorch kernels shown below are differentiable.

• The PyTorch snippets below use the mask-map routing variant. Other variants (for example index-map routing) follow the same pattern, see the PyTorch API reference.

Token dispatch

Token dispatch takes the token tensor and a routing map describing each token’s destination experts, and returns a permuted token buffer in which all rows assigned to the same expert are stored contiguously. This is the layout the grouped GEMM consumes, together with the per-expert token counts.

Figure 3. Token dispatch consumes the input token tensor together with the routing map and produces an expert-contiguous token tensor; rows assigned to the same expert are stored back-to-back.

A typical call looks like:

PyTorchJAX

import torch
from transformer_engine.pytorch import moe_permute

# tokens:      [num_tokens, hidden_size]
# routing_map: [num_tokens, num_experts] mask, 1 if token routed to expert
#
# num_out_tokens is the number of rows in the permuted buffer. Reading it from
# the routing map (``int(routing_map.sum())``) triggers a device-to-host sync;
# when the value is known statically (e.g. ``num_tokens * top_k`` for dropless
# routing), prefer passing that constant directly.
permuted, row_id_map = moe_permute(
    tokens,
    routing_map,
    num_out_tokens=int(routing_map.sum()),
)

# permuted:    [num_out_tokens, hidden_size], expert-contiguous
# row_id_map:  opaque tensor used by ``moe_unpermute`` to reverse the
#              permutation after the experts have run.

The PyTorch call returns the permuted token buffer of shape [num_out_tokens, hidden_size] together with a row_id_map that token combine uses to restore the original token order after the experts have run.

Token combine

Token combine is the inverse operation: it takes the expert-contiguous output of the grouped GEMM and the row_id_map returned by token dispatch, and returns a tensor of shape [num_tokens, hidden_size] with the rows written back into the original token order.

Pass the routing weights as merging_probs; the kernel then computes the weighted sum of the per-expert contributions in the same fused pass. Without them the contributions are summed unweighted. They may be omitted for top-1 only when the router guarantees that every selected weight is one, for example post-top-k softmax with scaling_factor=1. Pre-softmax, sigmoid, sqrtsoftplus, or an additional scaling factor can produce non-unit top-1 weights that must still be applied.

Figure 4. Token combine reads the expert-contiguous output tensor and the row_id_map, and writes each row back to its original token slot. With merging_probs, contributions from multiple experts to the same token are combined in the same fused kernel.

A typical call looks like:

PyTorchJAX

from transformer_engine.pytorch import moe_unpermute

# expert_out:    [num_out_tokens, hidden_size], expert-contiguous,
#                produced by an unweighted GroupedLinear or expert MLP.
# row_id_map:    returned by moe_permute.
# merging_probs: [num_tokens, num_experts]; routing probabilities used to
#                weight the per-expert contributions to each token. They may
#                be omitted when every selected routing weight is one or the
#                expert outputs are already weighted, as in the scaled grouped MLP.
# restore_shape: the original [num_tokens, hidden_size]. Required when the
#                permuted buffer has more rows than the input (top-k routing);
#                for top-1 it can be omitted and is inferred from expert_out.
tokens_out = moe_unpermute(
    expert_out,
    row_id_map,
    merging_probs=merging_probs,
    restore_shape=(num_tokens, hidden_size),
)

# tokens_out: [num_tokens, hidden_size], in the original token order

Token probabilities

The routing weights must be applied exactly once:

• In token combine. Pass the original routing weights as merging_probs, as in the examples above.

• Before token combine. Pass the weights to token dispatch as probs to receive them in expert-contiguous order, multiply the completed expert outputs by the permuted weights, and call token combine without merging_probs. An implementation may instead fold the weights into a mathematically equivalent point, such as immediately before the expert’s final bias-free linear layer.

Token dispatch only permutes probs; it does not apply them. In particular, scaling the input to a nonlinear expert is not equivalent to weighting that expert’s output.

Padding and alignment

Grouped GEMM backends require or prefer each expert’s token block to start at an aligned offset (for example, a multiple of 128 rows). Token dispatch can pad each block up to a multiple of align_size in the same kernel.

Figure 5. Each expert’s block is rounded up to a multiple of align_size. The per-expert padding offsets are returned so that token combine can drop the padding again.

The padded dispatch returns the padded token buffer, the aligned per-expert token counts to pass to the grouped GEMM, and the per-expert pad_offsets that token combine needs to remove the padding.

PyTorchJAX

from transformer_engine.pytorch import moe_permute_and_pad_with_probs, moe_unpermute

# tokens:      [num_tokens, hidden_size]
# probs:       [num_tokens, num_experts] routing probabilities
# routing_map: [num_tokens, num_experts] int32 mask
#
# Pad each expert's token block up to a multiple of align_size (here 256) so the
# grouped GEMM sees aligned blocks. Permutation and padding happen in one kernel.
tokens_per_expert = routing_map.sum(dim=0)  # [num_experts]
padded, permuted_probs, row_id_map, pad_offsets, padded_tokens_per_expert = (
    moe_permute_and_pad_with_probs(
        tokens,
        probs,
        routing_map,
        tokens_per_expert,
        align_size=256,
    )
)

# padded:                   [sum(padded_tokens_per_expert), hidden_size]
# pad_offsets:              per-expert cumulative padding (None if already aligned)
# padded_tokens_per_expert: aligned per-expert counts -> m_splits for GroupedLinear

# ... run an unweighted expert block on `padded`, producing expert_out ...

# Apply weights once. Skip this multiplication for the scaled grouped MLP,
# which already weights its outputs through ScaledSwiGLU and scale_bias=True.
expert_out = expert_out * permuted_probs[:, None]

# Combine already weighted outputs without merging_probs. Pass pad_offsets to
# remove padding and restore_shape for the original [num_tokens, hidden_size].
output = moe_unpermute(
    expert_out,
    row_id_map,
    restore_shape=tokens.shape,
    pad_offsets=pad_offsets,
)

Reordering expert chunks

sort_chunks_by_index reorders whole blocks of rows:

• the input [num_tokens, hidden_size] is split along the first dimension into chunks of the given split_sizes;

• the chunks are concatenated again in the order given by sorted_indices: output chunk i is input chunk sorted_indices[i], rows inside a chunk keep their order;

• the operation is differentiable, and a _with_probs variant moves a per-row probability tensor along with the rows.

The typical use is expert parallelism over a generic all-to-all, where the received buffer is ordered by source rank and then by expert, while the grouped GEMM needs all rows of one expert together. With two source ranks and two local experts:

• received chunks: (rank 0, E4), (rank 0, E5), (rank 1, E4), (rank 1, E5);

• sorted_indices = [0, 2, 1, 3] regroups them into E4, E4, E5, E5 for the grouped GEMM;

• after the experts have run, the inverse permutation restores the rank-major order for the combine all-to-all.

Grouped GEMM

The grouped GEMM applies the per-expert linear layers in one call, replacing a loop of one Linear call per expert and producing the same outputs.

Let G be the number of experts. For expert i, X_i is the routed token block, W_i is the expert weight, and b_i is the optional bias:

\[Y_i = X_i W_i^T + b_i,\quad i = 0, \ldots, G - 1\]

The full layer output is the concatenation of all expert outputs:

\[Y = \mathrm{concat}(Y_0, Y_1, \ldots, Y_{G-1})\]

The number of token rows belonging to each expert is passed as a per-expert token-count argument.

Figure 6. Both paths produce the same outputs from the same inputs. The baseline launches one Linear per expert; the grouped GEMM replaces the loop with one call.

Grouped GEMMs can use low-precision training recipes: the inputs are quantized per expert and the expert GEMMs run in the recipe’s precision.

The computation can be implemented in two ways:

• Per-expert GEMMs. The input is split by expert and each expert’s matrix multiplication is launched separately.

• Single grouped GEMM. The experts run in one grouped operation. Keeping token counts on the device avoids a host synchronization and helps with CUDA graph capture.

The snippets assume the tokens have already been permuted into expert-contiguous order.

PyTorchJAX

import torch
import transformer_engine.pytorch as te

# x:             [sum(m_splits), hidden_size], expert-contiguous tokens
# m_splits:      list[int] of length num_experts; m_splits[i] is the number
#                of routed tokens for expert i
# torch_experts: list[torch.nn.Linear] of length num_experts, one per expert
#                (used only by the baseline loop below)
x_by_expert = torch.split(x, m_splits, dim=0)

# Baseline: one Linear call per expert.
loop_out = torch.cat(
    [expert(x_i) for expert, x_i in zip(torch_experts, x_by_expert)],
    dim=0,
)

# Transformer Engine: one grouped linear call.
grouped_linear = te.GroupedLinear(
    num_experts,
    hidden_size,
    ffn_hidden_size,
    bias=True,
    params_dtype=torch.bfloat16,
).cuda()
grouped_out = grouped_linear(x, m_splits)

# Keep token counts on the device with use_grouped_tensor=True.
grouped_linear = te.GroupedLinear(
    num_experts,
    hidden_size,
    ffn_hidden_size,
    bias=True,
    params_dtype=torch.bfloat16,
    use_grouped_tensor=True,
).cuda()
m_splits_dev = torch.tensor(m_splits, dtype=torch.int64, device="cuda")
grouped_out = grouped_linear(x, m_splits_dev)

Grouped MLP

An expert MLP is two grouped GEMMs with an activation between them. Transformer Engine can fuse the activation and quantization into the expert GEMMs. In the fused path illustrated below, FC1 includes the activation and FC2 runs as a separate grouped GEMM, reading the intermediate tensor from device memory.

Figure 7. The operation fuser replaces the two grouped GEMMs and the activation between them with one grouped-MLP operation that launches separate FC1 and FC2 kernels.

The fusion is applied by the operation fuser: a grouped linear, a scaled GLU (or SReLU) activation and another grouped linear in sequence are replaced with one fused grouped-MLP operation. Fusion requires compatible layer dimensions, activation layout, and precision recipe.

Fusion eligibility may change between releases; this guide summarizes the requirements, while the implementation contains the current checks. For example, see the fusion eligibility checks in Transformer Engine v2.19.

For GLU fusion, use glu_interleave_size=32: FC1 must produce alternating blocks of 32 gate features and 32 value features. When loading a checkpoint whose FC1 outputs store all gate features before all value features, convert each expert’s FC1 weight and bias with te.interleave_glu_tensor(tensor, 32). ScaledClampedQGeGLU uses the same layout. For ScaledSReLU, FC1 instead produces ffn_hidden_size features, without the GLU’s doubled width.

ScaledSwiGLU multiplies the activation output by the routing weights before FC2. To weight the entire expert output, FC2 must also scale its bias:

\[p\,(h W_2 + b_2) = (p\,h) W_2 + p\,b_2,\]

where \(h\) is the activation output and \(p\) is the routing weight. Set scale_bias=True on FC2 and pass the same routing weights to both the activation and FC2. With bias=False on FC2, bias scaling is unnecessary.

For grouped MLP fusion, set split_alignment=256 on both GroupedLinear operations. Pad each expert’s input token block to a multiple of 256 and pass the corresponding padded counts as m_splits to both operations. For example, moe_permute_and_pad_with_probs with align_size=256 produces the padded input and counts; setting split_alignment alone does not pad them.

PyTorch

import transformer_engine.pytorch as te

# FC1 produces gate and value features interleaved in blocks of 32.
expert_mlp = te.ops.Sequential(
    te.ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size, split_alignment=256),
    te.ops.ScaledSwiGLU(glu_interleave_size=32),
    te.ops.GroupedLinear(
        num_experts, ffn_hidden_size, hidden_size, scale_bias=True, split_alignment=256
    ),
)

# Inputs from padded dispatch with align_size=256.
expert_out = expert_mlp(
    padded,
    padded_tokens_per_expert,
    permuted_probs,
    padded_tokens_per_expert,
    permuted_probs,
)

The arguments after padded are consumed in operation order: FC1 takes the padded per-expert token counts, the activation takes the routing weights, and FC2 takes the counts and the same weights for its bias. expert_out is already weighted; the subsequent combine must not apply the routing weights again.

Example: MoE layer on a single device

The runnable PyTorch example below wires the blocks together for top-k routing on a single supported NVIDIA GPU: route, dispatch, run the experts, combine. To keep the example short, each expert is represented by one grouped linear layer; a full expert MLP uses the same routing around two grouped linear layers and an activation.

PyTorchJAX

import torch
import transformer_engine.pytorch as te
from transformer_engine.pytorch.router import fused_topk_with_score_function

num_tokens, hidden_size, num_experts, top_k = 16, 64, 4, 2
hidden_states = torch.randn(
    num_tokens,
    hidden_size,
    device="cuda",
    dtype=torch.bfloat16,
    requires_grad=True,
)
gate = torch.nn.Linear(
    hidden_size,
    num_experts,
    bias=False,
    device="cuda",
    dtype=torch.bfloat16,
)
experts = te.GroupedLinear(
    num_experts,
    hidden_size,
    hidden_size,
    bias=False,
    params_dtype=torch.bfloat16,
    device="cuda",
)

# 1. Router: score the experts and pick the top-k for each token.
logits = gate(hidden_states)
probs, routing_map = fused_topk_with_score_function(
    logits,
    topk=top_k,
    use_pre_softmax=False,
    num_groups=None,
    group_topk=None,
    scaling_factor=None,
    score_function="softmax",
    expert_bias=None,
)

# 2. Dispatch: gather tokens into expert-contiguous order.
routing_map = routing_map.to(torch.int32)
permuted, row_id_map = te.moe_permute(
    hidden_states,
    routing_map,
    num_out_tokens=num_tokens * top_k,
)

# 3. Experts: one grouped call over all expert token blocks.
m_splits = routing_map.sum(dim=0).tolist()  # tokens routed to each expert
expert_out = experts(permuted, m_splits)

# 4. Combine: scatter the outputs back and merge the top-k contributions.
output = te.moe_unpermute(
    expert_out,
    row_id_map,
    merging_probs=probs,
    restore_shape=(num_tokens, hidden_size),
)
output.square().mean().backward()

The example uses dropless routing (num_out_tokens = num_tokens * top_k), so the dispatch buffer is sized statically rather than from a device-to-host sync. Every stage is differentiable, so the assembled layer trains end to end.

Expert parallelism

With expert parallelism (EP) the experts are sharded across ranks: every rank keeps its own shard of the tokens and holds only a slice of the experts.

Figure 8. Expert placement: each rank holds its token shard and a subset of the experts. Tokens t1, t3 and t5 are routed to experts on the other rank.

A token routed to an expert on another rank has to travel there and back. Two all-to-all collectives wrap the local expert computation:

• Dispatch sends each token to the rank that owns its expert. It takes the router output (expert indices and weights) directly and delivers a receive buffer grouped by local expert.

• The local grouped GEMM runs on the receive buffer.

• Combine returns the results to the source rank and writes them back in the original token order.

No separate token dispatch or token combine is needed. Shared experts, which every token passes through, are not part of the dispatch: they run as a regular dense MLP on the local tokens on every rank.

Figure 9. Dispatch routes each token to the rank owning its expert, the local experts run on the receive buffer, and combine returns the outputs to the source rank.

Transformer Engine provides dispatch and combine as ready, differentiable operations built on the NCCL EP library (libnccl_ep, loaded at runtime), so an MoE layer with expert parallelism is just router, dispatch, local experts and combine.

Communication and permutation in one step

Over a generic all-to-all the tokens have to be permuted twice: sorted by destination rank before the collective and regrouped by expert after it. The NCCL EP kernels fold both permutations into the transfer:

• dispatch computes the destination slot of every token from the top-k indices and writes it there directly, so no permutation kernel runs before or after the communication;

• combine does the inverse in one pass: it returns each expert output to its source rank and sums it into the original token position.

Figure 10. Dispatch over a generic all-to-all needs a permutation on each side; the NCCL EP dispatch writes every token straight into its expert slot.

Receive buffer

The number of tokens a rank receives depends on the routing. Local expert blocks are packed consecutively, with each block optionally padded to the configured alignment. Their offsets are the cumulative per-expert counts; any remaining capacity forms an unused tail.

• Fixed capacity, recv_capacity_per_rank (an integer): the total buffer size is fixed, while the expert block sizes and offsets depend on the routing.

  • Sizing needs no host synchronization, allowing CUDA graph capture when the expert operations also support it. The low-level zero-copy EP path additionally requires persistent caller-owned receive and gradient buffers.

  • Without per-expert padding, ep_size * max_tokens_per_rank * top_k is a dropless upper bound. With alignment enabled, capacity must also include padding for every local expert. For example, counts of 129 and 127 need 384 rows at alignment 128, even though only 256 tokens arrive. JAX provides get_moe_recv_capacity_per_rank to compute an aligned bound.

  • With balanced routing, a rank receives about max_tokens_per_rank * top_k assignments. A smaller capacity based on this estimate saves memory but can overflow; include alignment padding when sizing it.

  • On overflow the dispatch fails with a device-side error. With drop_on_overflow=True it instead drops the tokens that do not fit; they are not sent, so their experts contribute nothing to the output, like token dropping in capacity-limited MoE.

  • total_recv_tokens reports the required receive rows before dropping, including alignment padding; compare it with capacity after the step to detect an overflow.

• Eager (PyTorch), no capacity given: the buffer is sized from the actual receive count each step.

  • This costs a host synchronization per step and is not CUDA-graph capturable; drop_on_overflow does not apply.

JAX requires a static receive capacity. In the high-level moe API, recv_capacity_per_rank=None selects an aligned dropless worst-case bound; it does not size the buffer from each step’s actual receive count. Pass the same static capacity to ep_bootstrap.

Framework APIs

PyTorchJAX

te.ops.MoeDispatch and te.ops.MoeCombine compose with the local expert MLP in one te.ops.Sequential. The sequence runs dispatch, grouped linear, scaled activation, grouped linear, and combine, with autograd through both the tokens and routing weights.

Initialize NCCL EP once per process with ep_bootstrap(ep_group, ...). Both operations take the same immutable EpConfig and EpBuffer. Their settings must match the bootstrap configuration, including the process group, token limit, receive capacity, and overflow policy. EpBuffer stores the routing handle and per-expert counts until backward completes; use a separate buffer for each concurrently in-flight layer call, including overlapping pipeline microbatches.

Dispatch produces two extra tensors: the per-expert row counts (including alignment padding) and the received routing weights. Named channels pass the counts to both grouped linear operations and the weights to the activation. Bind these channels before constructing or calling the sequence. output_to_caller=False keeps the metadata internal, so the sequence returns only the combined output.

This assembly example assumes an initialized NCCL process group, the current CUDA device set for each rank, and router outputs for the local BF16 tokens. num_experts must be divisible by the EP group size, and max_tokens_per_rank must bound each rank’s input token count. The expert MLP uses BF16 without bias or per-expert padding.

The example uses eager sizing (recv_capacity_per_rank=None) so the receive buffer matches the actual per-expert counts. This synchronizes with the host each step and prevents CUDA graph capture.

import torch
import torch.distributed as dist
import transformer_engine.pytorch as te
from transformer_engine.pytorch.ep import EpBuffer, EpConfig, ep_bootstrap

# ep_group:  process group the experts are sharded over
# tokens:    [num_tokens, hidden_size] bf16 tokens local to this rank
# topk_idx:  [num_tokens, top_k] global expert index per selected expert
# topk_w:    [num_tokens, top_k] fp32 routing weights from the router
ep_size = dist.get_world_size(ep_group)
num_local_experts = num_experts // ep_size
config = EpConfig(
    top_k=top_k,
    hidden_dim=hidden_size,
    num_local_experts=num_local_experts,
    max_tokens_per_rank=max_tokens_per_rank,
    recv_capacity_per_rank=None,
    ep_group=ep_group,
    alignment=0,
)

# Once per process: sets up NCCL EP on ep_group's communicator.
ep_bootstrap(
    ep_group,
    num_experts=num_experts,
    max_tokens_per_rank=config.max_tokens_per_rank,
    hidden_dim=hidden_size,
    num_topk=top_k,
    recv_capacity_per_rank=config.recv_capacity_per_rank,
)
# One buffer per in-flight layer call (e.g. per pipeline microbatch).
buffer = EpBuffer(
    top_k=config.top_k,
    max_tokens_per_rank=config.max_tokens_per_rank,
    recv_capacity_per_rank=config.recv_capacity_per_rank,
    hidden_dim=config.hidden_dim,
    num_local_experts=config.num_local_experts,
    alignment=config.alignment,
    device=tokens.device,
)

dispatch = te.ops.MoeDispatch(config, buffer)
fc1 = te.ops.GroupedLinear(
    num_local_experts,
    hidden_size,
    2 * ffn_hidden_size,
    bias=False,
    device=tokens.device,
    dtype=torch.bfloat16,
)
activation = te.ops.ScaledSwiGLU()
fc2 = te.ops.GroupedLinear(
    num_local_experts,
    ffn_hidden_size,
    hidden_size,
    bias=False,
    device=tokens.device,
    dtype=torch.bfloat16,
)
combine = te.ops.MoeCombine(config, buffer)

# Bind dispatch metadata before constructing or calling the sequence.
dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False)
dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False)
fc1.set_extra_input_channel(0, "tokens_per_expert")
activation.set_extra_input_channel(0, "routing_weights")
fc2.set_extra_input_channel(0, "tokens_per_expert")
moe = te.ops.Sequential(dispatch, fc1, activation, fc2, combine)

# Dispatch takes indices and weights; combine takes the same indices.
output = moe(tokens, topk_idx, topk_w, topk_idx)  # [num_tokens, hidden_size]

The public extra arguments follow operation order: dispatch consumes topk_idx and topk_w, and combine consumes topk_idx again. The grouped linear and activation inputs come from the internal channels. NCCL EP uses the shared buffer’s routing state for combine, but the operation still requires the routing-index argument.

ScaledSwiGLU applies the routing weights before the bias-free FC2, so combine only sums the already weighted expert outputs. For an FC2 with bias, set scale_bias=True and bind its second extra input to the same routing_weights channel, as described in Grouped MLP.

PyTorch dispatch can quantize tokens before sending them, reducing communication volume; expert GEMMs can consume the quantized groups. Communication precision is configured separately from the expert GEMMs: set dispatch_fwd_quant_recipe and combine_bwd_quant_recipe in both EpConfig and EpBuffer for the corresponding communication paths. Use te.autocast to select the expert-compute recipe, and choose an alignment and receive capacity compatible with both the communication and expert kernels. See the configurations in tests/pytorch/distributed/run_ep.py.

The lower-level ep_dispatch and ep_combine functions remain available for custom expert code or explicit receive-buffer management. ep_dispatch returns (recv_tokens, recv_topk_weights, tokens_per_expert); unlike the operation’s extra outputs, the weights precede the counts. Weight the expert outputs before ep_combine, which returns the summed results in the original token order. When operating on a full fixed-capacity buffer, clear unused tail rows and weights before multiplying so uninitialized values cannot introduce NaNs. NCCL EP zeroes alignment padding, and combine ignores slots without a routed token.

Complete runnable examples: examples/pytorch/ep and examples/jax/ep.