From 36375087a0db18bc7bf4a634060ddcdc44687194 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 4 May 2026 15:00:05 +0200 Subject: [PATCH 01/52] [Docs] Add MoE feature overview Signed-off-by: Pawel Gadzinski --- docs/features/mixture_of_experts.rst | 74 ++++++++++++++++++++++++++++ docs/index.rst | 1 + 2 files changed, 75 insertions(+) create mode 100644 docs/features/mixture_of_experts.rst diff --git a/docs/features/mixture_of_experts.rst b/docs/features/mixture_of_experts.rst new file mode 100644 index 0000000000..91f5089415 --- /dev/null +++ b/docs/features/mixture_of_experts.rst @@ -0,0 +1,74 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Mixture of Experts +================== + +Mixture of Experts (MoE) layers replace a dense feed-forward network with a set +of expert networks and a router that sends each token to one or more experts. +This keeps the activated parameter count per token small while allowing the +model to scale to many more total parameters. + +Efficient MoE execution is mostly about data movement and batching. Tokens must +be dispatched so that tokens assigned to the same expert are contiguous, expert +linear layers must be executed without launching many tiny GEMMs, and expert +outputs must be combined back into the original token order. Transformer Engine +provides optimized building blocks for these steps. + +Grouped Linear +-------------- + +``GroupedLinear`` applies several independent linear transformations in one +grouped GEMM call. In an MoE layer, each group typically corresponds to one +expert, and the input is the concatenation of the token blocks routed to those +experts. The ``m_splits`` argument passed to ``forward`` gives the number of +tokens in each expert block. + +This is equivalent to splitting the input along the token dimension, applying a +separate linear layer to each split, and concatenating the outputs, but it avoids +the overhead of launching one GEMM per expert. This is especially important when +token counts per expert are small or imbalanced. + +The PyTorch module is available as +``transformer_engine.pytorch.GroupedLinear``. The operation-fuser variant is +available as ``transformer_engine.pytorch.ops.GroupedLinear`` for users building +custom fused graphs. + +``GroupedLinear`` supports Transformer Engine's low precision execution paths, +including FP8 autocast where supported. Tensor-parallel weight shapes can be +configured through the module arguments, but MoE dispatch and combine +communications are handled outside of the module. + +Permutation Kernels +------------------- + +Permutation kernels implement the token dispatch and combine stages around the +expert computation. Dispatch takes the original token tensor and a routing map, +then produces an expert-contiguous tensor suitable for grouped GEMMs. Combine +takes the expert outputs, restores the original token order, and optionally +merges multiple expert contributions using router probabilities. + +For PyTorch, Transformer Engine exposes: + +* ``transformer_engine.pytorch.moe_permute`` for token dispatch. +* ``transformer_engine.pytorch.moe_permute_with_probs`` for dispatching tokens + and router probabilities together. +* ``transformer_engine.pytorch.moe_permute_and_pad_with_probs`` for dispatching + with per-expert padding to satisfy alignment requirements. +* ``transformer_engine.pytorch.moe_unpermute`` for combining expert outputs. +* ``transformer_engine.pytorch.moe_sort_chunks_by_index`` and + ``transformer_engine.pytorch.moe_sort_chunks_by_index_with_probs`` for + reordering already chunked token buffers. + +For JAX, the same dispatch/combine pattern is available through +``transformer_engine.jax.permutation.token_dispatch``, +``transformer_engine.jax.permutation.token_combine``, and +``transformer_engine.jax.permutation.sort_chunks_by_index``. + +The permutation APIs support routing maps that describe which experts receive +each token. PyTorch supports both mask-based maps and index-based maps for +``moe_permute``/``moe_unpermute``; probability-aware and padding-aware variants +use mask-based routing maps. JAX dispatch accepts mask-based routing maps and can +compute padding internally when an alignment size is provided. diff --git a/docs/index.rst b/docs/index.rst index 7389553679..6be821fff8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -46,6 +46,7 @@ Transformer Engine documentation features/low_precision_training/index.rst features/other_optimizations/index.rst + features/mixture_of_experts.rst .. toctree:: From c497c81c1989dd21d9751ccf84390ab90a3a132d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 5 May 2026 12:32:04 +0200 Subject: [PATCH 02/52] [Docs] Refine MoE overview, add supporting snippets/figures and JAX API entries - Add code snippets and SVG figures referenced by mixture_of_experts.rst (moe_permute / moe_unpermute / grouped_linear tabbed examples for both PyTorch and JAX) - Add JAX API reference entries for token_dispatch, token_combine and grouped_dense so the cross-references from the MoE page resolve - Make wording framework-neutral where it was PyTorch-only (Grouped GEMM instead of GroupedLinear/grouped linear in shared sections, both m_splits and group_sizes mentioned, figure labels generalized) - Tighten routing-kernel intro: consolidate the redundant "multiple variants exist / see API ref" notes into one paragraph next to the example, and explicitly state that the kernels are differentiable - Sharpen merging_probs explanation (top-1 vs top-k) and explicitly describe what token_dispatch / token_combine return - Snippet cleanups: define previously undefined symbols, drop the JAX probs= argument from the basic example and explain its purpose in a comment, document the ignored permuted_probs / pad_offsets outputs - Reorder MoE entry in the docs/index.rst toctree Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor --- docs/api/jax.rst | 12 ++ docs/features/grouped_linear_jax.py | 30 +++ docs/features/grouped_linear_pytorch.py | 31 +++ docs/features/img/grouped_linear.svg | 150 ++++++++++++++ docs/features/img/moe_permute.svg | 87 ++++++++ docs/features/img/moe_unpermute.svg | 82 ++++++++ docs/features/mixture_of_experts.rst | 259 ++++++++++++++++++------ docs/features/moe_permute_jax.py | 30 +++ docs/features/moe_permute_pytorch.py | 25 +++ docs/features/moe_unpermute_jax.py | 21 ++ docs/features/moe_unpermute_pytorch.py | 21 ++ docs/index.rst | 2 +- 12 files changed, 688 insertions(+), 62 deletions(-) create mode 100644 docs/features/grouped_linear_jax.py create mode 100644 docs/features/grouped_linear_pytorch.py create mode 100644 docs/features/img/grouped_linear.svg create mode 100644 docs/features/img/moe_permute.svg create mode 100644 docs/features/img/moe_unpermute.svg create mode 100644 docs/features/moe_permute_jax.py create mode 100644 docs/features/moe_permute_pytorch.py create mode 100644 docs/features/moe_unpermute_jax.py create mode 100644 docs/features/moe_unpermute_pytorch.py diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 7a31c9d379..9ff29e5498 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -59,3 +59,15 @@ Modules :members: __call__ .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules + + +Mixture of Experts +------------------ +Routing kernels and grouped dense for MoE layers. See +:doc:`Mixture of Experts <../features/mixture_of_experts>` for an overview. + +.. autoapifunction:: transformer_engine.jax.permutation.token_dispatch + +.. autoapifunction:: transformer_engine.jax.permutation.token_combine + +.. autoapifunction:: transformer_engine.jax.dense.grouped_dense diff --git a/docs/features/grouped_linear_jax.py b/docs/features/grouped_linear_jax.py new file mode 100644 index 0000000000..94042cba21 --- /dev/null +++ b/docs/features/grouped_linear_jax.py @@ -0,0 +1,30 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_GROUPED_LINEAR_JAX +import jax.numpy as jnp +from transformer_engine.jax import dense as te_dense + +# x: [sum(group_sizes), hidden_size], expert-contiguous tokens +# kernel: [num_experts, hidden_size, ffn_hidden_size], stacked per-expert weights +# bias: [num_experts, ffn_hidden_size], stacked per-expert biases +# group_sizes: [num_experts] int array; group_sizes[i] is the number of routed +# tokens for expert i +split_indices = jnp.cumsum(group_sizes)[:-1] +x_by_expert = jnp.split(x, split_indices, axis=0) + +# Baseline: one matmul per expert. +loop_out = jnp.concatenate( + [x_i @ kernel_i + bias_i for x_i, kernel_i, bias_i in zip(x_by_expert, kernel, bias)], + axis=0, +) + +# Transformer Engine: one grouped dense call. +grouped_out = te_dense.grouped_dense( + x, + kernel, + group_sizes=group_sizes, + bias=bias, +) +# END_GROUPED_LINEAR_JAX diff --git a/docs/features/grouped_linear_pytorch.py b/docs/features/grouped_linear_pytorch.py new file mode 100644 index 0000000000..f927e0ac9a --- /dev/null +++ b/docs/features/grouped_linear_pytorch.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_GROUPED_LINEAR_PYTORCH +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) +# END_GROUPED_LINEAR_PYTORCH diff --git a/docs/features/img/grouped_linear.svg b/docs/features/img/grouped_linear.svg new file mode 100644 index 0000000000..82c0be436e --- /dev/null +++ b/docs/features/img/grouped_linear.svg @@ -0,0 +1,150 @@ + + + + + + + + + + + + + + + Loop over experts + + + + X0 + + W0 + + b0 + + + + Expert 0 (Linear) + + + Y0 + + + + + X1 + + W1 + + b1 + + + + Expert 1 (Linear) + + + Y1 + + + + + X2 + + W2 + + b2 + + + + Expert 2 (Linear) + + + Y2 + + + + + + + Grouped GEMM + + + + + W0 + + b0 + + W1 + + b1 + + W2 + + b2 + + + + + + + + + + + X0 + X1 + X2 + + + tokens for expert 0 + + tokens for expert 1 + + tokens for expert 2 + + + + + + + + Grouped GEMM + multiple experts + (GroupedLinear / grouped_dense) + + + + tokens per expert + + + + + + + + + + + Y0 + Y1 + Y2 + + + diff --git a/docs/features/img/moe_permute.svg b/docs/features/img/moe_permute.svg new file mode 100644 index 0000000000..ed50b69b6a --- /dev/null +++ b/docs/features/img/moe_permute.svg @@ -0,0 +1,87 @@ + + + + + + + + + + Token Dispatch + + + + tokens + + t0 + + t1 + + t2 + + t3 + + t4 + + t5 + + + + + + + + Token Dispatch + group tokens by + destination expert + + + + + + + routing_map + + + + + row_id_map + + + + + permuted tokens + + t2 + + t3 + + t0 + + t4 + + t1 + + t5 + + + expert 0 + + expert 1 + + expert 2 + + diff --git a/docs/features/img/moe_unpermute.svg b/docs/features/img/moe_unpermute.svg new file mode 100644 index 0000000000..93b44b17ed --- /dev/null +++ b/docs/features/img/moe_unpermute.svg @@ -0,0 +1,82 @@ + + + + + + + + + + Token Combine + + + + expert outputs + + y2 + + y3 + + y0 + + y4 + + y1 + + y5 + + + expert 0 + + expert 1 + + expert 2 + + + + + + + + Token Combine + restore original + token order + + + + + + + row_id_map + + + + + tokens (original order) + + y0 + + y1 + + y2 + + y3 + + y4 + + y5 + + diff --git a/docs/features/mixture_of_experts.rst b/docs/features/mixture_of_experts.rst index 91f5089415..cb1145318d 100644 --- a/docs/features/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts.rst @@ -3,6 +3,8 @@ See LICENSE for license information. +.. _moe-overview: + Mixture of Experts ================== @@ -11,64 +13,199 @@ of expert networks and a router that sends each token to one or more experts. This keeps the activated parameter count per token small while allowing the model to scale to many more total parameters. -Efficient MoE execution is mostly about data movement and batching. Tokens must -be dispatched so that tokens assigned to the same expert are contiguous, expert -linear layers must be executed without launching many tiny GEMMs, and expert -outputs must be combined back into the original token order. Transformer Engine -provides optimized building blocks for these steps. - -Grouped Linear --------------- - -``GroupedLinear`` applies several independent linear transformations in one -grouped GEMM call. In an MoE layer, each group typically corresponds to one -expert, and the input is the concatenation of the token blocks routed to those -experts. The ``m_splits`` argument passed to ``forward`` gives the number of -tokens in each expert block. - -This is equivalent to splitting the input along the token dimension, applying a -separate linear layer to each split, and concatenating the outputs, but it avoids -the overhead of launching one GEMM per expert. This is especially important when -token counts per expert are small or imbalanced. - -The PyTorch module is available as -``transformer_engine.pytorch.GroupedLinear``. The operation-fuser variant is -available as ``transformer_engine.pytorch.ops.GroupedLinear`` for users building -custom fused graphs. - -``GroupedLinear`` supports Transformer Engine's low precision execution paths, -including FP8 autocast where supported. Tensor-parallel weight shapes can be -configured through the module arguments, but MoE dispatch and combine -communications are handled outside of the module. - -Permutation Kernels -------------------- - -Permutation kernels implement the token dispatch and combine stages around the -expert computation. Dispatch takes the original token tensor and a routing map, -then produces an expert-contiguous tensor suitable for grouped GEMMs. Combine -takes the expert outputs, restores the original token order, and optionally -merges multiple expert contributions using router probabilities. - -For PyTorch, Transformer Engine exposes: - -* ``transformer_engine.pytorch.moe_permute`` for token dispatch. -* ``transformer_engine.pytorch.moe_permute_with_probs`` for dispatching tokens - and router probabilities together. -* ``transformer_engine.pytorch.moe_permute_and_pad_with_probs`` for dispatching - with per-expert padding to satisfy alignment requirements. -* ``transformer_engine.pytorch.moe_unpermute`` for combining expert outputs. -* ``transformer_engine.pytorch.moe_sort_chunks_by_index`` and - ``transformer_engine.pytorch.moe_sort_chunks_by_index_with_probs`` for - reordering already chunked token buffers. - -For JAX, the same dispatch/combine pattern is available through -``transformer_engine.jax.permutation.token_dispatch``, -``transformer_engine.jax.permutation.token_combine``, and -``transformer_engine.jax.permutation.sort_chunks_by_index``. - -The permutation APIs support routing maps that describe which experts receive -each token. PyTorch supports both mask-based maps and index-based maps for -``moe_permute``/``moe_unpermute``; probability-aware and padding-aware variants -use mask-based routing maps. JAX dispatch accepts mask-based routing maps and can -compute padding internally when an alignment size is provided. +Transformer Engine provides two complementary groups of MoE building blocks: + +* **Routing kernels** dispatch tokens to experts and combine expert outputs back + into the original token order. +* **Grouped GEMM** primitives execute the expert linear layers efficiently + once tokens are laid out in expert-contiguous blocks. + +Routing Kernels +--------------- + +Transformer Engine provides routing kernels that move tokens between their +original order and the expert-contiguous layout expected by the grouped GEMM +(``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX). These paths use +optimized kernels instead of Python-level gather / sort / cat chains. The +rest of this section focuses on the two core operations - token dispatch and +token combine - because they illustrate the layout transformation used by the +other variants. + +The snippets below show one concrete instance of this pattern: the mask-map +routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` +and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as +``transformer_engine.jax.permutation.token_dispatch`` and +``transformer_engine.jax.permutation.token_combine``. Other routing variants +(for example, index-map routing in PyTorch via ``map_type="index"``) are +available in both frameworks and follow the same pattern; see the +:doc:`PyTorch API reference <../api/pytorch>` and +:doc:`JAX API reference <../api/jax>` for the complete list and signatures. +The mask-map APIs have different framework-specific wrappers, but lower to +the same shared Triton permutation kernels, and both pairs are differentiable +so they can be used directly inside training graphs. + +Token Dispatch +~~~~~~~~~~~~~~ + +Token dispatch is the canonical routing operation: given the original token +tensor and a routing map describing each token's destination expert, it returns +a permuted token buffer in which all rows assigned to the same expert are +stored contiguously. In PyTorch this operation is exposed as ``moe_permute``; +in JAX it is exposed as ``token_dispatch``. This is exactly the layout that +the grouped linear layer consumes via its per-expert token-count argument +(``m_splits`` in PyTorch ``GroupedLinear``, ``group_sizes`` in JAX +``grouped_dense``), so token dispatch followed by the grouped GEMM forms a +typical MoE forward block. + +.. figure:: img/moe_permute.svg + :align: center + :alt: Token dispatch reorders tokens so that all tokens assigned to the same expert are contiguous + + Figure 1: 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: + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_permute_pytorch.py + :language: python + :start-after: # START_MOE_PERMUTE_PYTORCH + :end-before: # END_MOE_PERMUTE_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_permute_jax.py + :language: python + :start-after: # START_MOE_PERMUTE_JAX + :end-before: # END_MOE_PERMUTE_JAX + +Both variants return the permuted token buffer of shape +``[num_out_tokens, hidden_size]`` together with a ``row_id_map`` that +carries enough information for token combine to restore the original token +order once the expert computation is done. Token dispatch and token combine +are typically used as a matched pair around the grouped GEMM call. + +Token Combine +~~~~~~~~~~~~~ + +Token combine is the inverse routing operation: it takes the expert-contiguous +output produced by the grouped GEMM (or any per-expert computation) and the +``row_id_map`` returned by token dispatch, and returns a single tensor of +shape ``[num_tokens, hidden_size]`` with the rows written back into the +original token order. In PyTorch this operation is exposed as +``moe_unpermute``; in JAX it is exposed as ``token_combine``. + +For top-1 routing each token has exactly one expert contribution, so +``merging_probs`` is omitted. For top-k routing pass the per-token expert +weights as ``merging_probs`` and the kernel computes a weighted sum of the +per-expert contributions in the same fused pass; without it the per-expert +contributions are summed unweighted. + +.. figure:: img/moe_unpermute.svg + :align: center + :alt: Token combine restores expert outputs back into the original token order + + Figure 2: 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: + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_unpermute_pytorch.py + :language: python + :start-after: # START_MOE_UNPERMUTE_PYTORCH + :end-before: # END_MOE_UNPERMUTE_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_unpermute_jax.py + :language: python + :start-after: # START_MOE_UNPERMUTE_JAX + :end-before: # END_MOE_UNPERMUTE_JAX + +Together, token dispatch -> grouped MLP -> token combine form the inner block +of an MoE layer: the first kernel lays the tokens out for grouped GEMMs, the +grouped MLP computes all expert outputs (typically using ``GroupedLinear`` in +PyTorch or ``grouped_dense`` in JAX), and the last kernel restores (and +optionally merges) the per-token results. + +Grouped GEMM +------------ + +The straightforward way to apply per-expert linear layers is to loop over the +experts and call a separate ``Linear`` for each one. This is correct, but it +is not the most efficient way to execute many expert GEMMs. + +Transformer Engine provides a grouped GEMM primitive +(``GroupedLinear`` in PyTorch and ``grouped_dense`` in JAX) - an optimized +replacement that produces the same outputs as the loop while using +implementations that are better suited for MoE workloads. + +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: + +.. math:: + + 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: + +.. math:: + + Y = \mathrm{concat}(Y_0, Y_1, \ldots, Y_{G-1}) + +The grouped GEMM is told how many token rows belong to each expert via a +per-expert token-count argument: ``m_splits`` in PyTorch ``GroupedLinear`` and +``group_sizes`` in JAX ``grouped_dense``. + +.. figure:: img/grouped_linear.svg + :align: center + :alt: Comparison between launching one Linear per expert and using one grouped GEMM call for all expert blocks + + Figure 3: Both paths produce the same outputs from the same inputs. The + baseline launches one ``Linear`` per expert, while the grouped GEMM + (``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX) is an optimized + grouped implementation that replaces the loop. + +The following snippets show how to replace the loop with the grouped GEMM. +They assume the tokens have already been permuted into expert-contiguous order. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: grouped_linear_pytorch.py + :language: python + :start-after: # START_GROUPED_LINEAR_PYTORCH + :end-before: # END_GROUPED_LINEAR_PYTORCH + + .. tab:: JAX + + .. literalinclude:: grouped_linear_jax.py + :language: python + :start-after: # START_GROUPED_LINEAR_JAX + :end-before: # END_GROUPED_LINEAR_JAX + +The grouped GEMM uses implementations tuned for grouped expert execution: + +* **Optimized backends:** Transformer Engine selects from several grouped GEMM + backends depending on the framework, datatype, and GPU architecture. This + can be, for example, cuBLAS GEMMs launched on multiple CUDA streams or a + single grouped GEMM kernel, among other backend-specific implementations. +* **Recipe compatibility:** the grouped GEMM is integrated with + Transformer Engine's :doc:`low-precision training stack + `, so the same recipes available to regular + ``Linear`` layers can be used for MoE experts. +* **Fused quantization:** Low-precision grouped GEMM paths can fuse + quantization-related work such as scale computation, casting, and + cast/transpose steps across experts instead of repeating the same work in a + Python loop. diff --git a/docs/features/moe_permute_jax.py b/docs/features/moe_permute_jax.py new file mode 100644 index 0000000000..ea79a5581e --- /dev/null +++ b/docs/features/moe_permute_jax.py @@ -0,0 +1,30 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_PERMUTE_JAX +import jax.numpy as jnp +from transformer_engine.jax import permutation as te_permutation + +# tokens: [num_tokens, hidden_size] +# routing_map: [num_tokens, num_experts] mask, 1 if token routed to expert +# num_out_tokens must be known at trace time when used under ``jit``. +permuted, _, row_id_map, _, group_sizes = te_permutation.token_dispatch( + tokens, + routing_map, + num_out_tokens=int(jnp.sum(routing_map)), +) + +# permuted: [num_out_tokens, hidden_size], expert-contiguous +# group_sizes: [num_experts], per-expert token counts; can be passed directly +# to ``grouped_dense`` as ``group_sizes``. +# row_id_map: opaque tensor used by ``token_combine`` to reverse the permutation. +# +# The two ignored outputs are ``permuted_probs`` and ``pad_offsets``: +# - ``permuted_probs`` (returned only when ``probs=`` is supplied) holds the +# routing probabilities permuted into expert-contiguous order. It is used +# for input-side scaling of expert inputs before the grouped GEMM, as an +# alternative to passing ``merging_probs`` to ``token_combine``. +# - ``pad_offsets`` is only used together with ``align_size`` for fused +# padding to expert-aligned blocks. +# END_MOE_PERMUTE_JAX diff --git a/docs/features/moe_permute_pytorch.py b/docs/features/moe_permute_pytorch.py new file mode 100644 index 0000000000..cb5a4e263d --- /dev/null +++ b/docs/features/moe_permute_pytorch.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_PERMUTE_PYTORCH +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. +# END_MOE_PERMUTE_PYTORCH diff --git a/docs/features/moe_unpermute_jax.py b/docs/features/moe_unpermute_jax.py new file mode 100644 index 0000000000..f7a5d42167 --- /dev/null +++ b/docs/features/moe_unpermute_jax.py @@ -0,0 +1,21 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_UNPERMUTE_JAX +from transformer_engine.jax import permutation as te_permutation + +# expert_out: [num_out_tokens, hidden_size], expert-contiguous, +# produced by grouped_dense (or a grouped MLP). +# row_id_map: returned by token_dispatch. +# router_probs: [num_tokens, num_experts]; the original (un-permuted) routing +# probabilities. Provide for top-k routing to weight per-expert +# contributions in the same fused pass; pass None for top-1. +tokens_out = te_permutation.token_combine( + expert_out, + row_id_map, + merging_probs=router_probs, +) + +# tokens_out: [num_tokens, hidden_size], in the original token order +# END_MOE_UNPERMUTE_JAX diff --git a/docs/features/moe_unpermute_pytorch.py b/docs/features/moe_unpermute_pytorch.py new file mode 100644 index 0000000000..7cefb61621 --- /dev/null +++ b/docs/features/moe_unpermute_pytorch.py @@ -0,0 +1,21 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_UNPERMUTE_PYTORCH +from transformer_engine.pytorch import moe_unpermute + +# expert_out: [num_out_tokens, hidden_size], expert-contiguous, +# produced by GroupedLinear (or a grouped 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. Provide +# for top-k routing; pass None for top-1. +tokens_out = moe_unpermute( + expert_out, + row_id_map, + merging_probs=merging_probs, +) + +# tokens_out: [num_tokens, hidden_size], in the original token order +# END_MOE_UNPERMUTE_PYTORCH diff --git a/docs/index.rst b/docs/index.rst index 6be821fff8..b18fafd40c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -45,8 +45,8 @@ Transformer Engine documentation :caption: Features features/low_precision_training/index.rst - features/other_optimizations/index.rst features/mixture_of_experts.rst + features/other_optimizations/index.rst .. toctree:: From f905bbd64de7ea88b7dfecb0ce1c80a8076a94f8 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 24 Jun 2026 16:39:04 +0200 Subject: [PATCH 03/52] [Docs] Extend MoE overview: router, end-to-end layer, padding, low-precision Add Router (score function + top-k + load-balancing loss) and Putting-it-together sections, plus token-probabilities / padding-and-alignment / chunk-sort subsections and a fused-expert-MLP note. New SVG figures and PyTorch/JAX snippets. Add the router and moe_permute_and_pad_with_probs API reference entries (PyTorch and JAX) and sort_chunks_by_index (JAX). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- docs/api/jax.rst | 8 +- docs/api/pytorch.rst | 12 + docs/features/img/moe_layer.svg | 58 +++++ docs/features/img/moe_padding.svg | 57 +++++ docs/features/img/moe_router.svg | 104 +++++++++ docs/features/mixture_of_experts.rst | 266 ++++++++++++++++++++--- docs/features/moe_layer_jax.py | 33 +++ docs/features/moe_layer_pytorch.py | 41 ++++ docs/features/moe_permute_pad_jax.py | 31 +++ docs/features/moe_permute_pad_pytorch.py | 36 +++ docs/features/moe_unpermute_pytorch.py | 4 + docs/features/router_jax.py | 45 ++++ docs/features/router_pytorch.py | 56 +++++ 13 files changed, 725 insertions(+), 26 deletions(-) create mode 100644 docs/features/img/moe_layer.svg create mode 100644 docs/features/img/moe_padding.svg create mode 100644 docs/features/img/moe_router.svg create mode 100644 docs/features/moe_layer_jax.py create mode 100644 docs/features/moe_layer_pytorch.py create mode 100644 docs/features/moe_permute_pad_jax.py create mode 100644 docs/features/moe_permute_pad_pytorch.py create mode 100644 docs/features/router_jax.py create mode 100644 docs/features/router_pytorch.py diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 9ff29e5498..24fd6d25d2 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -63,11 +63,17 @@ Modules Mixture of Experts ------------------ -Routing kernels and grouped dense for MoE layers. See +Router, routing kernels, and grouped dense for MoE layers. See :doc:`Mixture of Experts <../features/mixture_of_experts>` for an overview. +.. autoapifunction:: transformer_engine.jax.router.fused_topk_with_score_function + +.. autoapifunction:: transformer_engine.jax.router.fused_moe_aux_loss + .. autoapifunction:: transformer_engine.jax.permutation.token_dispatch .. autoapifunction:: transformer_engine.jax.permutation.token_combine +.. autoapifunction:: transformer_engine.jax.permutation.sort_chunks_by_index + .. autoapifunction:: transformer_engine.jax.dense.grouped_dense diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index db86498005..9073b8b18a 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -77,16 +77,28 @@ Recipe availability Mixture of Experts (MoE) functions ---------------------------------- +See :doc:`Mixture of Experts <../features/mixture_of_experts>` for an overview of +how these functions fit together. ``GroupedLinear`` (the grouped GEMM used for the +expert layers) is documented above with the other modules. + .. autoapifunction:: transformer_engine.pytorch.moe_permute .. autoapifunction:: transformer_engine.pytorch.moe_permute_with_probs +.. autoapifunction:: transformer_engine.pytorch.moe_permute_and_pad_with_probs + .. autoapifunction:: transformer_engine.pytorch.moe_unpermute .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index_with_probs +.. autoapifunction:: transformer_engine.pytorch.router.fused_topk_with_score_function + +.. autoapifunction:: transformer_engine.pytorch.router.fused_compute_score_for_moe_aux_loss + +.. autoapifunction:: transformer_engine.pytorch.router.fused_moe_aux_loss + Communication-computation overlap --------------------------------- diff --git a/docs/features/img/moe_layer.svg b/docs/features/img/moe_layer.svg new file mode 100644 index 0000000000..8d80677b9a --- /dev/null +++ b/docs/features/img/moe_layer.svg @@ -0,0 +1,58 @@ + + + + + + + + + + + + + Mixture of Experts layer + + + + Router + + + Token + Dispatch + + + Grouped MLP + (experts) + + + Token + Combine + + + + tokens + + + routing_map + + + permuted + + + expert out + + + output + + + + probs (routing weights) + diff --git a/docs/features/img/moe_padding.svg b/docs/features/img/moe_padding.svg new file mode 100644 index 0000000000..038f2744d8 --- /dev/null +++ b/docs/features/img/moe_padding.svg @@ -0,0 +1,57 @@ + + + + + + + + + + Fused padding to expert-aligned blocks + + + permuted + expert 0 + + + + expert 1 + + expert 2 + + + tokens_per_expert = [3, 1, 2] + + + + moe_permute_and_pad_with_probs + align_size = 2 + + + padded + expert 0 + + + + pad + expert 1 + + pad + expert 2 + + + padded = [4, 2, 2] + diff --git a/docs/features/img/moe_router.svg b/docs/features/img/moe_router.svg new file mode 100644 index 0000000000..7f26288642 --- /dev/null +++ b/docs/features/img/moe_router.svg @@ -0,0 +1,104 @@ + + + + + + + + + + Router: scoring and top-k selection + + + logits + e0 + e1 + e2 + e3 + + t0 + 1.2 + 3.1 + 0.4 + 2.7 + + t1 + 2.9 + 0.8 + 2.2 + 1.0 + + t2 + 0.5 + 1.1 + 3.3 + 2.5 + + + + + + + score function + + top-k + + + + + + + routing_map + e0 + e1 + e2 + e3 + + 0 + 1 + 0 + 1 + + 1 + 0 + 1 + 0 + + 0 + 0 + 1 + 1 + + + probs + e0 + e1 + e2 + e3 + + 0 + .6 + 0 + .4 + + .7 + 0 + .3 + 0 + + 0 + 0 + .6 + .4 + diff --git a/docs/features/mixture_of_experts.rst b/docs/features/mixture_of_experts.rst index cb1145318d..5e73de5d83 100644 --- a/docs/features/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts.rst @@ -13,23 +13,133 @@ of expert networks and a router that sends each token to one or more experts. This keeps the activated parameter count per token small while allowing the model to scale to many more total parameters. -Transformer Engine provides two complementary groups of MoE building blocks: +A token passes through an MoE layer in four stages: + +#. The **router** scores the experts for each token and selects the top-k of + them. +#. **Token dispatch** gathers the tokens into expert-contiguous order. +#. The **grouped MLP** (the experts) runs a single batched computation over all + expert blocks. +#. **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. + +.. figure:: img/moe_layer.svg + :align: center + :alt: The four stages of an MoE layer: router, token dispatch, grouped MLP, token combine + + Figure 1: The four stages of an MoE layer. The router produces the + ``routing_map`` consumed by token dispatch and the ``probs`` used as merging + weights in token combine. + +Transformer Engine provides an optimized 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: + +* The **router** fuses the score function with the top-k selection, and provides + a fused load-balancing loss. +* **Token dispatch and combine** move tokens between their original order and the + expert-contiguous layout using optimized kernels instead of Python-level + gather / sort / concatenate chains. +* **Grouped GEMM** primitives execute the expert linear layers efficiently once + the tokens are laid out in expert-contiguous blocks. + +The :ref:`end-to-end example ` at the bottom of this +page wires the four stages together; the sections in between describe each +building block on its own. + +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`` - 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. + +Transformer Engine fuses the score function and the top-k selection into a single +differentiable kernel, exposed as ``fused_topk_with_score_function`` in both +``transformer_engine.pytorch.router`` and ``transformer_engine.jax.router``. All +internal math runs in FP32 for numerical stability, regardless of the logits +dtype. + +.. figure:: img/moe_router.svg + :align: center + :alt: The router scores experts per token, keeps the top-k, and fills routing_map and probs + + 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. + +The kernel covers the score functions and selection variants used by common MoE +architectures: + +* **Score function:** ``"softmax"`` or ``"sigmoid"`` (the PyTorch API also offers + ``"sqrtsoftplus"``). With softmax, ``use_pre_softmax`` selects whether the + softmax is applied before or after the top-k. +* **Grouped (device-limited) routing:** ``num_groups`` and ``group_topk`` restrict + selection to a subset of expert groups, as in DeepSeek-style routing. +* **Expert bias:** with the sigmoid score function, ``expert_bias`` shifts the + selection without changing the returned weights - the bias-adjustment scheme + used for auxiliary-loss-free load balancing. +* **Scaling:** ``scaling_factor`` rescales the returned probabilities. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: router_pytorch.py + :language: python + :start-after: # START_ROUTER_PYTORCH + :end-before: # END_ROUTER_PYTORCH + + .. tab:: JAX + + .. literalinclude:: router_jax.py + :language: python + :start-after: # START_ROUTER_JAX + :end-before: # END_ROUTER_JAX -* **Routing kernels** dispatch tokens to experts and combine expert outputs back - into the original token order. -* **Grouped GEMM** primitives execute the expert linear layers efficiently - once tokens are laid out in expert-contiguous blocks. +Load balancing +~~~~~~~~~~~~~~~ + +Left unconstrained, a router tends to collapse onto a handful of experts. The +usual remedy is an auxiliary load-balancing loss that rewards spreading tokens +evenly across experts. Transformer Engine computes it with ``fused_moe_aux_loss`` +from the per-expert token counts and the *dense* routing scores - one value per +expert rather than only the selected top-k - so the loss has a gradient with +respect to every expert's logit. Those dense scores come from +``fused_compute_score_for_moe_aux_loss`` in PyTorch, or from +``fused_topk_with_score_function(..., compute_aux_scores=True)`` in JAX. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: router_pytorch.py + :language: python + :start-after: # START_ROUTER_AUX_PYTORCH + :end-before: # END_ROUTER_AUX_PYTORCH + + .. tab:: JAX + + .. literalinclude:: router_jax.py + :language: python + :start-after: # START_ROUTER_AUX_JAX + :end-before: # END_ROUTER_AUX_JAX Routing Kernels --------------- -Transformer Engine provides routing kernels that move tokens between their -original order and the expert-contiguous layout expected by the grouped GEMM -(``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX). These paths use -optimized kernels instead of Python-level gather / sort / cat chains. The -rest of this section focuses on the two core operations - token dispatch and -token combine - because they illustrate the layout transformation used by the -other variants. +Once the router has produced a routing map, the tokens must be moved into the +expert-contiguous layout expected by the grouped GEMM (``GroupedLinear`` in +PyTorch, ``grouped_dense`` in JAX) and, afterwards, moved back. Transformer +Engine provides differentiable kernels for both directions. The rest of this +section focuses on the two core operations - token dispatch and token combine - +because they illustrate the layout transformation used by the other variants. The snippets below show one concrete instance of this pattern: the mask-map routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` @@ -40,9 +150,9 @@ and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as available in both frameworks and follow the same pattern; see the :doc:`PyTorch API reference <../api/pytorch>` and :doc:`JAX API reference <../api/jax>` for the complete list and signatures. -The mask-map APIs have different framework-specific wrappers, but lower to -the same shared Triton permutation kernels, and both pairs are differentiable -so they can be used directly inside training graphs. +The mask-map APIs have different framework-specific wrappers, but lower to the +same shared Triton permutation kernels, and both pairs are differentiable so they +can be used directly inside training graphs. Token Dispatch ~~~~~~~~~~~~~~ @@ -61,7 +171,7 @@ typical MoE forward block. :align: center :alt: Token dispatch reorders tokens so that all tokens assigned to the same expert are contiguous - Figure 1: Token dispatch consumes the input token tensor together with the + 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. @@ -103,13 +213,16 @@ For top-1 routing each token has exactly one expert contribution, so ``merging_probs`` is omitted. For top-k routing pass the per-token expert weights as ``merging_probs`` and the kernel computes a weighted sum of the per-expert contributions in the same fused pass; without it the per-expert -contributions are summed unweighted. +contributions are summed unweighted. In PyTorch, also pass +``restore_shape=(num_tokens, hidden_size)`` whenever the permuted buffer has more +rows than the original tokens (top-k routing); JAX infers the original token +count from the ``row_id_map``. .. figure:: img/moe_unpermute.svg :align: center :alt: Token combine restores expert outputs back into the original token order - Figure 2: Token combine reads the expert-contiguous output tensor and the + 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. @@ -132,11 +245,72 @@ A typical call looks like: :start-after: # START_MOE_UNPERMUTE_JAX :end-before: # END_MOE_UNPERMUTE_JAX -Together, token dispatch -> grouped MLP -> token combine form the inner block -of an MoE layer: the first kernel lays the tokens out for grouped GEMMs, the -grouped MLP computes all expert outputs (typically using ``GroupedLinear`` in -PyTorch or ``grouped_dense`` in JAX), and the last kernel restores (and -optionally merges) the per-token results. +Token probabilities +~~~~~~~~~~~~~~~~~~~~ + +In top-k routing each token contributes to several experts, and those +contributions are recombined using the routing weights. There are two equivalent +places to apply the weights: + +* **At combine (output side).** Pass the routing weights to token combine as + ``merging_probs``; it forms the weighted sum of the per-expert contributions in + the same fused pass. This is the path used in the examples above. +* **At dispatch (input side).** Scale each expert's input by its routing weight + before the grouped GEMM. ``moe_permute_with_probs`` (PyTorch) and the ``probs`` + argument of ``token_dispatch`` (JAX) permute a probability tensor alongside the + tokens, so the weights arrive already aligned with the expert-contiguous + layout. + +Padding and alignment +~~~~~~~~~~~~~~~~~~~~~~~ + +Grouped GEMM backends are most efficient when each expert's token block starts at +an aligned offset (for example, a multiple of 128 rows). Because the number of +tokens routed to an expert is data dependent, the blocks are generally ragged. +Transformer Engine can pad each block up to a multiple of ``align_size`` as part +of the dispatch kernel, avoiding a separate padding pass. + +.. figure:: img/moe_padding.svg + :align: center + :alt: Each expert block is padded up to a multiple of align_size during dispatch + + 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. + +In PyTorch this is ``moe_permute_and_pad_with_probs``; in JAX it is the +``align_size`` argument of ``token_dispatch``. Both return the padded token +buffer, the aligned per-expert token counts (used as ``m_splits`` / +``group_sizes`` for the grouped GEMM), and the per-expert ``pad_offsets`` that +token combine needs in order to remove the padding. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_permute_pad_pytorch.py + :language: python + :start-after: # START_MOE_PERMUTE_PAD_PYTORCH + :end-before: # END_MOE_PERMUTE_PAD_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_permute_pad_jax.py + :language: python + :start-after: # START_MOE_PERMUTE_PAD_JAX + :end-before: # END_MOE_PERMUTE_PAD_JAX + +Reordering expert chunks +~~~~~~~~~~~~~~~~~~~~~~~~~ + +When experts are sharded across devices, the per-expert token blocks often have +to be reordered - for example, to regroup tokens by destination rank before an +all-to-all, or to restore the original grouping afterwards. +``moe_sort_chunks_by_index`` (PyTorch) and ``sort_chunks_by_index`` (JAX) permute +contiguous chunks of a token tensor according to a list of chunk sizes and a +permutation of chunk indices, without falling back to Python-level slicing and +concatenation. ``moe_sort_chunks_by_index_with_probs`` reorders an accompanying +probability tensor in the same call. Grouped GEMM ------------ @@ -171,7 +345,7 @@ per-expert token-count argument: ``m_splits`` in PyTorch ``GroupedLinear`` and :align: center :alt: Comparison between launching one Linear per expert and using one grouped GEMM call for all expert blocks - Figure 3: Both paths produce the same outputs from the same inputs. The + Figure 6: Both paths produce the same outputs from the same inputs. The baseline launches one ``Linear`` per expert, while the grouped GEMM (``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX) is an optimized grouped implementation that replaces the loop. @@ -204,8 +378,50 @@ The grouped GEMM uses implementations tuned for grouped expert execution: * **Recipe compatibility:** the grouped GEMM is integrated with Transformer Engine's :doc:`low-precision training stack `, so the same recipes available to regular - ``Linear`` layers can be used for MoE experts. + ``Linear`` layers - FP8 (delayed, current, and blockwise scaling), MXFP8, and + NVFP4 - can be used for MoE experts. * **Fused quantization:** Low-precision grouped GEMM paths can fuse quantization-related work such as scale computation, casting, and cast/transpose steps across experts instead of repeating the same work in a Python loop. +* **Fused expert MLP:** Through the :doc:`operation-based API + <../examples/op_fuser/op_fuser>`, the two expert GEMMs and the activation + between them can be fused into a single grouped operation on recent + architectures, removing the intermediate round trips to memory. + +The PyTorch ``GroupedLinear`` module also supports the features expected of a +Transformer Engine linear layer - tensor and sequence parallelism, gradient +accumulation fusion, and FP8 weight caching - so it can serve as a drop-in expert +layer. See the :doc:`PyTorch API reference <../api/pytorch>` for the full +signature. + +.. _moe-putting-it-together: + +Putting it together +------------------- + +The building blocks assemble into the four-stage MoE layer from Figure 1: route, +dispatch, run the experts, and combine. The example below wires them together for +top-k routing. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_layer_pytorch.py + :language: python + :start-after: # START_MOE_LAYER_PYTORCH + :end-before: # END_MOE_LAYER_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_layer_jax.py + :language: python + :start-after: # START_MOE_LAYER_JAX + :end-before: # END_MOE_LAYER_JAX + +This 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. The +expert step is the grouped MLP from the previous section; a full expert MLP +stacks two grouped GEMMs around an activation. Every stage is differentiable, so +the assembled layer trains end to end. diff --git a/docs/features/moe_layer_jax.py b/docs/features/moe_layer_jax.py new file mode 100644 index 0000000000..218a53b7d0 --- /dev/null +++ b/docs/features/moe_layer_jax.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_LAYER_JAX +import jax.numpy as jnp +from transformer_engine.jax import permutation as te_permutation +from transformer_engine.jax import dense as te_dense +from transformer_engine.jax.router import fused_topk_with_score_function + +# hidden_states: [num_tokens, hidden_size] +# gate_kernel: [hidden_size, num_experts], the router projection +# kernel, bias: stacked per-expert weights/biases for the grouped GEMM that +# stands in for the expert MLP here (see "Grouped GEMM"). +top_k = 2 +num_tokens = hidden_states.shape[0] + +# 1. Router: score the experts and pick the top-k for each token. +logits = hidden_states @ gate_kernel +probs, routing_map = fused_topk_with_score_function(logits, topk=top_k, score_function="softmax") + +# 2. Dispatch: gather tokens into expert-contiguous order. +permuted, _, row_id_map, _, group_sizes = te_permutation.token_dispatch( + hidden_states, routing_map.astype(jnp.int32), num_out_tokens=num_tokens * top_k, +) + +# 3. Experts: one grouped GEMM over all expert token blocks. +expert_out = te_dense.grouped_dense(permuted, kernel, group_sizes=group_sizes, bias=bias) + +# 4. Combine: scatter the outputs back and merge the top-k contributions. +output = te_permutation.token_combine(expert_out, row_id_map, merging_probs=probs) +# output: [num_tokens, hidden_size] +# END_MOE_LAYER_JAX diff --git a/docs/features/moe_layer_pytorch.py b/docs/features/moe_layer_pytorch.py new file mode 100644 index 0000000000..bf0cecc540 --- /dev/null +++ b/docs/features/moe_layer_pytorch.py @@ -0,0 +1,41 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_LAYER_PYTORCH +import torch +from transformer_engine.pytorch import moe_permute, moe_unpermute +from transformer_engine.pytorch.router import fused_topk_with_score_function + +# hidden_states: [num_tokens, hidden_size] +# gate: torch.nn.Linear(hidden_size, num_experts), the router projection +# experts: the per-expert MLP, built from te.GroupedLinear (see "Grouped GEMM"); +# a full expert MLP stacks two grouped GEMMs around an activation. +top_k = 2 +num_tokens, hidden_size = hidden_states.shape + +# 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 = moe_permute( + hidden_states, routing_map, num_out_tokens=num_tokens * top_k, +) + +# 3. Experts: one grouped MLP 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. +# restore_shape is the original token shape; it is needed whenever the permuted +# buffer has more rows than the input (top-k routing: num_out_tokens > num_tokens). +output = moe_unpermute( + expert_out, row_id_map, merging_probs=probs, restore_shape=(num_tokens, hidden_size), +) +# output: [num_tokens, hidden_size] +# END_MOE_LAYER_PYTORCH diff --git a/docs/features/moe_permute_pad_jax.py b/docs/features/moe_permute_pad_jax.py new file mode 100644 index 0000000000..0bdc95d6f3 --- /dev/null +++ b/docs/features/moe_permute_pad_jax.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_PERMUTE_PAD_JAX +from transformer_engine.jax import permutation as te_permutation + +# tokens: [num_tokens, hidden_size] +# probs: [num_tokens, num_experts] routing probabilities +# routing_map: [num_tokens, num_experts] int32 mask +# +# Passing align_size enables the same fused padding. token_dispatch allocates a +# fixed worst-case buffer (so it stays jit-compatible) and reports the aligned +# per-expert counts together with the padding offsets. +padded, permuted_probs, row_id_map, pad_offsets, tokens_per_expert = te_permutation.token_dispatch( + tokens, + routing_map, + num_out_tokens=num_tokens * top_k, + probs=probs, + align_size=128, +) + +# tokens_per_expert: aligned per-expert counts -> group_sizes for grouped_dense + +# ... run the grouped GEMM on `padded`, producing expert_out ... + +# Pass pad_offsets so token combine removes the padding it added. +output = te_permutation.token_combine( + expert_out, row_id_map, merging_probs=probs, pad_offsets=pad_offsets, +) +# END_MOE_PERMUTE_PAD_JAX diff --git a/docs/features/moe_permute_pad_pytorch.py b/docs/features/moe_permute_pad_pytorch.py new file mode 100644 index 0000000000..8be4b1b9b0 --- /dev/null +++ b/docs/features/moe_permute_pad_pytorch.py @@ -0,0 +1,36 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_PERMUTE_PAD_PYTORCH +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 128) 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=128, + ) +) + +# 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 the grouped MLP on `padded`, producing expert_out ... + +# Pass pad_offsets so token combine removes the padding it added, and +# restore_shape so the result has the original [num_tokens, hidden_size] shape. +output = moe_unpermute( + expert_out, + row_id_map, + merging_probs=probs, + restore_shape=tokens.shape, + pad_offsets=pad_offsets, +) +# END_MOE_PERMUTE_PAD_PYTORCH diff --git a/docs/features/moe_unpermute_pytorch.py b/docs/features/moe_unpermute_pytorch.py index 7cefb61621..72f45af89c 100644 --- a/docs/features/moe_unpermute_pytorch.py +++ b/docs/features/moe_unpermute_pytorch.py @@ -11,10 +11,14 @@ # merging_probs: [num_tokens, num_experts]; routing probabilities used to # weight the per-expert contributions to each token. Provide # for top-k routing; pass None for top-1. +# 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 diff --git a/docs/features/router_jax.py b/docs/features/router_jax.py new file mode 100644 index 0000000000..fa13b90254 --- /dev/null +++ b/docs/features/router_jax.py @@ -0,0 +1,45 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_ROUTER_JAX +from transformer_engine.jax.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. The score function and the top-k +# selection run in a single fused kernel. Most arguments have defaults, so a +# basic call only needs the logits, topk and score_function. +probs, routing_map = fused_topk_with_score_function( + logits, + topk=2, + score_function="softmax", # "softmax" or "sigmoid" +) + +# probs: [num_tokens, num_experts], non-zero only at the selected experts. +# Pass to token_combine as merging_probs. +# routing_map: [num_tokens, num_experts] bool mask. Cast to int32 for token_dispatch. +# END_ROUTER_JAX + + +# START_ROUTER_AUX_JAX +from transformer_engine.jax.router import fused_moe_aux_loss + +# The load-balancing auxiliary loss uses the dense scores over all experts. In +# JAX the same router function returns them when compute_aux_scores=True (the +# bias / grouping / scaling arguments are ignored in this mode). +scores, routing_map = fused_topk_with_score_function( + logits, + topk=2, + score_function="softmax", + compute_aux_scores=True, +) +tokens_per_expert = routing_map.sum(axis=0) # [num_experts] + +aux_loss = fused_moe_aux_loss( + scores, + tokens_per_expert, + topk=2, + coeff=1e-2, # loss weight; add aux_loss to the training loss +) +# END_ROUTER_AUX_JAX diff --git a/docs/features/router_pytorch.py b/docs/features/router_pytorch.py new file mode 100644 index 0000000000..1b99e1f8ce --- /dev/null +++ b/docs/features/router_pytorch.py @@ -0,0 +1,56 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_ROUTER_PYTORCH +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). +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="softmax", # "softmax", "sigmoid" or "sqrtsoftplus" + expert_bias=None, # [num_experts] selection bias, only with score_function="sigmoid" +) + +# 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. +# END_ROUTER_PYTORCH + + +# START_ROUTER_AUX_PYTORCH +from transformer_engine.pytorch.router import ( + fused_compute_score_for_moe_aux_loss, + fused_moe_aux_loss, +) + +# The load-balancing auxiliary loss is computed from the *dense* scores over all +# experts (not from the sparse top-k probs above), so its gradient reaches every +# expert's logit. fused_compute_score_for_moe_aux_loss returns those dense scores +# together with the same routing map. +routing_map, scores = fused_compute_score_for_moe_aux_loss( + logits, + topk=2, + score_function="softmax", +) +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 +) +# END_ROUTER_AUX_PYTORCH From 82ec7d5ece2662610a2a1ad6b6d14bafd85b4c93 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 12:49:56 +0200 Subject: [PATCH 04/52] [Docs] MoE: fused grouped MLP and expert parallelism sections Signed-off-by: Pawel Gadzinski --- docs/api/jax.rst | 2 + docs/features/grouped_mlp_pytorch.py | 21 ++++++ docs/features/img/moe_expert_parallel.svg | 64 ++++++++++++++++ docs/features/img/moe_grouped_mlp.svg | 53 ++++++++++++++ docs/features/mixture_of_experts.rst | 89 +++++++++++++++++++++++ docs/features/moe_expert_parallel_jax.py | 28 +++++++ 6 files changed, 257 insertions(+) create mode 100644 docs/features/grouped_mlp_pytorch.py create mode 100644 docs/features/img/moe_expert_parallel.svg create mode 100644 docs/features/img/moe_grouped_mlp.svg create mode 100644 docs/features/moe_expert_parallel_jax.py diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 24fd6d25d2..d9ec5b270b 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -77,3 +77,5 @@ Router, routing kernels, and grouped dense for MoE layers. See .. autoapifunction:: transformer_engine.jax.permutation.sort_chunks_by_index .. autoapifunction:: transformer_engine.jax.dense.grouped_dense + +.. autoapifunction:: transformer_engine.jax.moe.moe diff --git a/docs/features/grouped_mlp_pytorch.py b/docs/features/grouped_mlp_pytorch.py new file mode 100644 index 0000000000..2495f99959 --- /dev/null +++ b/docs/features/grouped_mlp_pytorch.py @@ -0,0 +1,21 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_GROUPED_MLP_PYTORCH +import transformer_engine.pytorch as te + +# Build the expert MLP from the operation-based API: two grouped linear layers +# with a scaled GLU activation in between. FC1 produces 2 * ffn_hidden_size +# features (gate and value) for the GLU. +expert_mlp = te.ops.Sequential( + te.ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size), + te.ops.ScaledSwiGLU(), # or ScaledClampedQGeGLU; ScaledSReLU for the unary variant + te.ops.GroupedLinear(num_experts, ffn_hidden_size, hidden_size), +) + +# When this sequence runs under a block-scaled recipe (MXFP8 or NVFP4) on a +# Blackwell (SM100) GPU with NVTE_CUTEDSL_FUSED_GROUPED_MLP=1, the operation +# fuser transparently replaces the three ops with a single fused grouped-MLP +# kernel (GroupedMLP_CuTeGEMMGLU). No code change is needed to opt in. +# END_GROUPED_MLP_PYTORCH diff --git a/docs/features/img/moe_expert_parallel.svg b/docs/features/img/moe_expert_parallel.svg new file mode 100644 index 0000000000..886b66d510 --- /dev/null +++ b/docs/features/img/moe_expert_parallel.svg @@ -0,0 +1,64 @@ + + + + + + + + Expert parallelism: all-to-all dispatch and combine + + Rank 0 + Rank 1 + + + →E0 + →E2 + + →E1 + →E3 + tokens + + + + + + all-to-all + dispatch + + + + + E0 + E1 + E2 + E3 + local experts + + + + + + all-to-all + combine + + + + + output + output + + Each token is sent to the rank that owns its expert, computed locally, and returned to its source rank. + diff --git a/docs/features/img/moe_grouped_mlp.svg b/docs/features/img/moe_grouped_mlp.svg new file mode 100644 index 0000000000..e731163ace --- /dev/null +++ b/docs/features/img/moe_grouped_mlp.svg @@ -0,0 +1,53 @@ + + + + + + + + + Fused grouped MLP + + + unfused: three kernels, with the intermediate written to and re-read from HBM + + FC1 + grouped GEMM + + + SwiGLU + + + FC2 + grouped GEMM + + + + HBM + + + + HBM + + + + + op fuser + + + + Fused grouped MLP — single CuTe DSL kernel + FC1 + SwiGLU + FC2, intermediate kept on-chip + Blackwell (SM100), MXFP8 or NVFP4; activation also GeGLU or SReLU + diff --git a/docs/features/mixture_of_experts.rst b/docs/features/mixture_of_experts.rst index 5e73de5d83..9a88024fce 100644 --- a/docs/features/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts.rst @@ -425,3 +425,92 @@ dispatch buffer is sized statically rather than from a device-to-host sync. The expert step is the grouped MLP from the previous section; a full expert MLP stacks two grouped GEMMs around an activation. Every stage is differentiable, so the assembled layer trains end to end. + +Fused grouped MLP +----------------- + +An expert MLP is two grouped GEMMs with an activation between them: the first +projects into the (gated) feed-forward dimension, the activation is applied, and +the second projects back. Running these as separate kernels writes the large +intermediate activation out to HBM and reads it back for the second GEMM, and +re-quantizes it in a separate pass. + +On Blackwell (SM100) GPUs, Transformer Engine can fuse the whole expert MLP - +both grouped GEMMs and the activation - into a single CuTe DSL kernel. The +intermediate stays on chip and the cross-expert quantization is folded into the +GEMMs, removing the HBM round-trip and the extra kernel launches. + +.. figure:: img/moe_grouped_mlp.svg + :align: center + :alt: The two expert grouped GEMMs and the activation between them fused into one kernel + + Figure 7: The operation fuser replaces the first grouped GEMM, the activation, + and the second grouped GEMM with a single fused grouped-MLP kernel that keeps + the intermediate on chip. + +The fusion is exposed through the operation-based API and applied automatically +by the :doc:`operation fuser <../examples/op_fuser/op_fuser>`: when it sees a +grouped linear, a scaled GLU (or SReLU) activation, and another grouped linear in +sequence, it replaces them with one fused grouped-MLP operation. No change to the +forward code is needed to opt in. + +.. literalinclude:: grouped_mlp_pytorch.py + :language: python + :start-after: # START_GROUPED_MLP_PYTORCH + :end-before: # END_GROUPED_MLP_PYTORCH + +The fused path is taken when all of the following hold; otherwise the three ops +run separately and produce identical results: + +* **Architecture:** Blackwell (SM100) with cuDNN frontend 1.23 or newer. +* **Recipe:** a block-scaled low-precision recipe - MXFP8, or NVFP4 with the + randomized Hadamard transform enabled. +* **Opt-in:** the environment variable ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``. +* **Activation:** a scaled ``SwiGLU`` / ``GeGLU`` (gated) or ``SReLU`` (unary), + with feature dimensions aligned to 64 and the token count to 128. + +Expert parallelism +------------------ + +The grouped GEMM keeps all experts on a single device. When the experts no longer +fit there - or to add another dimension of parallelism - they are sharded across +devices, a scheme called expert parallelism (EP). Each device then owns only a +slice of the experts, so a token routed to a non-local expert has to travel to +the device that owns it. + +That data movement is two all-to-all collectives wrapped around the local expert +computation: a **dispatch** all-to-all sends each token to the rank that owns its +expert, the local grouped GEMM runs, and a **combine** all-to-all returns the +results to the source rank. It is the distributed counterpart of the token +dispatch and token combine kernels described above. + +.. figure:: img/moe_expert_parallel.svg + :align: center + :alt: Tokens are exchanged across ranks by all-to-all so each is processed by the rank owning its expert + + Figure 8: With experts sharded across ranks, a dispatch all-to-all routes each + token to the rank owning its expert and a combine all-to-all returns the + outputs to the source rank. + +Transformer Engine provides expert parallelism at two levels: + +* **JAX MoE layer.** ``transformer_engine.jax.moe.moe`` runs the entire layer - + router, dispatch, grouped expert GEMMs, and combine - as a single + differentiable call. Naming a mesh axis with ``ep_axis`` turns the dispatch and + combine steps into ``jax.lax.ragged_all_to_all`` collectives over that axis. + This API is currently experimental. +* **NCCL EP backend.** A common C API (``nvte_ep_dispatch`` / ``nvte_ep_combine`` + and their backward passes, declared in + ``transformer_engine/common/include/transformer_engine/ep.h``) implements the + dispatch and combine all-to-alls directly on NCCL, using NCCL symmetric-memory + windows for zero-copy transfers. It is compiled in with ``NVTE_WITH_NCCL_EP`` + (Hopper or newer, NCCL 2.30.4+) and provides the high-performance communication + path that the framework layers build on. + +A minimal call to the JAX layer, with experts sharded over the ``"ep"`` mesh +axis: + +.. literalinclude:: moe_expert_parallel_jax.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_JAX + :end-before: # END_MOE_EXPERT_PARALLEL_JAX diff --git a/docs/features/moe_expert_parallel_jax.py b/docs/features/moe_expert_parallel_jax.py new file mode 100644 index 0000000000..bce8abdf4f --- /dev/null +++ b/docs/features/moe_expert_parallel_jax.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MOE_EXPERT_PARALLEL_JAX +from transformer_engine.jax.moe import moe # experimental + +# x: [num_tokens, hidden_size] +# gate_kernel:[hidden_size, num_experts] router projection +# wi_0, wi_1: [num_experts, hidden_size, ffn] expert gate / value projections (SwiGLU) +# wo: [num_experts, ffn, hidden_size] expert output projection +# +# moe() runs the whole layer - router, dispatch, grouped expert GEMMs and +# combine - as a single differentiable call. When ep_axis names a mesh axis, +# the dispatch and combine steps become all-to-all collectives over that axis, +# so experts can be sharded across devices (expert parallelism). +output, aux_loss = moe( + x, gate_kernel, wi_0, wi_1, wo, + num_experts=8, + num_experts_per_tok=2, # top-k + activation_type="silu", + score_function="softmax", + aux_loss_coeff=1e-2, # load-balancing loss; 0 disables it + ep_axis="ep", # mesh axis for expert parallelism (None = no EP) +) +# output: [num_tokens, hidden_size] +# aux_loss: scalar load-balancing loss (None when aux_loss_coeff == 0) +# END_MOE_EXPERT_PARALLEL_JAX From f66df65232688825acc82fe8fe4bff56cc54d482 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 12:59:56 +0200 Subject: [PATCH 05/52] [Docs] MoE: split into per-topic pages, match current docs style, document EP APIs Signed-off-by: Pawel Gadzinski --- docs/_static/css/diagram-colors.css | 32 ++ docs/api/jax.rst | 10 +- docs/api/pytorch.rst | 18 +- docs/features/mixture_of_experts.rst | 516 ------------------ .../expert_parallelism/expert_parallelism.rst | 95 ++++ .../img/moe_expert_parallel.svg | 1 - .../moe_expert_parallel_jax.py | 24 + .../moe_expert_parallel_pytorch.py | 40 ++ .../grouped_gemm/grouped_gemm.rst | 140 +++++ .../grouped_gemm}/grouped_linear_jax.py | 0 .../grouped_gemm}/grouped_linear_pytorch.py | 0 .../grouped_gemm}/grouped_mlp_pytorch.py | 0 .../grouped_gemm}/img/grouped_linear.svg | 1 - .../grouped_gemm}/img/moe_grouped_mlp.svg | 5 +- docs/features/mixture_of_experts/index.rst | 16 + .../introduction}/img/moe_layer.svg | 1 - .../introduction/introduction.rst | 50 ++ .../moe_layer/moe_layer.rst | 40 ++ .../moe_layer}/moe_layer_jax.py | 0 .../moe_layer}/moe_layer_pytorch.py | 0 .../router}/img/moe_router.svg | 1 - .../mixture_of_experts/router/router.rst | 86 +++ .../router}/router_jax.py | 0 .../router}/router_pytorch.py | 0 .../routing_kernels}/img/moe_padding.svg | 1 - .../routing_kernels}/img/moe_permute.svg | 1 - .../routing_kernels}/img/moe_unpermute.svg | 1 - .../routing_kernels}/moe_permute_jax.py | 0 .../routing_kernels}/moe_permute_pad_jax.py | 0 .../moe_permute_pad_pytorch.py | 0 .../routing_kernels}/moe_permute_pytorch.py | 0 .../routing_kernels}/moe_unpermute_jax.py | 0 .../routing_kernels}/moe_unpermute_pytorch.py | 0 .../routing_kernels/routing_kernels.rst | 182 ++++++ docs/features/moe_expert_parallel_jax.py | 28 - docs/index.rst | 2 +- 36 files changed, 734 insertions(+), 557 deletions(-) delete mode 100644 docs/features/mixture_of_experts.rst create mode 100644 docs/features/mixture_of_experts/expert_parallelism/expert_parallelism.rst rename docs/features/{ => mixture_of_experts/expert_parallelism}/img/moe_expert_parallel.svg (98%) create mode 100644 docs/features/mixture_of_experts/expert_parallelism/moe_expert_parallel_jax.py create mode 100644 docs/features/mixture_of_experts/expert_parallelism/moe_expert_parallel_pytorch.py create mode 100644 docs/features/mixture_of_experts/grouped_gemm/grouped_gemm.rst rename docs/features/{ => mixture_of_experts/grouped_gemm}/grouped_linear_jax.py (100%) rename docs/features/{ => mixture_of_experts/grouped_gemm}/grouped_linear_pytorch.py (100%) rename docs/features/{ => mixture_of_experts/grouped_gemm}/grouped_mlp_pytorch.py (100%) rename docs/features/{ => mixture_of_experts/grouped_gemm}/img/grouped_linear.svg (99%) rename docs/features/{ => mixture_of_experts/grouped_gemm}/img/moe_grouped_mlp.svg (93%) create mode 100644 docs/features/mixture_of_experts/index.rst rename docs/features/{ => mixture_of_experts/introduction}/img/moe_layer.svg (97%) create mode 100644 docs/features/mixture_of_experts/introduction/introduction.rst create mode 100644 docs/features/mixture_of_experts/moe_layer/moe_layer.rst rename docs/features/{ => mixture_of_experts/moe_layer}/moe_layer_jax.py (100%) rename docs/features/{ => mixture_of_experts/moe_layer}/moe_layer_pytorch.py (100%) rename docs/features/{ => mixture_of_experts/router}/img/moe_router.svg (99%) create mode 100644 docs/features/mixture_of_experts/router/router.rst rename docs/features/{ => mixture_of_experts/router}/router_jax.py (100%) rename docs/features/{ => mixture_of_experts/router}/router_pytorch.py (100%) rename docs/features/{ => mixture_of_experts/routing_kernels}/img/moe_padding.svg (98%) rename docs/features/{ => mixture_of_experts/routing_kernels}/img/moe_permute.svg (98%) rename docs/features/{ => mixture_of_experts/routing_kernels}/img/moe_unpermute.svg (98%) rename docs/features/{ => mixture_of_experts/routing_kernels}/moe_permute_jax.py (100%) rename docs/features/{ => mixture_of_experts/routing_kernels}/moe_permute_pad_jax.py (100%) rename docs/features/{ => mixture_of_experts/routing_kernels}/moe_permute_pad_pytorch.py (100%) rename docs/features/{ => mixture_of_experts/routing_kernels}/moe_permute_pytorch.py (100%) rename docs/features/{ => mixture_of_experts/routing_kernels}/moe_unpermute_jax.py (100%) rename docs/features/{ => mixture_of_experts/routing_kernels}/moe_unpermute_pytorch.py (100%) create mode 100644 docs/features/mixture_of_experts/routing_kernels/routing_kernels.rst delete mode 100644 docs/features/moe_expert_parallel_jax.py diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index f5dc7da4dd..88f7f90f11 100644 --- a/docs/_static/css/diagram-colors.css +++ b/docs/_static/css/diagram-colors.css @@ -279,3 +279,35 @@ html[data-theme="dark"] .subtitle, html[data-theme="dark"] .memory-label { fill: #e0e0e0; } html[data-theme="dark"] .connector { stroke: #bdbdbd; } +/* mixture_of_experts diagrams */ +html[data-theme="dark"] .small-label, +html[data-theme="dark"] .tiny-label, +html[data-theme="dark"] .row-label, +html[data-theme="dark"] .rank-label { fill: #bdbdbd; } +html[data-theme="dark"] .grid-title, +html[data-theme="dark"] .brace-label, +html[data-theme="dark"] .row-label-l, +html[data-theme="dark"] .row-label-r, +html[data-theme="dark"] .mono { fill: #e0e0e0; } +html[data-theme="dark"] .expert-e0 { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .expert-e1 { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .expert-e2 { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .expert-e3 { fill: #5c1f38; stroke: #f48fb1; } +html[data-theme="dark"] .stage, +html[data-theme="dark"] .grouped, +html[data-theme="dark"] .a2a, +html[data-theme="dark"] .fused-mlp { fill: #3a2f5c; stroke: #b39ddb; } +html[data-theme="dark"] .op { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .act, +html[data-theme="dark"] .sel { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .linear, +html[data-theme="dark"] .param, +html[data-theme="dark"] .hbm, +html[data-theme="dark"] .cell, +html[data-theme="dark"] .box { fill: #2b2b2b; stroke: #9e9e9e; } +html[data-theme="dark"] .unsel, +html[data-theme="dark"] .pad { fill: #1f1f1f; stroke: #616161; } +html[data-theme="dark"] .soft-arrow, +html[data-theme="dark"] .brace { stroke: #bdbdbd; } +html[data-theme="dark"] .skip, +html[data-theme="dark"] .farrow { stroke: #b39ddb; } diff --git a/docs/api/jax.rst b/docs/api/jax.rst index d9ec5b270b..daa8639641 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -64,7 +64,7 @@ Modules Mixture of Experts ------------------ Router, routing kernels, and grouped dense for MoE layers. See -:doc:`Mixture of Experts <../features/mixture_of_experts>` for an overview. +:doc:`Mixture of Experts <../features/mixture_of_experts/index>` for an overview. .. autoapifunction:: transformer_engine.jax.router.fused_topk_with_score_function @@ -79,3 +79,11 @@ Router, routing kernels, and grouped dense for MoE layers. See .. autoapifunction:: transformer_engine.jax.dense.grouped_dense .. autoapifunction:: transformer_engine.jax.moe.moe + +.. autoapifunction:: transformer_engine.jax.ep.ep_bootstrap + +.. autoapifunction:: transformer_engine.jax.ep.ep_finalize + +.. autoapifunction:: transformer_engine.jax.ep.ep_dispatch + +.. autoapifunction:: transformer_engine.jax.ep.ep_combine diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 0523853c99..b178b74198 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -90,7 +90,7 @@ Recipe availability Mixture of Experts (MoE) functions ---------------------------------- -See :doc:`Mixture of Experts <../features/mixture_of_experts>` for an overview of +See :doc:`Mixture of Experts <../features/mixture_of_experts/index>` for an overview of how these functions fit together. ``GroupedLinear`` (the grouped GEMM used for the expert layers) is documented above with the other modules. @@ -112,6 +112,22 @@ expert layers) is documented above with the other modules. .. autoapifunction:: transformer_engine.pytorch.router.fused_moe_aux_loss +Expert parallelism +~~~~~~~~~~~~~~~~~~ + +NCCL-based dispatch and combine for experts sharded across ranks. See +:doc:`Expert parallelism <../features/mixture_of_experts/expert_parallelism/expert_parallelism>`. + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_bootstrap + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_finalize + +.. autoapiclass:: transformer_engine.pytorch.ep.EpBuffer + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_dispatch + +.. autoapifunction:: transformer_engine.pytorch.ep.ep_combine + Communication-computation overlap --------------------------------- diff --git a/docs/features/mixture_of_experts.rst b/docs/features/mixture_of_experts.rst deleted file mode 100644 index 9a88024fce..0000000000 --- a/docs/features/mixture_of_experts.rst +++ /dev/null @@ -1,516 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -.. _moe-overview: - -Mixture of Experts -================== - -Mixture of Experts (MoE) layers replace a dense feed-forward network with a set -of expert networks and a router that sends each token to one or more experts. -This keeps the activated parameter count per token small while allowing the -model to scale to many more total parameters. - -A token passes through an MoE layer in four stages: - -#. The **router** scores the experts for each token and selects the top-k of - them. -#. **Token dispatch** gathers the tokens into expert-contiguous order. -#. The **grouped MLP** (the experts) runs a single batched computation over all - expert blocks. -#. **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. - -.. figure:: img/moe_layer.svg - :align: center - :alt: The four stages of an MoE layer: router, token dispatch, grouped MLP, token combine - - Figure 1: The four stages of an MoE layer. The router produces the - ``routing_map`` consumed by token dispatch and the ``probs`` used as merging - weights in token combine. - -Transformer Engine provides an optimized 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: - -* The **router** fuses the score function with the top-k selection, and provides - a fused load-balancing loss. -* **Token dispatch and combine** move tokens between their original order and the - expert-contiguous layout using optimized kernels instead of Python-level - gather / sort / concatenate chains. -* **Grouped GEMM** primitives execute the expert linear layers efficiently once - the tokens are laid out in expert-contiguous blocks. - -The :ref:`end-to-end example ` at the bottom of this -page wires the four stages together; the sections in between describe each -building block on its own. - -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`` - 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. - -Transformer Engine fuses the score function and the top-k selection into a single -differentiable kernel, exposed as ``fused_topk_with_score_function`` in both -``transformer_engine.pytorch.router`` and ``transformer_engine.jax.router``. All -internal math runs in FP32 for numerical stability, regardless of the logits -dtype. - -.. figure:: img/moe_router.svg - :align: center - :alt: The router scores experts per token, keeps the top-k, and fills routing_map and probs - - 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. - -The kernel covers the score functions and selection variants used by common MoE -architectures: - -* **Score function:** ``"softmax"`` or ``"sigmoid"`` (the PyTorch API also offers - ``"sqrtsoftplus"``). With softmax, ``use_pre_softmax`` selects whether the - softmax is applied before or after the top-k. -* **Grouped (device-limited) routing:** ``num_groups`` and ``group_topk`` restrict - selection to a subset of expert groups, as in DeepSeek-style routing. -* **Expert bias:** with the sigmoid score function, ``expert_bias`` shifts the - selection without changing the returned weights - the bias-adjustment scheme - used for auxiliary-loss-free load balancing. -* **Scaling:** ``scaling_factor`` rescales the returned probabilities. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: router_pytorch.py - :language: python - :start-after: # START_ROUTER_PYTORCH - :end-before: # END_ROUTER_PYTORCH - - .. tab:: JAX - - .. literalinclude:: router_jax.py - :language: python - :start-after: # START_ROUTER_JAX - :end-before: # END_ROUTER_JAX - -Load balancing -~~~~~~~~~~~~~~~ - -Left unconstrained, a router tends to collapse onto a handful of experts. The -usual remedy is an auxiliary load-balancing loss that rewards spreading tokens -evenly across experts. Transformer Engine computes it with ``fused_moe_aux_loss`` -from the per-expert token counts and the *dense* routing scores - one value per -expert rather than only the selected top-k - so the loss has a gradient with -respect to every expert's logit. Those dense scores come from -``fused_compute_score_for_moe_aux_loss`` in PyTorch, or from -``fused_topk_with_score_function(..., compute_aux_scores=True)`` in JAX. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: router_pytorch.py - :language: python - :start-after: # START_ROUTER_AUX_PYTORCH - :end-before: # END_ROUTER_AUX_PYTORCH - - .. tab:: JAX - - .. literalinclude:: router_jax.py - :language: python - :start-after: # START_ROUTER_AUX_JAX - :end-before: # END_ROUTER_AUX_JAX - -Routing Kernels ---------------- - -Once the router has produced a routing map, the tokens must be moved into the -expert-contiguous layout expected by the grouped GEMM (``GroupedLinear`` in -PyTorch, ``grouped_dense`` in JAX) and, afterwards, moved back. Transformer -Engine provides differentiable kernels for both directions. The rest of this -section focuses on the two core operations - token dispatch and token combine - -because they illustrate the layout transformation used by the other variants. - -The snippets below show one concrete instance of this pattern: the mask-map -routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` -and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as -``transformer_engine.jax.permutation.token_dispatch`` and -``transformer_engine.jax.permutation.token_combine``. Other routing variants -(for example, index-map routing in PyTorch via ``map_type="index"``) are -available in both frameworks and follow the same pattern; see the -:doc:`PyTorch API reference <../api/pytorch>` and -:doc:`JAX API reference <../api/jax>` for the complete list and signatures. -The mask-map APIs have different framework-specific wrappers, but lower to the -same shared Triton permutation kernels, and both pairs are differentiable so they -can be used directly inside training graphs. - -Token Dispatch -~~~~~~~~~~~~~~ - -Token dispatch is the canonical routing operation: given the original token -tensor and a routing map describing each token's destination expert, it returns -a permuted token buffer in which all rows assigned to the same expert are -stored contiguously. In PyTorch this operation is exposed as ``moe_permute``; -in JAX it is exposed as ``token_dispatch``. This is exactly the layout that -the grouped linear layer consumes via its per-expert token-count argument -(``m_splits`` in PyTorch ``GroupedLinear``, ``group_sizes`` in JAX -``grouped_dense``), so token dispatch followed by the grouped GEMM forms a -typical MoE forward block. - -.. figure:: img/moe_permute.svg - :align: center - :alt: Token dispatch reorders tokens so that all tokens assigned to the same expert are contiguous - - 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: - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_permute_pytorch.py - :language: python - :start-after: # START_MOE_PERMUTE_PYTORCH - :end-before: # END_MOE_PERMUTE_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_permute_jax.py - :language: python - :start-after: # START_MOE_PERMUTE_JAX - :end-before: # END_MOE_PERMUTE_JAX - -Both variants return the permuted token buffer of shape -``[num_out_tokens, hidden_size]`` together with a ``row_id_map`` that -carries enough information for token combine to restore the original token -order once the expert computation is done. Token dispatch and token combine -are typically used as a matched pair around the grouped GEMM call. - -Token Combine -~~~~~~~~~~~~~ - -Token combine is the inverse routing operation: it takes the expert-contiguous -output produced by the grouped GEMM (or any per-expert computation) and the -``row_id_map`` returned by token dispatch, and returns a single tensor of -shape ``[num_tokens, hidden_size]`` with the rows written back into the -original token order. In PyTorch this operation is exposed as -``moe_unpermute``; in JAX it is exposed as ``token_combine``. - -For top-1 routing each token has exactly one expert contribution, so -``merging_probs`` is omitted. For top-k routing pass the per-token expert -weights as ``merging_probs`` and the kernel computes a weighted sum of the -per-expert contributions in the same fused pass; without it the per-expert -contributions are summed unweighted. In PyTorch, also pass -``restore_shape=(num_tokens, hidden_size)`` whenever the permuted buffer has more -rows than the original tokens (top-k routing); JAX infers the original token -count from the ``row_id_map``. - -.. figure:: img/moe_unpermute.svg - :align: center - :alt: Token combine restores expert outputs back into the original token order - - 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: - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_unpermute_pytorch.py - :language: python - :start-after: # START_MOE_UNPERMUTE_PYTORCH - :end-before: # END_MOE_UNPERMUTE_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_unpermute_jax.py - :language: python - :start-after: # START_MOE_UNPERMUTE_JAX - :end-before: # END_MOE_UNPERMUTE_JAX - -Token probabilities -~~~~~~~~~~~~~~~~~~~~ - -In top-k routing each token contributes to several experts, and those -contributions are recombined using the routing weights. There are two equivalent -places to apply the weights: - -* **At combine (output side).** Pass the routing weights to token combine as - ``merging_probs``; it forms the weighted sum of the per-expert contributions in - the same fused pass. This is the path used in the examples above. -* **At dispatch (input side).** Scale each expert's input by its routing weight - before the grouped GEMM. ``moe_permute_with_probs`` (PyTorch) and the ``probs`` - argument of ``token_dispatch`` (JAX) permute a probability tensor alongside the - tokens, so the weights arrive already aligned with the expert-contiguous - layout. - -Padding and alignment -~~~~~~~~~~~~~~~~~~~~~~~ - -Grouped GEMM backends are most efficient when each expert's token block starts at -an aligned offset (for example, a multiple of 128 rows). Because the number of -tokens routed to an expert is data dependent, the blocks are generally ragged. -Transformer Engine can pad each block up to a multiple of ``align_size`` as part -of the dispatch kernel, avoiding a separate padding pass. - -.. figure:: img/moe_padding.svg - :align: center - :alt: Each expert block is padded up to a multiple of align_size during dispatch - - 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. - -In PyTorch this is ``moe_permute_and_pad_with_probs``; in JAX it is the -``align_size`` argument of ``token_dispatch``. Both return the padded token -buffer, the aligned per-expert token counts (used as ``m_splits`` / -``group_sizes`` for the grouped GEMM), and the per-expert ``pad_offsets`` that -token combine needs in order to remove the padding. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_permute_pad_pytorch.py - :language: python - :start-after: # START_MOE_PERMUTE_PAD_PYTORCH - :end-before: # END_MOE_PERMUTE_PAD_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_permute_pad_jax.py - :language: python - :start-after: # START_MOE_PERMUTE_PAD_JAX - :end-before: # END_MOE_PERMUTE_PAD_JAX - -Reordering expert chunks -~~~~~~~~~~~~~~~~~~~~~~~~~ - -When experts are sharded across devices, the per-expert token blocks often have -to be reordered - for example, to regroup tokens by destination rank before an -all-to-all, or to restore the original grouping afterwards. -``moe_sort_chunks_by_index`` (PyTorch) and ``sort_chunks_by_index`` (JAX) permute -contiguous chunks of a token tensor according to a list of chunk sizes and a -permutation of chunk indices, without falling back to Python-level slicing and -concatenation. ``moe_sort_chunks_by_index_with_probs`` reorders an accompanying -probability tensor in the same call. - -Grouped GEMM ------------- - -The straightforward way to apply per-expert linear layers is to loop over the -experts and call a separate ``Linear`` for each one. This is correct, but it -is not the most efficient way to execute many expert GEMMs. - -Transformer Engine provides a grouped GEMM primitive -(``GroupedLinear`` in PyTorch and ``grouped_dense`` in JAX) - an optimized -replacement that produces the same outputs as the loop while using -implementations that are better suited for MoE workloads. - -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: - -.. math:: - - 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: - -.. math:: - - Y = \mathrm{concat}(Y_0, Y_1, \ldots, Y_{G-1}) - -The grouped GEMM is told how many token rows belong to each expert via a -per-expert token-count argument: ``m_splits`` in PyTorch ``GroupedLinear`` and -``group_sizes`` in JAX ``grouped_dense``. - -.. figure:: img/grouped_linear.svg - :align: center - :alt: Comparison between launching one Linear per expert and using one grouped GEMM call for all expert blocks - - Figure 6: Both paths produce the same outputs from the same inputs. The - baseline launches one ``Linear`` per expert, while the grouped GEMM - (``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX) is an optimized - grouped implementation that replaces the loop. - -The following snippets show how to replace the loop with the grouped GEMM. -They assume the tokens have already been permuted into expert-contiguous order. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: grouped_linear_pytorch.py - :language: python - :start-after: # START_GROUPED_LINEAR_PYTORCH - :end-before: # END_GROUPED_LINEAR_PYTORCH - - .. tab:: JAX - - .. literalinclude:: grouped_linear_jax.py - :language: python - :start-after: # START_GROUPED_LINEAR_JAX - :end-before: # END_GROUPED_LINEAR_JAX - -The grouped GEMM uses implementations tuned for grouped expert execution: - -* **Optimized backends:** Transformer Engine selects from several grouped GEMM - backends depending on the framework, datatype, and GPU architecture. This - can be, for example, cuBLAS GEMMs launched on multiple CUDA streams or a - single grouped GEMM kernel, among other backend-specific implementations. -* **Recipe compatibility:** the grouped GEMM is integrated with - Transformer Engine's :doc:`low-precision training stack - `, so the same recipes available to regular - ``Linear`` layers - FP8 (delayed, current, and blockwise scaling), MXFP8, and - NVFP4 - can be used for MoE experts. -* **Fused quantization:** Low-precision grouped GEMM paths can fuse - quantization-related work such as scale computation, casting, and - cast/transpose steps across experts instead of repeating the same work in a - Python loop. -* **Fused expert MLP:** Through the :doc:`operation-based API - <../examples/op_fuser/op_fuser>`, the two expert GEMMs and the activation - between them can be fused into a single grouped operation on recent - architectures, removing the intermediate round trips to memory. - -The PyTorch ``GroupedLinear`` module also supports the features expected of a -Transformer Engine linear layer - tensor and sequence parallelism, gradient -accumulation fusion, and FP8 weight caching - so it can serve as a drop-in expert -layer. See the :doc:`PyTorch API reference <../api/pytorch>` for the full -signature. - -.. _moe-putting-it-together: - -Putting it together -------------------- - -The building blocks assemble into the four-stage MoE layer from Figure 1: route, -dispatch, run the experts, and combine. The example below wires them together for -top-k routing. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_layer_pytorch.py - :language: python - :start-after: # START_MOE_LAYER_PYTORCH - :end-before: # END_MOE_LAYER_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_layer_jax.py - :language: python - :start-after: # START_MOE_LAYER_JAX - :end-before: # END_MOE_LAYER_JAX - -This 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. The -expert step is the grouped MLP from the previous section; a full expert MLP -stacks two grouped GEMMs around an activation. Every stage is differentiable, so -the assembled layer trains end to end. - -Fused grouped MLP ------------------ - -An expert MLP is two grouped GEMMs with an activation between them: the first -projects into the (gated) feed-forward dimension, the activation is applied, and -the second projects back. Running these as separate kernels writes the large -intermediate activation out to HBM and reads it back for the second GEMM, and -re-quantizes it in a separate pass. - -On Blackwell (SM100) GPUs, Transformer Engine can fuse the whole expert MLP - -both grouped GEMMs and the activation - into a single CuTe DSL kernel. The -intermediate stays on chip and the cross-expert quantization is folded into the -GEMMs, removing the HBM round-trip and the extra kernel launches. - -.. figure:: img/moe_grouped_mlp.svg - :align: center - :alt: The two expert grouped GEMMs and the activation between them fused into one kernel - - Figure 7: The operation fuser replaces the first grouped GEMM, the activation, - and the second grouped GEMM with a single fused grouped-MLP kernel that keeps - the intermediate on chip. - -The fusion is exposed through the operation-based API and applied automatically -by the :doc:`operation fuser <../examples/op_fuser/op_fuser>`: when it sees a -grouped linear, a scaled GLU (or SReLU) activation, and another grouped linear in -sequence, it replaces them with one fused grouped-MLP operation. No change to the -forward code is needed to opt in. - -.. literalinclude:: grouped_mlp_pytorch.py - :language: python - :start-after: # START_GROUPED_MLP_PYTORCH - :end-before: # END_GROUPED_MLP_PYTORCH - -The fused path is taken when all of the following hold; otherwise the three ops -run separately and produce identical results: - -* **Architecture:** Blackwell (SM100) with cuDNN frontend 1.23 or newer. -* **Recipe:** a block-scaled low-precision recipe - MXFP8, or NVFP4 with the - randomized Hadamard transform enabled. -* **Opt-in:** the environment variable ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``. -* **Activation:** a scaled ``SwiGLU`` / ``GeGLU`` (gated) or ``SReLU`` (unary), - with feature dimensions aligned to 64 and the token count to 128. - -Expert parallelism ------------------- - -The grouped GEMM keeps all experts on a single device. When the experts no longer -fit there - or to add another dimension of parallelism - they are sharded across -devices, a scheme called expert parallelism (EP). Each device then owns only a -slice of the experts, so a token routed to a non-local expert has to travel to -the device that owns it. - -That data movement is two all-to-all collectives wrapped around the local expert -computation: a **dispatch** all-to-all sends each token to the rank that owns its -expert, the local grouped GEMM runs, and a **combine** all-to-all returns the -results to the source rank. It is the distributed counterpart of the token -dispatch and token combine kernels described above. - -.. figure:: img/moe_expert_parallel.svg - :align: center - :alt: Tokens are exchanged across ranks by all-to-all so each is processed by the rank owning its expert - - Figure 8: With experts sharded across ranks, a dispatch all-to-all routes each - token to the rank owning its expert and a combine all-to-all returns the - outputs to the source rank. - -Transformer Engine provides expert parallelism at two levels: - -* **JAX MoE layer.** ``transformer_engine.jax.moe.moe`` runs the entire layer - - router, dispatch, grouped expert GEMMs, and combine - as a single - differentiable call. Naming a mesh axis with ``ep_axis`` turns the dispatch and - combine steps into ``jax.lax.ragged_all_to_all`` collectives over that axis. - This API is currently experimental. -* **NCCL EP backend.** A common C API (``nvte_ep_dispatch`` / ``nvte_ep_combine`` - and their backward passes, declared in - ``transformer_engine/common/include/transformer_engine/ep.h``) implements the - dispatch and combine all-to-alls directly on NCCL, using NCCL symmetric-memory - windows for zero-copy transfers. It is compiled in with ``NVTE_WITH_NCCL_EP`` - (Hopper or newer, NCCL 2.30.4+) and provides the high-performance communication - path that the framework layers build on. - -A minimal call to the JAX layer, with experts sharded over the ``"ep"`` mesh -axis: - -.. literalinclude:: moe_expert_parallel_jax.py - :language: python - :start-after: # START_MOE_EXPERT_PARALLEL_JAX - :end-before: # END_MOE_EXPERT_PARALLEL_JAX diff --git a/docs/features/mixture_of_experts/expert_parallelism/expert_parallelism.rst b/docs/features/mixture_of_experts/expert_parallelism/expert_parallelism.rst new file mode 100644 index 0000000000..d14dfa9769 --- /dev/null +++ b/docs/features/mixture_of_experts/expert_parallelism/expert_parallelism.rst @@ -0,0 +1,95 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Expert parallelism +=================================== + +.. note:: + + NCCL-based expert parallelism requires Hopper (SM90) or later and NCCL 2.30.4 + or newer. It is compiled in by default when Transformer Engine is built for + these architectures; set ``NVTE_WITH_NCCL_EP=0`` at build time to disable it. + +The grouped GEMM keeps all experts on a single device. When the experts no longer +fit there - or to add another dimension of parallelism - they are sharded across +devices, a scheme called expert parallelism (EP). Each device then owns only a +slice of the experts, so a token routed to a non-local expert has to travel to +the device that owns it. + +That data movement is two all-to-all collectives wrapped around the local expert +computation: a **dispatch** all-to-all sends each token to the rank that owns its +expert, the local grouped GEMM runs, and a **combine** all-to-all returns the +results to the source rank. It is the distributed counterpart of the +:doc:`token dispatch and token combine kernels <../routing_kernels/routing_kernels>`. + +.. raw:: html + :file: img/moe_expert_parallel.svg + +*Figure 1. With experts sharded across ranks, a dispatch all-to-all routes each +token to the rank owning its expert and a combine all-to-all returns the +outputs to the source rank.* + +Transformer Engine implements dispatch and combine directly on NCCL, using +NCCL symmetric-memory windows for zero-copy transfers. The backend is a common +C API (``nvte_ep_dispatch`` / ``nvte_ep_combine`` and their backward passes, +declared in ``transformer_engine/common/include/transformer_engine/ep.h``) that +both frameworks build on: + +* **PyTorch.** ``transformer_engine.pytorch.ep`` exposes the primitives with + autograd support. ``ep_bootstrap`` initializes EP once per process on an + existing process group, an ``EpBuffer`` holds the per-call state, and + ``ep_dispatch`` / ``ep_combine`` perform the two all-to-alls. The routing + itself comes from the :doc:`router <../router/router>`; the local experts run + on the receive buffer between the two calls. +* **JAX.** ``transformer_engine.jax.moe.moe`` runs the entire layer - router, + dispatch, grouped expert GEMMs, and combine - as a single differentiable call. + ``ep_axis`` names the mesh axis the experts are sharded over, and the dispatch + and combine steps become all-to-all collectives over that axis. The underlying + primitives are also available separately in ``transformer_engine.jax.ep``. + This API is currently experimental. + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM90 (Hopper) or later +
+ + .. literalinclude:: moe_expert_parallel_pytorch.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH + :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH + + .. tab:: JAX + + .. raw:: html + +
+ Requires SM90 (Hopper) or later +
+ + .. literalinclude:: moe_expert_parallel_jax.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_JAX + :end-before: # END_MOE_EXPERT_PARALLEL_JAX + +Sizing the receive buffer +------------------------- + +Each rank receives a data-dependent number of tokens per step. Passing +``recv_capacity_per_rank`` fixes the size of the receive buffer up front, so the +step needs no device-to-host synchronization and can be captured in a CUDA graph; +the dropless worst case is ``ep_size * max_tokens_per_rank * top_k``. Omitting it +selects eager mode, which sizes the buffer from the actual receive count each +step at the cost of a host sync. + +In PyTorch, ``ep_dispatch`` can quantize the tokens on the fly when the +``EpBuffer`` is created with an MXFP8 ``dispatch_fwd_quant_recipe``, so the +all-to-all moves the low-precision payload and the local grouped GEMM consumes +it directly. Complete runnable examples live in ``examples/pytorch/ep/`` and +``examples/jax/ep/`` in the repository. diff --git a/docs/features/img/moe_expert_parallel.svg b/docs/features/mixture_of_experts/expert_parallelism/img/moe_expert_parallel.svg similarity index 98% rename from docs/features/img/moe_expert_parallel.svg rename to docs/features/mixture_of_experts/expert_parallelism/img/moe_expert_parallel.svg index 886b66d510..4336de7289 100644 --- a/docs/features/img/moe_expert_parallel.svg +++ b/docs/features/mixture_of_experts/expert_parallelism/img/moe_expert_parallel.svg @@ -1,4 +1,3 @@ -
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: grouped_mlp_pytorch.py + :language: python + :start-after: # START_GROUPED_MLP_PYTORCH + :end-before: # END_GROUPED_MLP_PYTORCH + +The fused path is taken when all of the following hold; otherwise the three ops +run separately and produce identical results: + +* **Architecture:** Blackwell (SM100) with cuDNN frontend 1.23 or newer. +* **Recipe:** a block-scaled low-precision recipe - MXFP8, or NVFP4 with the + randomized Hadamard transform enabled. +* **Opt-in:** the environment variable ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``. +* **Activation:** a scaled ``SwiGLU`` / ``GeGLU`` (gated) or ``SReLU`` (unary), + with feature dimensions aligned to 64 and the token count to 128. diff --git a/docs/features/grouped_linear_jax.py b/docs/features/mixture_of_experts/grouped_gemm/grouped_linear_jax.py similarity index 100% rename from docs/features/grouped_linear_jax.py rename to docs/features/mixture_of_experts/grouped_gemm/grouped_linear_jax.py diff --git a/docs/features/grouped_linear_pytorch.py b/docs/features/mixture_of_experts/grouped_gemm/grouped_linear_pytorch.py similarity index 100% rename from docs/features/grouped_linear_pytorch.py rename to docs/features/mixture_of_experts/grouped_gemm/grouped_linear_pytorch.py diff --git a/docs/features/grouped_mlp_pytorch.py b/docs/features/mixture_of_experts/grouped_gemm/grouped_mlp_pytorch.py similarity index 100% rename from docs/features/grouped_mlp_pytorch.py rename to docs/features/mixture_of_experts/grouped_gemm/grouped_mlp_pytorch.py diff --git a/docs/features/img/grouped_linear.svg b/docs/features/mixture_of_experts/grouped_gemm/img/grouped_linear.svg similarity index 99% rename from docs/features/img/grouped_linear.svg rename to docs/features/mixture_of_experts/grouped_gemm/img/grouped_linear.svg index 82c0be436e..e438d10ba3 100644 --- a/docs/features/img/grouped_linear.svg +++ b/docs/features/mixture_of_experts/grouped_gemm/img/grouped_linear.svg @@ -1,4 +1,3 @@ -
- Requires SM90 (Hopper) or later -
- - .. literalinclude:: moe_expert_parallel_pytorch.py - :language: python - :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH - :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH - - .. tab:: JAX - - .. raw:: html - -
- Requires SM90 (Hopper) or later -
- - .. literalinclude:: moe_expert_parallel_jax.py - :language: python - :start-after: # START_MOE_EXPERT_PARALLEL_JAX - :end-before: # END_MOE_EXPERT_PARALLEL_JAX - -Sizing the receive buffer -------------------------- - -Each rank receives a data-dependent number of tokens per step. Passing -``recv_capacity_per_rank`` fixes the size of the receive buffer up front, so the -step needs no device-to-host synchronization and can be captured in a CUDA graph; -the dropless worst case is ``ep_size * max_tokens_per_rank * top_k``. Omitting it -selects eager mode, which sizes the buffer from the actual receive count each -step at the cost of a host sync. - -In PyTorch, ``ep_dispatch`` can quantize the tokens on the fly when the -``EpBuffer`` is created with an MXFP8 ``dispatch_fwd_quant_recipe``, so the -all-to-all moves the low-precision payload and the local grouped GEMM consumes -it directly. Complete runnable examples live in ``examples/pytorch/ep/`` and -``examples/jax/ep/`` in the repository. diff --git a/docs/features/mixture_of_experts/expert_parallelism/img/moe_expert_parallel.svg b/docs/features/mixture_of_experts/expert_parallelism/img/moe_expert_parallel.svg deleted file mode 100644 index 4336de7289..0000000000 --- a/docs/features/mixture_of_experts/expert_parallelism/img/moe_expert_parallel.svg +++ /dev/null @@ -1,63 +0,0 @@ - - - - - - - Expert parallelism: all-to-all dispatch and combine - - Rank 0 - Rank 1 - - - →E0 - →E2 - - →E1 - →E3 - tokens - - - - - - all-to-all - dispatch - - - - - E0 - E1 - E2 - E3 - local experts - - - - - - all-to-all - combine - - - - - output - output - - Each token is sent to the rank that owns its expert, computed locally, and returned to its source rank. - diff --git a/docs/features/mixture_of_experts/grouped_gemm/grouped_gemm.rst b/docs/features/mixture_of_experts/grouped_gemm/grouped_gemm.rst deleted file mode 100644 index 61f07f2230..0000000000 --- a/docs/features/mixture_of_experts/grouped_gemm/grouped_gemm.rst +++ /dev/null @@ -1,140 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -Grouped GEMM -=================================== - -The straightforward way to apply per-expert linear layers is to loop over the -experts and call a separate ``Linear`` for each one. This is correct, but it -is not the most efficient way to execute many expert GEMMs. - -Transformer Engine provides a grouped GEMM primitive -(``GroupedLinear`` in PyTorch and ``grouped_dense`` in JAX) - an optimized -replacement that produces the same outputs as the loop while using -implementations that are better suited for MoE workloads. - -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: - -.. math:: - - 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: - -.. math:: - - Y = \mathrm{concat}(Y_0, Y_1, \ldots, Y_{G-1}) - -The grouped GEMM is told how many token rows belong to each expert via a -per-expert token-count argument: ``m_splits`` in PyTorch ``GroupedLinear`` and -``group_sizes`` in JAX ``grouped_dense``. - -.. raw:: html - :file: img/grouped_linear.svg - -*Figure 1. Both paths produce the same outputs from the same inputs. The -baseline launches one* ``Linear`` *per expert, while the grouped GEMM -(*\ ``GroupedLinear`` *in PyTorch,* ``grouped_dense`` *in JAX) is an optimized -grouped implementation that replaces the loop.* - -The following snippets show how to replace the loop with the grouped GEMM. -They assume the tokens have already been permuted into expert-contiguous order. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: grouped_linear_pytorch.py - :language: python - :start-after: # START_GROUPED_LINEAR_PYTORCH - :end-before: # END_GROUPED_LINEAR_PYTORCH - - .. tab:: JAX - - .. literalinclude:: grouped_linear_jax.py - :language: python - :start-after: # START_GROUPED_LINEAR_JAX - :end-before: # END_GROUPED_LINEAR_JAX - -The grouped GEMM uses implementations tuned for grouped expert execution: - -* **Optimized backends:** Transformer Engine selects from several grouped GEMM - backends depending on the framework, datatype, and GPU architecture. This - can be, for example, cuBLAS GEMMs launched on multiple CUDA streams or a - single grouped GEMM kernel, among other backend-specific implementations. -* **Recipe compatibility:** the grouped GEMM is integrated with - Transformer Engine's :doc:`low-precision training stack - `, so the same recipes available to - regular ``Linear`` layers - FP8 (delayed, current, and blockwise scaling), - MXFP8, and NVFP4 - can be used for MoE experts. -* **Fused quantization:** Low-precision grouped GEMM paths can fuse - quantization-related work such as scale computation, casting, and - cast/transpose steps across experts instead of repeating the same work in a - Python loop. -* **Fused expert MLP:** Through the :doc:`operation-based API - `, the two expert GEMMs and the activation - between them can be fused into a single grouped operation on recent - architectures; see :ref:`moe-fused-grouped-mlp` below. - -The PyTorch ``GroupedLinear`` module also supports the features expected of a -Transformer Engine linear layer - tensor and sequence parallelism, gradient -accumulation fusion, and FP8 weight caching - so it can serve as a drop-in expert -layer. See the :doc:`PyTorch API reference ` for the full -signature. - -.. _moe-fused-grouped-mlp: - -Fused grouped MLP ------------------ - -An expert MLP is two grouped GEMMs with an activation between them: the first -projects into the (gated) feed-forward dimension, the activation is applied, and -the second projects back. Running these as separate kernels writes the large -intermediate activation out to HBM and reads it back for the second GEMM, and -re-quantizes it in a separate pass. - -On Blackwell (SM100) GPUs, Transformer Engine can fuse the whole expert MLP - -both grouped GEMMs and the activation - into a single CuTe DSL kernel. The -intermediate stays on chip and the cross-expert quantization is folded into the -GEMMs, removing the HBM round-trip and the extra kernel launches. - -.. raw:: html - :file: img/moe_grouped_mlp.svg - -*Figure 2. The operation fuser replaces the first grouped GEMM, the activation, -and the second grouped GEMM with a single fused grouped-MLP kernel that keeps -the intermediate on chip.* - -The fusion is exposed through the operation-based API and applied automatically -by the :doc:`operation fuser `: when it sees a -grouped linear, a scaled GLU (or SReLU) activation, and another grouped linear in -sequence, it replaces them with one fused grouped-MLP operation. No change to the -forward code is needed to opt in. - -.. tabs:: - - .. tab:: PyTorch - - .. raw:: html - -
- Requires SM100 (Blackwell) or later -
- - .. literalinclude:: grouped_mlp_pytorch.py - :language: python - :start-after: # START_GROUPED_MLP_PYTORCH - :end-before: # END_GROUPED_MLP_PYTORCH - -The fused path is taken when all of the following hold; otherwise the three ops -run separately and produce identical results: - -* **Architecture:** Blackwell (SM100) with cuDNN frontend 1.23 or newer. -* **Recipe:** a block-scaled low-precision recipe - MXFP8, or NVFP4 with the - randomized Hadamard transform enabled. -* **Opt-in:** the environment variable ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``. -* **Activation:** a scaled ``SwiGLU`` / ``GeGLU`` (gated) or ``SReLU`` (unary), - with feature dimensions aligned to 64 and the token count to 128. diff --git a/docs/features/mixture_of_experts/grouped_gemm/grouped_linear_jax.py b/docs/features/mixture_of_experts/grouped_linear_jax.py similarity index 100% rename from docs/features/mixture_of_experts/grouped_gemm/grouped_linear_jax.py rename to docs/features/mixture_of_experts/grouped_linear_jax.py diff --git a/docs/features/mixture_of_experts/grouped_gemm/grouped_linear_pytorch.py b/docs/features/mixture_of_experts/grouped_linear_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/grouped_gemm/grouped_linear_pytorch.py rename to docs/features/mixture_of_experts/grouped_linear_pytorch.py diff --git a/docs/features/mixture_of_experts/grouped_gemm/grouped_mlp_pytorch.py b/docs/features/mixture_of_experts/grouped_mlp_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/grouped_gemm/grouped_mlp_pytorch.py rename to docs/features/mixture_of_experts/grouped_mlp_pytorch.py diff --git a/docs/features/mixture_of_experts/grouped_gemm/img/grouped_linear.svg b/docs/features/mixture_of_experts/img/grouped_linear.svg similarity index 100% rename from docs/features/mixture_of_experts/grouped_gemm/img/grouped_linear.svg rename to docs/features/mixture_of_experts/img/grouped_linear.svg diff --git a/docs/features/mixture_of_experts/img/moe_expert_parallel.svg b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg new file mode 100644 index 0000000000..8952b7ca7a --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg @@ -0,0 +1,84 @@ + + + + + + + Expert parallelism: experts E0, E1 on rank 0 and E2, E3 on rank 1 + + tokens + dispatch all-to-all + receive buffer + local experts + combine all-to-all + output + + + + Rank 0 + + Rank 1 + + + t0 → E0 + t1 → E2 + t2 → E0 + t3 → E1 + t4 → E3 + t5 → E2 + + + + + + + + + + + E0: t0 + E0: t2 + E1: t3 + E2: t1 + E2: t5 + E3: t4 + + + + + + + Grouped MLPE0, E1 + + Grouped MLPE2, E3 + + + + + + + + + + + y0 + y1 + y2 + y3 + y4 + y5 + diff --git a/docs/features/mixture_of_experts/grouped_gemm/img/moe_grouped_mlp.svg b/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg similarity index 100% rename from docs/features/mixture_of_experts/grouped_gemm/img/moe_grouped_mlp.svg rename to docs/features/mixture_of_experts/img/moe_grouped_mlp.svg diff --git a/docs/features/mixture_of_experts/introduction/img/moe_layer.svg b/docs/features/mixture_of_experts/img/moe_layer.svg similarity index 100% rename from docs/features/mixture_of_experts/introduction/img/moe_layer.svg rename to docs/features/mixture_of_experts/img/moe_layer.svg diff --git a/docs/features/mixture_of_experts/routing_kernels/img/moe_padding.svg b/docs/features/mixture_of_experts/img/moe_padding.svg similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/img/moe_padding.svg rename to docs/features/mixture_of_experts/img/moe_padding.svg diff --git a/docs/features/mixture_of_experts/routing_kernels/img/moe_permute.svg b/docs/features/mixture_of_experts/img/moe_permute.svg similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/img/moe_permute.svg rename to docs/features/mixture_of_experts/img/moe_permute.svg diff --git a/docs/features/mixture_of_experts/router/img/moe_router.svg b/docs/features/mixture_of_experts/img/moe_router.svg similarity index 100% rename from docs/features/mixture_of_experts/router/img/moe_router.svg rename to docs/features/mixture_of_experts/img/moe_router.svg diff --git a/docs/features/mixture_of_experts/routing_kernels/img/moe_unpermute.svg b/docs/features/mixture_of_experts/img/moe_unpermute.svg similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/img/moe_unpermute.svg rename to docs/features/mixture_of_experts/img/moe_unpermute.svg diff --git a/docs/features/mixture_of_experts/index.rst b/docs/features/mixture_of_experts/index.rst deleted file mode 100644 index 4ba2dbce0e..0000000000 --- a/docs/features/mixture_of_experts/index.rst +++ /dev/null @@ -1,16 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -Mixture of Experts -=================================== - -.. toctree:: - - introduction/introduction.rst - router/router.rst - routing_kernels/routing_kernels.rst - grouped_gemm/grouped_gemm.rst - moe_layer/moe_layer.rst - expert_parallelism/expert_parallelism.rst diff --git a/docs/features/mixture_of_experts/introduction/introduction.rst b/docs/features/mixture_of_experts/introduction/introduction.rst deleted file mode 100644 index da4820bc2b..0000000000 --- a/docs/features/mixture_of_experts/introduction/introduction.rst +++ /dev/null @@ -1,50 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -.. _moe-overview: - -Introduction -=================================== - -Mixture of Experts (MoE) layers replace a dense feed-forward network with a set -of expert networks and a router that sends each token to one or more experts. -This keeps the activated parameter count per token small while allowing the -model to scale to many more total parameters. - -A token passes through an MoE layer in four stages: - -#. The **router** scores the experts for each token and selects the top-k of - them. -#. **Token dispatch** gathers the tokens into expert-contiguous order. -#. The **grouped MLP** (the experts) runs a single batched computation over all - expert blocks. -#. **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. - -.. raw:: html - :file: img/moe_layer.svg - -*Figure 1. The four stages of an MoE layer. The router produces the* -``routing_map`` *consumed by token dispatch and the* ``probs`` *used as merging -weights in token combine.* - -Transformer Engine provides an optimized 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: - -* The :doc:`router <../router/router>` fuses the score function with the top-k - selection, and provides a fused load-balancing loss. -* :doc:`Token dispatch and combine <../routing_kernels/routing_kernels>` move - tokens between their original order and the expert-contiguous layout using - optimized kernels instead of Python-level gather / sort / concatenate chains. -* :doc:`Grouped GEMM <../grouped_gemm/grouped_gemm>` primitives execute the - expert linear layers efficiently once the tokens are laid out in - expert-contiguous blocks. - -:doc:`Building an MoE layer <../moe_layer/moe_layer>` wires the four stages -together into a complete layer, and :doc:`Expert parallelism -<../expert_parallelism/expert_parallelism>` covers sharding the experts across -devices. diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst new file mode 100644 index 0000000000..d0c501416f --- /dev/null +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -0,0 +1,581 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _moe-overview: + +Mixture of Experts +=================================== + +Mixture of Experts (MoE) layers replace a dense feed-forward network with a set +of expert networks and a router that sends each token to one or more experts. +This keeps the activated parameter count per token small while allowing the +model to scale to many more total parameters. + +A token passes through an MoE layer in four stages: + +#. The **router** scores the experts for each token and selects the top-k of + them. +#. **Token dispatch** gathers the tokens into expert-contiguous order. +#. The **grouped MLP** (the experts) runs a single batched computation over all + expert blocks. +#. **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. + +.. raw:: html + :file: img/moe_layer.svg + +*Figure 1. The four stages of an MoE layer. The router produces the* +``routing_map`` *consumed by token dispatch and the* ``probs`` *used as merging +weights in token combine.* + +Transformer Engine provides an optimized 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: + +* The :ref:`router ` fuses the score function with the top-k + selection, and provides a fused load-balancing loss. +* :ref:`Token dispatch and combine ` move + tokens between their original order and the expert-contiguous layout using + optimized kernels instead of Python-level gather / sort / concatenate chains. +* :ref:`Grouped GEMM ` primitives execute the + expert linear layers efficiently once the tokens are laid out in + expert-contiguous blocks. + +:ref:`Building an MoE layer ` wires the four stages +together into a complete layer, and :ref:`Expert parallelism +` covers sharding the experts across +devices. + +.. _moe-router: + +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`` - 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. + +Transformer Engine fuses the score function and the top-k selection into a single +differentiable kernel, exposed as ``fused_topk_with_score_function`` in both +``transformer_engine.pytorch.router`` and ``transformer_engine.jax.router``. All +internal math runs in FP32 for numerical stability, regardless of the logits +dtype. + +.. raw:: html + :file: img/moe_router.svg + +*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.* + +The kernel covers the score functions and selection variants used by common MoE +architectures: + +* **Score function:** ``"softmax"`` or ``"sigmoid"`` (the PyTorch API also offers + ``"sqrtsoftplus"``). With softmax, ``use_pre_softmax`` selects whether the + softmax is applied before or after the top-k. +* **Grouped (device-limited) routing:** ``num_groups`` and ``group_topk`` restrict + selection to a subset of expert groups, as in DeepSeek-style routing. +* **Expert bias:** with the sigmoid score function, ``expert_bias`` shifts the + selection without changing the returned weights - the bias-adjustment scheme + used for auxiliary-loss-free load balancing. +* **Scaling:** ``scaling_factor`` rescales the returned probabilities. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: router_pytorch.py + :language: python + :start-after: # START_ROUTER_PYTORCH + :end-before: # END_ROUTER_PYTORCH + + .. tab:: JAX + + .. literalinclude:: router_jax.py + :language: python + :start-after: # START_ROUTER_JAX + :end-before: # END_ROUTER_JAX + +Load balancing +~~~~~~~~~~~~~~ + +Left unconstrained, a router tends to collapse onto a handful of experts. The +usual remedy is an auxiliary load-balancing loss that rewards spreading tokens +evenly across experts. Transformer Engine computes it with ``fused_moe_aux_loss`` +from the per-expert token counts and the *dense* routing scores - one value per +expert rather than only the selected top-k - so the loss has a gradient with +respect to every expert's logit. Those dense scores come from +``fused_compute_score_for_moe_aux_loss`` in PyTorch, or from +``fused_topk_with_score_function(..., compute_aux_scores=True)`` in JAX. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: router_pytorch.py + :language: python + :start-after: # START_ROUTER_AUX_PYTORCH + :end-before: # END_ROUTER_AUX_PYTORCH + + .. tab:: JAX + + .. literalinclude:: router_jax.py + :language: python + :start-after: # START_ROUTER_AUX_JAX + :end-before: # END_ROUTER_AUX_JAX + +.. _moe-routing-kernels: + +Routing kernels +--------------- + +Once the router has produced a routing map, the tokens must be moved into the +expert-contiguous layout expected by the grouped GEMM (``GroupedLinear`` in +PyTorch, ``grouped_dense`` in JAX) and, afterwards, moved back. Transformer +Engine provides differentiable kernels for both directions. This section focuses on +the two core operations - token dispatch and token combine - because they +illustrate the layout transformation used by the other variants. + +The snippets below show one concrete instance of this pattern: the mask-map +routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` +and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as +``transformer_engine.jax.permutation.token_dispatch`` and +``transformer_engine.jax.permutation.token_combine``. Other routing variants +(for example, index-map routing in PyTorch via ``map_type="index"``) are +available in both frameworks and follow the same pattern; see the +:doc:`PyTorch API reference ` and +:doc:`JAX API reference ` for the complete list and signatures. +The mask-map APIs have different framework-specific wrappers, but lower to the +same shared Triton permutation kernels, and both pairs are differentiable so they +can be used directly inside training graphs. + +Token dispatch +~~~~~~~~~~~~~~ + +Token dispatch is the canonical routing operation: given the original token +tensor and a routing map describing each token's destination expert, it returns +a permuted token buffer in which all rows assigned to the same expert are +stored contiguously. In PyTorch this operation is exposed as ``moe_permute``; +in JAX it is exposed as ``token_dispatch``. This is exactly the layout that +the grouped linear layer consumes via its per-expert token-count argument +(``m_splits`` in PyTorch ``GroupedLinear``, ``group_sizes`` in JAX +``grouped_dense``), so token dispatch followed by the grouped GEMM forms a +typical MoE forward block. + +.. raw:: html + :file: img/moe_permute.svg + +*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: + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_permute_pytorch.py + :language: python + :start-after: # START_MOE_PERMUTE_PYTORCH + :end-before: # END_MOE_PERMUTE_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_permute_jax.py + :language: python + :start-after: # START_MOE_PERMUTE_JAX + :end-before: # END_MOE_PERMUTE_JAX + +Both variants return the permuted token buffer of shape +``[num_out_tokens, hidden_size]`` together with a ``row_id_map`` that +carries enough information for token combine to restore the original token +order once the expert computation is done. Token dispatch and token combine +are typically used as a matched pair around the grouped GEMM call. + +Token combine +~~~~~~~~~~~~~ + +Token combine is the inverse routing operation: it takes the expert-contiguous +output produced by the grouped GEMM (or any per-expert computation) and the +``row_id_map`` returned by token dispatch, and returns a single tensor of +shape ``[num_tokens, hidden_size]`` with the rows written back into the +original token order. In PyTorch this operation is exposed as +``moe_unpermute``; in JAX it is exposed as ``token_combine``. + +For top-1 routing each token has exactly one expert contribution, so +``merging_probs`` is omitted. For top-k routing pass the per-token expert +weights as ``merging_probs`` and the kernel computes a weighted sum of the +per-expert contributions in the same fused pass; without it the per-expert +contributions are summed unweighted. In PyTorch, also pass +``restore_shape=(num_tokens, hidden_size)`` whenever the permuted buffer has more +rows than the original tokens (top-k routing); JAX infers the original token +count from the ``row_id_map``. + +.. raw:: html + :file: img/moe_unpermute.svg + +*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: + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_unpermute_pytorch.py + :language: python + :start-after: # START_MOE_UNPERMUTE_PYTORCH + :end-before: # END_MOE_UNPERMUTE_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_unpermute_jax.py + :language: python + :start-after: # START_MOE_UNPERMUTE_JAX + :end-before: # END_MOE_UNPERMUTE_JAX + +Token probabilities +~~~~~~~~~~~~~~~~~~~ + +In top-k routing each token contributes to several experts, and those +contributions are recombined using the routing weights. There are two equivalent +places to apply the weights: + +* **At combine (output side).** Pass the routing weights to token combine as + ``merging_probs``; it forms the weighted sum of the per-expert contributions in + the same fused pass. This is the path used in the examples above. +* **At dispatch (input side).** Scale each expert's input by its routing weight + before the grouped GEMM. ``moe_permute_with_probs`` (PyTorch) and the ``probs`` + argument of ``token_dispatch`` (JAX) permute a probability tensor alongside the + tokens, so the weights arrive already aligned with the expert-contiguous + layout. + +Padding and alignment +~~~~~~~~~~~~~~~~~~~~~ + +Grouped GEMM backends are most efficient when each expert's token block starts at +an aligned offset (for example, a multiple of 128 rows). Because the number of +tokens routed to an expert is data dependent, the blocks are generally ragged. +Transformer Engine can pad each block up to a multiple of ``align_size`` as part +of the dispatch kernel, avoiding a separate padding pass. + +.. raw:: html + :file: img/moe_padding.svg + +*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.* + +In PyTorch this is ``moe_permute_and_pad_with_probs``; in JAX it is the +``align_size`` argument of ``token_dispatch``. Both return the padded token +buffer, the aligned per-expert token counts (used as ``m_splits`` / +``group_sizes`` for the grouped GEMM), and the per-expert ``pad_offsets`` that +token combine needs in order to remove the padding. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_permute_pad_pytorch.py + :language: python + :start-after: # START_MOE_PERMUTE_PAD_PYTORCH + :end-before: # END_MOE_PERMUTE_PAD_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_permute_pad_jax.py + :language: python + :start-after: # START_MOE_PERMUTE_PAD_JAX + :end-before: # END_MOE_PERMUTE_PAD_JAX + +Reordering expert chunks +~~~~~~~~~~~~~~~~~~~~~~~~ + +When experts are sharded across devices, the per-expert token blocks often have +to be reordered - for example, to regroup tokens by destination rank before an +all-to-all, or to restore the original grouping afterwards. +``moe_sort_chunks_by_index`` (PyTorch) and ``sort_chunks_by_index`` (JAX) permute +contiguous chunks of a token tensor according to a list of chunk sizes and a +permutation of chunk indices, without falling back to Python-level slicing and +concatenation. ``moe_sort_chunks_by_index_with_probs`` reorders an accompanying +probability tensor in the same call. + +.. _moe-grouped-gemm: + +Grouped GEMM +------------ + +The straightforward way to apply per-expert linear layers is to loop over the +experts and call a separate ``Linear`` for each one. This is correct, but it +is not the most efficient way to execute many expert GEMMs. + +Transformer Engine provides a grouped GEMM primitive +(``GroupedLinear`` in PyTorch and ``grouped_dense`` in JAX) - an optimized +replacement that produces the same outputs as the loop while using +implementations that are better suited for MoE workloads. + +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: + +.. math:: + + 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: + +.. math:: + + Y = \mathrm{concat}(Y_0, Y_1, \ldots, Y_{G-1}) + +The grouped GEMM is told how many token rows belong to each expert via a +per-expert token-count argument: ``m_splits`` in PyTorch ``GroupedLinear`` and +``group_sizes`` in JAX ``grouped_dense``. + +.. raw:: html + :file: img/grouped_linear.svg + +*Figure 6. Both paths produce the same outputs from the same inputs. The +baseline launches one* ``Linear`` *per expert, while the grouped GEMM +(*\ ``GroupedLinear`` *in PyTorch,* ``grouped_dense`` *in JAX) is an optimized +grouped implementation that replaces the loop.* + +The following snippets show how to replace the loop with the grouped GEMM. +They assume the tokens have already been permuted into expert-contiguous order. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: grouped_linear_pytorch.py + :language: python + :start-after: # START_GROUPED_LINEAR_PYTORCH + :end-before: # END_GROUPED_LINEAR_PYTORCH + + .. tab:: JAX + + .. literalinclude:: grouped_linear_jax.py + :language: python + :start-after: # START_GROUPED_LINEAR_JAX + :end-before: # END_GROUPED_LINEAR_JAX + +The grouped GEMM uses implementations tuned for grouped expert execution: + +* **Optimized backends:** Transformer Engine selects from several grouped GEMM + backends depending on the framework, datatype, and GPU architecture. This + can be, for example, cuBLAS GEMMs launched on multiple CUDA streams or a + single grouped GEMM kernel, among other backend-specific implementations. +* **Recipe compatibility:** the grouped GEMM is integrated with + Transformer Engine's :doc:`low-precision training stack + `, so the same recipes available to + regular ``Linear`` layers - FP8 (delayed, current, and blockwise scaling), + MXFP8, and NVFP4 - can be used for MoE experts. +* **Fused quantization:** Low-precision grouped GEMM paths can fuse + quantization-related work such as scale computation, casting, and + cast/transpose steps across experts instead of repeating the same work in a + Python loop. +* **Fused expert MLP:** Through the :doc:`operation-based API + `, the two expert GEMMs and the activation + between them can be fused into a single grouped operation on recent + architectures; see :ref:`moe-fused-grouped-mlp` below. + +The PyTorch ``GroupedLinear`` module also supports the features expected of a +Transformer Engine linear layer - tensor and sequence parallelism, gradient +accumulation fusion, and FP8 weight caching - so it can serve as a drop-in expert +layer. See the :doc:`PyTorch API reference ` for the full +signature. + +.. _moe-fused-grouped-mlp: + +Fused grouped MLP +~~~~~~~~~~~~~~~~~ + +An expert MLP is two grouped GEMMs with an activation between them: the first +projects into the (gated) feed-forward dimension, the activation is applied, and +the second projects back. Running these as separate kernels writes the large +intermediate activation out to HBM and reads it back for the second GEMM, and +re-quantizes it in a separate pass. + +On Blackwell (SM100) GPUs, Transformer Engine can fuse the whole expert MLP - +both grouped GEMMs and the activation - into a single CuTe DSL kernel. The +intermediate stays on chip and the cross-expert quantization is folded into the +GEMMs, removing the HBM round-trip and the extra kernel launches. + +.. raw:: html + :file: img/moe_grouped_mlp.svg + +*Figure 7. The operation fuser replaces the first grouped GEMM, the activation, +and the second grouped GEMM with a single fused grouped-MLP kernel that keeps +the intermediate on chip.* + +The fusion is exposed through the operation-based API and applied automatically +by the :doc:`operation fuser `: when it sees a +grouped linear, a scaled GLU (or SReLU) activation, and another grouped linear in +sequence, it replaces them with one fused grouped-MLP operation. No change to the +forward code is needed to opt in. + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: grouped_mlp_pytorch.py + :language: python + :start-after: # START_GROUPED_MLP_PYTORCH + :end-before: # END_GROUPED_MLP_PYTORCH + +The fused path is taken when all of the following hold; otherwise the three ops +run separately and produce identical results: + +* **Architecture:** Blackwell (SM100) with cuDNN frontend 1.23 or newer. +* **Recipe:** a block-scaled low-precision recipe - MXFP8, or NVFP4 with the + randomized Hadamard transform enabled. +* **Opt-in:** the environment variable ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``. +* **Activation:** a scaled ``SwiGLU`` / ``GeGLU`` (gated) or ``SReLU`` (unary), + with feature dimensions aligned to 64 and the token count to 128. + +.. _moe-putting-it-together: + +Building an MoE layer +--------------------- + +The building blocks assemble into the four-stage MoE layer from the +:ref:`introduction `: route, dispatch, run the +experts, and combine. The example below wires them together for top-k routing. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_layer_pytorch.py + :language: python + :start-after: # START_MOE_LAYER_PYTORCH + :end-before: # END_MOE_LAYER_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_layer_jax.py + :language: python + :start-after: # START_MOE_LAYER_JAX + :end-before: # END_MOE_LAYER_JAX + +This 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. The +expert step is the grouped MLP from :ref:`Grouped GEMM +`; a full expert MLP stacks two grouped GEMMs +around an activation. Every stage is differentiable, so the assembled layer +trains end to end. + +When the experts are sharded across devices, the dispatch and combine steps +become collectives; see :ref:`Expert parallelism +`. + +.. _moe-expert-parallelism: + +Expert parallelism +------------------ + +.. note:: + + NCCL-based expert parallelism requires Hopper (SM90) or later and NCCL 2.30.4 + or newer. It is compiled in by default when Transformer Engine is built for + these architectures; set ``NVTE_WITH_NCCL_EP=0`` at build time to disable it. + +The grouped GEMM keeps all experts on a single device. When the experts no longer +fit there - or to add another dimension of parallelism - they are sharded across +devices, a scheme called expert parallelism (EP). Each device then owns only a +slice of the experts, so a token routed to a non-local expert has to travel to +the device that owns it. + +That data movement is two all-to-all collectives wrapped around the local expert +computation: a **dispatch** all-to-all sends each token to the rank that owns its +expert, the local grouped GEMM runs, and a **combine** all-to-all returns the +results to the source rank. It is the distributed counterpart of the +:ref:`token dispatch and token combine kernels `. + +.. raw:: html + :file: img/moe_expert_parallel.svg + +*Figure 8. With experts sharded across ranks, a dispatch all-to-all routes each +token to the rank owning its expert and a combine all-to-all returns the +outputs to the source rank.* + +Transformer Engine implements dispatch and combine directly on NCCL, using +NCCL symmetric-memory windows for zero-copy transfers. The backend is a common +C API (``nvte_ep_dispatch`` / ``nvte_ep_combine`` and their backward passes, +declared in ``transformer_engine/common/include/transformer_engine/ep.h``) that +both frameworks build on: + +* **PyTorch.** ``transformer_engine.pytorch.ep`` exposes the primitives with + autograd support. ``ep_bootstrap`` initializes EP once per process on an + existing process group, an ``EpBuffer`` holds the per-call state, and + ``ep_dispatch`` / ``ep_combine`` perform the two all-to-alls. The routing + itself comes from the :ref:`router `; the local experts run + on the receive buffer between the two calls. +* **JAX.** ``transformer_engine.jax.moe.moe`` runs the entire layer - router, + dispatch, grouped expert GEMMs, and combine - as a single differentiable call. + ``ep_axis`` names the mesh axis the experts are sharded over, and the dispatch + and combine steps become all-to-all collectives over that axis. The underlying + primitives are also available separately in ``transformer_engine.jax.ep``. + This API is currently experimental. + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM90 (Hopper) or later +
+ + .. literalinclude:: moe_expert_parallel_pytorch.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH + :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH + + .. tab:: JAX + + .. raw:: html + +
+ Requires SM90 (Hopper) or later +
+ + .. literalinclude:: moe_expert_parallel_jax.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_JAX + :end-before: # END_MOE_EXPERT_PARALLEL_JAX + +Sizing the receive buffer +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Each rank receives a data-dependent number of tokens per step. Passing +``recv_capacity_per_rank`` fixes the size of the receive buffer up front, so the +step needs no device-to-host synchronization and can be captured in a CUDA graph; +the dropless worst case is ``ep_size * max_tokens_per_rank * top_k``. Omitting it +selects eager mode, which sizes the buffer from the actual receive count each +step at the cost of a host sync. + +In PyTorch, ``ep_dispatch`` can quantize the tokens on the fly when the +``EpBuffer`` is created with an MXFP8 ``dispatch_fwd_quant_recipe``, so the +all-to-all moves the low-precision payload and the local grouped GEMM consumes +it directly. Complete runnable examples live in ``examples/pytorch/ep/`` and +``examples/jax/ep/`` in the repository. diff --git a/docs/features/mixture_of_experts/expert_parallelism/moe_expert_parallel_jax.py b/docs/features/mixture_of_experts/moe_expert_parallel_jax.py similarity index 100% rename from docs/features/mixture_of_experts/expert_parallelism/moe_expert_parallel_jax.py rename to docs/features/mixture_of_experts/moe_expert_parallel_jax.py diff --git a/docs/features/mixture_of_experts/expert_parallelism/moe_expert_parallel_pytorch.py b/docs/features/mixture_of_experts/moe_expert_parallel_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/expert_parallelism/moe_expert_parallel_pytorch.py rename to docs/features/mixture_of_experts/moe_expert_parallel_pytorch.py diff --git a/docs/features/mixture_of_experts/moe_layer/moe_layer.rst b/docs/features/mixture_of_experts/moe_layer/moe_layer.rst deleted file mode 100644 index b2bff2a863..0000000000 --- a/docs/features/mixture_of_experts/moe_layer/moe_layer.rst +++ /dev/null @@ -1,40 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -.. _moe-putting-it-together: - -Building an MoE layer -=================================== - -The building blocks assemble into the four-stage MoE layer from the -:doc:`introduction <../introduction/introduction>`: route, dispatch, run the -experts, and combine. The example below wires them together for top-k routing. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_layer_pytorch.py - :language: python - :start-after: # START_MOE_LAYER_PYTORCH - :end-before: # END_MOE_LAYER_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_layer_jax.py - :language: python - :start-after: # START_MOE_LAYER_JAX - :end-before: # END_MOE_LAYER_JAX - -This 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. The -expert step is the grouped MLP from :doc:`Grouped GEMM -<../grouped_gemm/grouped_gemm>`; a full expert MLP stacks two grouped GEMMs -around an activation. Every stage is differentiable, so the assembled layer -trains end to end. - -When the experts are sharded across devices, the dispatch and combine steps -become collectives; see :doc:`Expert parallelism -<../expert_parallelism/expert_parallelism>`. diff --git a/docs/features/mixture_of_experts/moe_layer/moe_layer_jax.py b/docs/features/mixture_of_experts/moe_layer_jax.py similarity index 100% rename from docs/features/mixture_of_experts/moe_layer/moe_layer_jax.py rename to docs/features/mixture_of_experts/moe_layer_jax.py diff --git a/docs/features/mixture_of_experts/moe_layer/moe_layer_pytorch.py b/docs/features/mixture_of_experts/moe_layer_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/moe_layer/moe_layer_pytorch.py rename to docs/features/mixture_of_experts/moe_layer_pytorch.py diff --git a/docs/features/mixture_of_experts/routing_kernels/moe_permute_jax.py b/docs/features/mixture_of_experts/moe_permute_jax.py similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/moe_permute_jax.py rename to docs/features/mixture_of_experts/moe_permute_jax.py diff --git a/docs/features/mixture_of_experts/routing_kernels/moe_permute_pad_jax.py b/docs/features/mixture_of_experts/moe_permute_pad_jax.py similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/moe_permute_pad_jax.py rename to docs/features/mixture_of_experts/moe_permute_pad_jax.py diff --git a/docs/features/mixture_of_experts/routing_kernels/moe_permute_pad_pytorch.py b/docs/features/mixture_of_experts/moe_permute_pad_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/moe_permute_pad_pytorch.py rename to docs/features/mixture_of_experts/moe_permute_pad_pytorch.py diff --git a/docs/features/mixture_of_experts/routing_kernels/moe_permute_pytorch.py b/docs/features/mixture_of_experts/moe_permute_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/moe_permute_pytorch.py rename to docs/features/mixture_of_experts/moe_permute_pytorch.py diff --git a/docs/features/mixture_of_experts/routing_kernels/moe_unpermute_jax.py b/docs/features/mixture_of_experts/moe_unpermute_jax.py similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/moe_unpermute_jax.py rename to docs/features/mixture_of_experts/moe_unpermute_jax.py diff --git a/docs/features/mixture_of_experts/routing_kernels/moe_unpermute_pytorch.py b/docs/features/mixture_of_experts/moe_unpermute_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/routing_kernels/moe_unpermute_pytorch.py rename to docs/features/mixture_of_experts/moe_unpermute_pytorch.py diff --git a/docs/features/mixture_of_experts/router/router.rst b/docs/features/mixture_of_experts/router/router.rst deleted file mode 100644 index 82778a3598..0000000000 --- a/docs/features/mixture_of_experts/router/router.rst +++ /dev/null @@ -1,86 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -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`` - 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. - -Transformer Engine fuses the score function and the top-k selection into a single -differentiable kernel, exposed as ``fused_topk_with_score_function`` in both -``transformer_engine.pytorch.router`` and ``transformer_engine.jax.router``. All -internal math runs in FP32 for numerical stability, regardless of the logits -dtype. - -.. raw:: html - :file: img/moe_router.svg - -*Figure 1. 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.* - -The kernel covers the score functions and selection variants used by common MoE -architectures: - -* **Score function:** ``"softmax"`` or ``"sigmoid"`` (the PyTorch API also offers - ``"sqrtsoftplus"``). With softmax, ``use_pre_softmax`` selects whether the - softmax is applied before or after the top-k. -* **Grouped (device-limited) routing:** ``num_groups`` and ``group_topk`` restrict - selection to a subset of expert groups, as in DeepSeek-style routing. -* **Expert bias:** with the sigmoid score function, ``expert_bias`` shifts the - selection without changing the returned weights - the bias-adjustment scheme - used for auxiliary-loss-free load balancing. -* **Scaling:** ``scaling_factor`` rescales the returned probabilities. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: router_pytorch.py - :language: python - :start-after: # START_ROUTER_PYTORCH - :end-before: # END_ROUTER_PYTORCH - - .. tab:: JAX - - .. literalinclude:: router_jax.py - :language: python - :start-after: # START_ROUTER_JAX - :end-before: # END_ROUTER_JAX - -Load balancing --------------- - -Left unconstrained, a router tends to collapse onto a handful of experts. The -usual remedy is an auxiliary load-balancing loss that rewards spreading tokens -evenly across experts. Transformer Engine computes it with ``fused_moe_aux_loss`` -from the per-expert token counts and the *dense* routing scores - one value per -expert rather than only the selected top-k - so the loss has a gradient with -respect to every expert's logit. Those dense scores come from -``fused_compute_score_for_moe_aux_loss`` in PyTorch, or from -``fused_topk_with_score_function(..., compute_aux_scores=True)`` in JAX. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: router_pytorch.py - :language: python - :start-after: # START_ROUTER_AUX_PYTORCH - :end-before: # END_ROUTER_AUX_PYTORCH - - .. tab:: JAX - - .. literalinclude:: router_jax.py - :language: python - :start-after: # START_ROUTER_AUX_JAX - :end-before: # END_ROUTER_AUX_JAX diff --git a/docs/features/mixture_of_experts/router/router_jax.py b/docs/features/mixture_of_experts/router_jax.py similarity index 100% rename from docs/features/mixture_of_experts/router/router_jax.py rename to docs/features/mixture_of_experts/router_jax.py diff --git a/docs/features/mixture_of_experts/router/router_pytorch.py b/docs/features/mixture_of_experts/router_pytorch.py similarity index 100% rename from docs/features/mixture_of_experts/router/router_pytorch.py rename to docs/features/mixture_of_experts/router_pytorch.py diff --git a/docs/features/mixture_of_experts/routing_kernels/routing_kernels.rst b/docs/features/mixture_of_experts/routing_kernels/routing_kernels.rst deleted file mode 100644 index 7c22d781f6..0000000000 --- a/docs/features/mixture_of_experts/routing_kernels/routing_kernels.rst +++ /dev/null @@ -1,182 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -Routing kernels -=================================== - -Once the router has produced a routing map, the tokens must be moved into the -expert-contiguous layout expected by the grouped GEMM (``GroupedLinear`` in -PyTorch, ``grouped_dense`` in JAX) and, afterwards, moved back. Transformer -Engine provides differentiable kernels for both directions. This page focuses on -the two core operations - token dispatch and token combine - because they -illustrate the layout transformation used by the other variants. - -The snippets below show one concrete instance of this pattern: the mask-map -routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` -and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as -``transformer_engine.jax.permutation.token_dispatch`` and -``transformer_engine.jax.permutation.token_combine``. Other routing variants -(for example, index-map routing in PyTorch via ``map_type="index"``) are -available in both frameworks and follow the same pattern; see the -:doc:`PyTorch API reference ` and -:doc:`JAX API reference ` for the complete list and signatures. -The mask-map APIs have different framework-specific wrappers, but lower to the -same shared Triton permutation kernels, and both pairs are differentiable so they -can be used directly inside training graphs. - -Token dispatch --------------- - -Token dispatch is the canonical routing operation: given the original token -tensor and a routing map describing each token's destination expert, it returns -a permuted token buffer in which all rows assigned to the same expert are -stored contiguously. In PyTorch this operation is exposed as ``moe_permute``; -in JAX it is exposed as ``token_dispatch``. This is exactly the layout that -the grouped linear layer consumes via its per-expert token-count argument -(``m_splits`` in PyTorch ``GroupedLinear``, ``group_sizes`` in JAX -``grouped_dense``), so token dispatch followed by the grouped GEMM forms a -typical MoE forward block. - -.. raw:: html - :file: img/moe_permute.svg - -*Figure 1. 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: - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_permute_pytorch.py - :language: python - :start-after: # START_MOE_PERMUTE_PYTORCH - :end-before: # END_MOE_PERMUTE_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_permute_jax.py - :language: python - :start-after: # START_MOE_PERMUTE_JAX - :end-before: # END_MOE_PERMUTE_JAX - -Both variants return the permuted token buffer of shape -``[num_out_tokens, hidden_size]`` together with a ``row_id_map`` that -carries enough information for token combine to restore the original token -order once the expert computation is done. Token dispatch and token combine -are typically used as a matched pair around the grouped GEMM call. - -Token combine -------------- - -Token combine is the inverse routing operation: it takes the expert-contiguous -output produced by the grouped GEMM (or any per-expert computation) and the -``row_id_map`` returned by token dispatch, and returns a single tensor of -shape ``[num_tokens, hidden_size]`` with the rows written back into the -original token order. In PyTorch this operation is exposed as -``moe_unpermute``; in JAX it is exposed as ``token_combine``. - -For top-1 routing each token has exactly one expert contribution, so -``merging_probs`` is omitted. For top-k routing pass the per-token expert -weights as ``merging_probs`` and the kernel computes a weighted sum of the -per-expert contributions in the same fused pass; without it the per-expert -contributions are summed unweighted. In PyTorch, also pass -``restore_shape=(num_tokens, hidden_size)`` whenever the permuted buffer has more -rows than the original tokens (top-k routing); JAX infers the original token -count from the ``row_id_map``. - -.. raw:: html - :file: img/moe_unpermute.svg - -*Figure 2. 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: - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_unpermute_pytorch.py - :language: python - :start-after: # START_MOE_UNPERMUTE_PYTORCH - :end-before: # END_MOE_UNPERMUTE_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_unpermute_jax.py - :language: python - :start-after: # START_MOE_UNPERMUTE_JAX - :end-before: # END_MOE_UNPERMUTE_JAX - -Token probabilities -------------------- - -In top-k routing each token contributes to several experts, and those -contributions are recombined using the routing weights. There are two equivalent -places to apply the weights: - -* **At combine (output side).** Pass the routing weights to token combine as - ``merging_probs``; it forms the weighted sum of the per-expert contributions in - the same fused pass. This is the path used in the examples above. -* **At dispatch (input side).** Scale each expert's input by its routing weight - before the grouped GEMM. ``moe_permute_with_probs`` (PyTorch) and the ``probs`` - argument of ``token_dispatch`` (JAX) permute a probability tensor alongside the - tokens, so the weights arrive already aligned with the expert-contiguous - layout. - -Padding and alignment ---------------------- - -Grouped GEMM backends are most efficient when each expert's token block starts at -an aligned offset (for example, a multiple of 128 rows). Because the number of -tokens routed to an expert is data dependent, the blocks are generally ragged. -Transformer Engine can pad each block up to a multiple of ``align_size`` as part -of the dispatch kernel, avoiding a separate padding pass. - -.. raw:: html - :file: img/moe_padding.svg - -*Figure 3. 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.* - -In PyTorch this is ``moe_permute_and_pad_with_probs``; in JAX it is the -``align_size`` argument of ``token_dispatch``. Both return the padded token -buffer, the aligned per-expert token counts (used as ``m_splits`` / -``group_sizes`` for the grouped GEMM), and the per-expert ``pad_offsets`` that -token combine needs in order to remove the padding. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_permute_pad_pytorch.py - :language: python - :start-after: # START_MOE_PERMUTE_PAD_PYTORCH - :end-before: # END_MOE_PERMUTE_PAD_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_permute_pad_jax.py - :language: python - :start-after: # START_MOE_PERMUTE_PAD_JAX - :end-before: # END_MOE_PERMUTE_PAD_JAX - -Reordering expert chunks ------------------------- - -When experts are sharded across devices, the per-expert token blocks often have -to be reordered - for example, to regroup tokens by destination rank before an -all-to-all, or to restore the original grouping afterwards. -``moe_sort_chunks_by_index`` (PyTorch) and ``sort_chunks_by_index`` (JAX) permute -contiguous chunks of a token tensor according to a list of chunk sizes and a -permutation of chunk indices, without falling back to Python-level slicing and -concatenation. ``moe_sort_chunks_by_index_with_probs`` reorders an accompanying -probability tensor in the same call. diff --git a/docs/index.rst b/docs/index.rst index de3669e3d6..7b63281efd 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -46,7 +46,7 @@ Transformer Engine documentation :caption: Features features/low_precision_training/index.rst - features/mixture_of_experts/index.rst + features/mixture_of_experts/mixture_of_experts.rst features/other_optimizations/index.rst From 2758490e3c71d33987c4011013059ce168c7d10c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:28:21 +0200 Subject: [PATCH 07/52] [Docs] MoE: warn about uneven low-precision support Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index d0c501416f..cb52802714 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -8,6 +8,15 @@ Mixture of Experts =================================== +.. warning:: + + Low-precision support across the MoE building blocks is uneven. Which + recipes (FP8 delayed, current, and blockwise scaling, MXFP8, NVFP4) a given + block accepts depends on the block, the framework, the GPU architecture, and + the cuBLAS/cuDNN version; some paths are BF16/FP16 only. Coverage is being + extended with each release, so check the API reference of the specific + function or module for the current constraints. + Mixture of Experts (MoE) layers replace a dense feed-forward network with a set of expert networks and a router that sends each token to one or more experts. This keeps the activated parameter count per token small while allowing the @@ -379,9 +388,11 @@ The grouped GEMM uses implementations tuned for grouped expert execution: single grouped GEMM kernel, among other backend-specific implementations. * **Recipe compatibility:** the grouped GEMM is integrated with Transformer Engine's :doc:`low-precision training stack - `, so the same recipes available to - regular ``Linear`` layers - FP8 (delayed, current, and blockwise scaling), - MXFP8, and NVFP4 - can be used for MoE experts. + `, so the recipes available to + regular ``Linear`` layers can also be used for MoE experts. The exact set + depends on the execution path, GPU architecture, and cuBLAS version; see the + ``GroupedLinear`` / ``grouped_dense`` API reference for the current + constraints. * **Fused quantization:** Low-precision grouped GEMM paths can fuse quantization-related work such as scale computation, casting, and cast/transpose steps across experts instead of repeating the same work in a From c850a45ce5e5ecb1b2322da4b1f5d3419d35e6bc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:29:43 +0200 Subject: [PATCH 08/52] [Docs] MoE: shorten low-precision warning Signed-off-by: Pawel Gadzinski --- docs/features/mixture_of_experts/mixture_of_experts.rst | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index cb52802714..3b7f9a1286 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -10,12 +10,8 @@ Mixture of Experts .. warning:: - Low-precision support across the MoE building blocks is uneven. Which - recipes (FP8 delayed, current, and blockwise scaling, MXFP8, NVFP4) a given - block accepts depends on the block, the framework, the GPU architecture, and - the cuBLAS/cuDNN version; some paths are BF16/FP16 only. Coverage is being - extended with each release, so check the API reference of the specific - function or module for the current constraints. + Not every MoE building block supports every low-precision recipe yet. + Support is being extended; see the API reference of each function for details. Mixture of Experts (MoE) layers replace a dense feed-forward network with a set of expert networks and a router that sends each token to one or more experts. From ae16f61a4fb2dd03d636712cfd20bc3950067494 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:31:47 +0200 Subject: [PATCH 09/52] [Docs] MoE: reword low-precision note Signed-off-by: Pawel Gadzinski --- docs/features/mixture_of_experts/mixture_of_experts.rst | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 3b7f9a1286..3f024c1229 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -8,10 +8,12 @@ Mixture of Experts =================================== -.. warning:: +.. note:: - Not every MoE building block supports every low-precision recipe yet. - Support is being extended; see the API reference of each function for details. + The MoE building blocks are designed to work with Transformer Engine's + low-precision recipes. This support is still being extended, so not every + block works with every recipe yet; see the API reference of each function + for details. Mixture of Experts (MoE) layers replace a dense feed-forward network with a set of expert networks and a router that sends each token to one or more experts. From b8f5625def06294d9e8aef599858ed1c80288c81 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:34:18 +0200 Subject: [PATCH 10/52] [Docs] MoE: drop API-reference pointer from low-precision note Signed-off-by: Pawel Gadzinski --- docs/features/mixture_of_experts/mixture_of_experts.rst | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 3f024c1229..e1658c5361 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -12,8 +12,7 @@ Mixture of Experts The MoE building blocks are designed to work with Transformer Engine's low-precision recipes. This support is still being extended, so not every - block works with every recipe yet; see the API reference of each function - for details. + block works with every recipe yet. Mixture of Experts (MoE) layers replace a dense feed-forward network with a set of expert networks and a router that sends each token to one or more experts. From 473273f21fecfe1a9975640ff3ef4de02ebf36d1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:39:21 +0200 Subject: [PATCH 11/52] [Docs] MoE: reorganize into introduction, routing kernels, grouped GEMM, grouped MLP, EP Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 115 +++++++++--------- 1 file changed, 57 insertions(+), 58 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index e1658c5361..65a841df27 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -14,6 +14,9 @@ Mixture of Experts low-precision recipes. This support is still being extended, so not every block works with every recipe yet. +Introduction +------------ + Mixture of Experts (MoE) layers replace a dense feed-forward network with a set of expert networks and a router that sends each token to one or more experts. This keeps the activated parameter count per token small while allowing the @@ -41,24 +44,61 @@ Transformer Engine provides an optimized 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: -* The :ref:`router ` fuses the score function with the top-k - selection, and provides a fused load-balancing loss. -* :ref:`Token dispatch and combine ` move - tokens between their original order and the expert-contiguous layout using - optimized kernels instead of Python-level gather / sort / concatenate chains. -* :ref:`Grouped GEMM ` primitives execute the - expert linear layers efficiently once the tokens are laid out in - expert-contiguous blocks. +* :ref:`Routing kernels `: the router fuses the score + function with the top-k selection, and token dispatch and combine move tokens + between their original order and the expert-contiguous layout with optimized + kernels instead of Python-level gather / sort / concatenate chains. +* :ref:`Grouped GEMM ` primitives execute the expert linear + layers efficiently once the tokens are laid out in expert-contiguous blocks, + and the :ref:`grouped MLP ` fuses the whole expert MLP into + one kernel. +* :ref:`Expert parallelism ` shards the experts across + devices with all-to-all dispatch and combine. + +.. _moe-putting-it-together: + +Putting it together +~~~~~~~~~~~~~~~~~~~ + +The building blocks assemble into the four stages above: route, dispatch, run +the experts, and combine. The example below wires them together for top-k +routing; the sections that follow describe each block on its own. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_layer_pytorch.py + :language: python + :start-after: # START_MOE_LAYER_PYTORCH + :end-before: # END_MOE_LAYER_PYTORCH + + .. tab:: JAX -:ref:`Building an MoE layer ` wires the four stages -together into a complete layer, and :ref:`Expert parallelism -` covers sharding the experts across -devices. + .. literalinclude:: moe_layer_jax.py + :language: python + :start-after: # START_MOE_LAYER_JAX + :end-before: # END_MOE_LAYER_JAX + +This 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. The +expert step is built from the :ref:`grouped GEMM `; a full +expert MLP stacks two grouped GEMMs around an activation. Every stage is differentiable, so the assembled layer +trains end to end. + +When the experts are sharded across devices, the dispatch and combine steps +become collectives; see :ref:`Expert parallelism +`. + +.. _moe-routing-kernels: + +Routing kernels +--------------- .. _moe-router: 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 @@ -139,11 +179,6 @@ respect to every expert's logit. Those dense scores come from :start-after: # START_ROUTER_AUX_JAX :end-before: # END_ROUTER_AUX_JAX -.. _moe-routing-kernels: - -Routing kernels ---------------- - Once the router has produced a routing map, the tokens must be moved into the expert-contiguous layout expected by the grouped GEMM (``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX) and, afterwards, moved back. Transformer @@ -397,7 +432,7 @@ The grouped GEMM uses implementations tuned for grouped expert execution: * **Fused expert MLP:** Through the :doc:`operation-based API `, the two expert GEMMs and the activation between them can be fused into a single grouped operation on recent - architectures; see :ref:`moe-fused-grouped-mlp` below. + architectures; see :ref:`Grouped MLP `. The PyTorch ``GroupedLinear`` module also supports the features expected of a Transformer Engine linear layer - tensor and sequence parallelism, gradient @@ -405,10 +440,10 @@ accumulation fusion, and FP8 weight caching - so it can serve as a drop-in exper layer. See the :doc:`PyTorch API reference ` for the full signature. -.. _moe-fused-grouped-mlp: +.. _moe-grouped-mlp: -Fused grouped MLP -~~~~~~~~~~~~~~~~~ +Grouped MLP +----------- An expert MLP is two grouped GEMMs with an activation between them: the first projects into the (gated) feed-forward dimension, the activation is applied, and @@ -459,42 +494,6 @@ run separately and produce identical results: * **Activation:** a scaled ``SwiGLU`` / ``GeGLU`` (gated) or ``SReLU`` (unary), with feature dimensions aligned to 64 and the token count to 128. -.. _moe-putting-it-together: - -Building an MoE layer ---------------------- - -The building blocks assemble into the four-stage MoE layer from the -:ref:`introduction `: route, dispatch, run the -experts, and combine. The example below wires them together for top-k routing. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_layer_pytorch.py - :language: python - :start-after: # START_MOE_LAYER_PYTORCH - :end-before: # END_MOE_LAYER_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_layer_jax.py - :language: python - :start-after: # START_MOE_LAYER_JAX - :end-before: # END_MOE_LAYER_JAX - -This 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. The -expert step is the grouped MLP from :ref:`Grouped GEMM -`; a full expert MLP stacks two grouped GEMMs -around an activation. Every stage is differentiable, so the assembled layer -trains end to end. - -When the experts are sharded across devices, the dispatch and combine steps -become collectives; see :ref:`Expert parallelism -`. - .. _moe-expert-parallelism: Expert parallelism From ff4da3483aeb214ce46cdf5e0f14e617e351ceba Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:50:07 +0200 Subject: [PATCH 12/52] [Docs] MoE: move end-to-end example to the end of the article Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 71 ++++++++++--------- 1 file changed, 37 insertions(+), 34 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 65a841df27..06bceb9552 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -55,40 +55,8 @@ layer or dropped into an existing implementation one piece at a time: * :ref:`Expert parallelism ` shards the experts across devices with all-to-all dispatch and combine. -.. _moe-putting-it-together: - -Putting it together -~~~~~~~~~~~~~~~~~~~ - -The building blocks assemble into the four stages above: route, dispatch, run -the experts, and combine. The example below wires them together for top-k -routing; the sections that follow describe each block on its own. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: moe_layer_pytorch.py - :language: python - :start-after: # START_MOE_LAYER_PYTORCH - :end-before: # END_MOE_LAYER_PYTORCH - - .. tab:: JAX - - .. literalinclude:: moe_layer_jax.py - :language: python - :start-after: # START_MOE_LAYER_JAX - :end-before: # END_MOE_LAYER_JAX - -This 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. The -expert step is built from the :ref:`grouped GEMM `; a full -expert MLP stacks two grouped GEMMs around an activation. Every stage is differentiable, so the assembled layer -trains end to end. - -When the experts are sharded across devices, the dispatch and combine steps -become collectives; see :ref:`Expert parallelism -`. +The :ref:`example at the end ` wires the blocks into a +complete MoE layer. .. _moe-routing-kernels: @@ -586,3 +554,38 @@ In PyTorch, ``ep_dispatch`` can quantize the tokens on the fly when the all-to-all moves the low-precision payload and the local grouped GEMM consumes it directly. Complete runnable examples live in ``examples/pytorch/ep/`` and ``examples/jax/ep/`` in the repository. + +.. _moe-putting-it-together: + +Example: putting it all together +-------------------------------- + +The building blocks assemble into the four stages from the introduction: route, +dispatch, run the experts, and combine. The example below wires them together +for top-k routing on a single device. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: moe_layer_pytorch.py + :language: python + :start-after: # START_MOE_LAYER_PYTORCH + :end-before: # END_MOE_LAYER_PYTORCH + + .. tab:: JAX + + .. literalinclude:: moe_layer_jax.py + :language: python + :start-after: # START_MOE_LAYER_JAX + :end-before: # END_MOE_LAYER_JAX + +This 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. The +expert step is built from the :ref:`grouped GEMM `; a full +expert MLP stacks two grouped GEMMs around an activation. Every stage is differentiable, so the assembled layer +trains end to end. + +When the experts are sharded across devices, the dispatch and combine steps +become the all-to-all collectives described in :ref:`Expert parallelism +`. From 90624cbdabb9210915503b40a13c63d431fdc405 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:55:04 +0200 Subject: [PATCH 13/52] [Docs] MoE: add expert-parallel layer diagram to the end-to-end example Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/img/moe_layer_ep.svg | 77 +++++++++++++++++++ .../mixture_of_experts/mixture_of_experts.rst | 24 +++++- 2 files changed, 98 insertions(+), 3 deletions(-) create mode 100644 docs/features/mixture_of_experts/img/moe_layer_ep.svg diff --git a/docs/features/mixture_of_experts/img/moe_layer_ep.svg b/docs/features/mixture_of_experts/img/moe_layer_ep.svg new file mode 100644 index 0000000000..e961772b86 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_layer_ep.svg @@ -0,0 +1,77 @@ + + + + + + + + + + + + Mixture of Experts layer with expert parallelism + + + + Router + + + Token + Dispatch + + + All-to-all + dispatch + + + Grouped MLP + (local experts) + + + All-to-all + combine + + + Token + Combine + + + + tokens + + + routing_map + + + by rank + + + by expert + + + expert out + + + returned + + + output + + + + probs (routing weights) + + + + expert parallelism: tokens cross ranks, experts stay local + diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 06bceb9552..1844c8bb53 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -586,6 +586,24 @@ expert step is built from the :ref:`grouped GEMM `; a full expert MLP stacks two grouped GEMMs around an activation. Every stage is differentiable, so the assembled layer trains end to end. -When the experts are sharded across devices, the dispatch and combine steps -become the all-to-all collectives described in :ref:`Expert parallelism -`. +When the experts are sharded across devices, the same layer gains two all-to-all +collectives around the local experts. + +.. raw:: html + :file: img/moe_layer_ep.svg + +*Figure 9. The MoE layer with expert parallelism. Token dispatch groups the tokens +by destination rank, the all-to-all dispatch moves them to the ranks owning their +experts, the local grouped MLP runs, and the all-to-all combine returns the +outputs before token combine restores the original order and applies the routing +weights.* + +The routing kernels and expert parallelism complement each other. With a generic +all-to-all, the routing kernels do the reordering on both sides of the +communication: tokens are sorted by destination rank before the all-to-all and +regrouped by local expert after it (see :ref:`Reordering expert chunks +`). With the NCCL-based :ref:`expert parallelism +` primitives, this permutation is folded into the +communication itself: the dispatch delivers an expert-contiguous receive buffer +and the combine writes the results straight back into the original token order, +so no separate permute or unpermute is needed on the local side. From c1460e686240fd8096054d4cbfcda168abaa97f6 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 13:59:41 +0200 Subject: [PATCH 14/52] [Docs] MoE: use expert-parallel layer diagram as Figure 1, mention EP in introduction Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/img/moe_layer.svg | 57 ------------------- .../mixture_of_experts/mixture_of_experts.rst | 31 +++++----- 2 files changed, 16 insertions(+), 72 deletions(-) delete mode 100644 docs/features/mixture_of_experts/img/moe_layer.svg diff --git a/docs/features/mixture_of_experts/img/moe_layer.svg b/docs/features/mixture_of_experts/img/moe_layer.svg deleted file mode 100644 index 4c65f08e57..0000000000 --- a/docs/features/mixture_of_experts/img/moe_layer.svg +++ /dev/null @@ -1,57 +0,0 @@ - - - - - - - - - - - - Mixture of Experts layer - - - - Router - - - Token - Dispatch - - - Grouped MLP - (experts) - - - Token - Combine - - - - tokens - - - routing_map - - - permuted - - - expert out - - - output - - - - probs (routing weights) - diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 1844c8bb53..bbdcbb008b 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -33,12 +33,18 @@ A token passes through an MoE layer in four stages: order, merging the contributions when a token was sent to more than one expert. +When the experts do not fit on one device, they are sharded across devices +(**expert parallelism**). The layer then gains two all-to-all collectives: a +dispatch that sends each token to the rank owning its expert, and a combine that +returns the results to the source rank. The experts themselves stay local. + .. raw:: html - :file: img/moe_layer.svg + :file: img/moe_layer_ep.svg -*Figure 1. The four stages of an MoE layer. The router produces the* -``routing_map`` *consumed by token dispatch and the* ``probs`` *used as merging -weights in token combine.* +*Figure 1. The stages of an MoE layer with expert parallelism. The router produces +the* ``routing_map`` *consumed by token dispatch and the* ``probs`` *used as merging +weights in token combine; the all-to-all dispatch and combine are only present +when the experts are sharded across ranks.* Transformer Engine provides an optimized building block for each stage. They are exposed as standalone functions, so they can be assembled into a complete MoE @@ -586,17 +592,12 @@ expert step is built from the :ref:`grouped GEMM `; a full expert MLP stacks two grouped GEMMs around an activation. Every stage is differentiable, so the assembled layer trains end to end. -When the experts are sharded across devices, the same layer gains two all-to-all -collectives around the local experts. - -.. raw:: html - :file: img/moe_layer_ep.svg - -*Figure 9. The MoE layer with expert parallelism. Token dispatch groups the tokens -by destination rank, the all-to-all dispatch moves them to the ranks owning their -experts, the local grouped MLP runs, and the all-to-all combine returns the -outputs before token combine restores the original order and applies the routing -weights.* +When the experts are sharded across devices, the same layer gains the two +all-to-all collectives from Figure 1 around the local experts: token dispatch +groups the tokens by destination rank, the all-to-all dispatch moves them to the +ranks owning their experts, the local grouped MLP runs, and the all-to-all +combine returns the outputs before token combine restores the original order and +applies the routing weights. The routing kernels and expert parallelism complement each other. With a generic all-to-all, the routing kernels do the reordering on both sides of the From bd5c9260099b0c7e05605d21be86e4be41afc2f6 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:18:34 +0200 Subject: [PATCH 15/52] [Docs] MoE: fold EP into the stage list, surface load balancing in the introduction Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index bbdcbb008b..eec2da6533 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -22,22 +22,21 @@ of expert networks and a router that sends each token to one or more experts. This keeps the activated parameter count per token small while allowing the model to scale to many more total parameters. -A token passes through an MoE layer in four stages: +A token passes through an MoE layer in the following stages: #. The **router** scores the experts for each token and selects the top-k of them. #. **Token dispatch** gathers the tokens into expert-contiguous order. +#. With **expert parallelism** - experts sharded across devices - an + **all-to-all dispatch** sends each token to the rank that owns its expert. #. The **grouped MLP** (the experts) runs a single batched computation over all - expert blocks. + expert blocks local to the device. +#. With expert parallelism, an **all-to-all combine** returns the expert outputs + to the rank the token came from. #. **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. -When the experts do not fit on one device, they are sharded across devices -(**expert parallelism**). The layer then gains two all-to-all collectives: a -dispatch that sends each token to the rank owning its expert, and a combine that -returns the results to the source rank. The experts themselves stay local. - .. raw:: html :file: img/moe_layer_ep.svg @@ -51,9 +50,11 @@ 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: * :ref:`Routing kernels `: the router fuses the score - function with the top-k selection, and token dispatch and combine move tokens - between their original order and the expert-contiguous layout with optimized - kernels instead of Python-level gather / sort / concatenate chains. + function with the top-k selection, a fused :ref:`load-balancing loss + ` keeps the routing spread evenly across experts, and token + dispatch and combine move tokens between their original order and the + expert-contiguous layout with optimized kernels instead of Python-level + gather / sort / concatenate chains. * :ref:`Grouped GEMM ` primitives execute the expert linear layers efficiently once the tokens are laid out in expert-contiguous blocks, and the :ref:`grouped MLP ` fuses the whole expert MLP into @@ -125,6 +126,8 @@ architectures: :start-after: # START_ROUTER_JAX :end-before: # END_ROUTER_JAX +.. _moe-load-balancing: + Load balancing ~~~~~~~~~~~~~~ @@ -137,6 +140,11 @@ respect to every expert's logit. Those dense scores come from ``fused_compute_score_for_moe_aux_loss`` in PyTorch, or from ``fused_topk_with_score_function(..., compute_aux_scores=True)`` in JAX. +An alternative that needs no auxiliary loss is to bias the selection directly: +with the sigmoid score function, the router's ``expert_bias`` shifts which experts +are selected without changing the returned weights, so it can be adjusted between +steps to steer load towards under-used experts. + .. tabs:: .. tab:: PyTorch From 23392d143b33ccb753930aa1f0e038a004cd1cc0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:23:09 +0200 Subject: [PATCH 16/52] [Docs] MoE: explain grouped routing, simplify router wording Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index eec2da6533..fac5ea93f3 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -27,7 +27,7 @@ A token passes through an MoE layer in the following stages: #. The **router** scores the experts for each token and selects the top-k of them. #. **Token dispatch** gathers the tokens into expert-contiguous order. -#. With **expert parallelism** - experts sharded across devices - an +#. With expert parallelism - experts sharded across devices - an **all-to-all dispatch** sends each token to the rank that owns its expert. #. The **grouped MLP** (the experts) runs a single batched computation over all expert blocks local to the device. @@ -85,10 +85,8 @@ the two tensors that drive the rest of the layer: these as merging weights when a token was routed to more than one expert. Transformer Engine fuses the score function and the top-k selection into a single -differentiable kernel, exposed as ``fused_topk_with_score_function`` in both -``transformer_engine.pytorch.router`` and ``transformer_engine.jax.router``. All -internal math runs in FP32 for numerical stability, regardless of the logits -dtype. +differentiable kernel, ``fused_topk_with_score_function``. All internal math runs +in FP32 for numerical stability, regardless of the logits dtype. .. raw:: html :file: img/moe_router.svg @@ -103,8 +101,12 @@ architectures: * **Score function:** ``"softmax"`` or ``"sigmoid"`` (the PyTorch API also offers ``"sqrtsoftplus"``). With softmax, ``use_pre_softmax`` selects whether the softmax is applied before or after the top-k. -* **Grouped (device-limited) routing:** ``num_groups`` and ``group_topk`` restrict - selection to a subset of expert groups, as in DeepSeek-style routing. +* **Grouped routing:** the experts are split into ``num_groups`` equal groups + (for example, one group per node). Each group is scored by the sum of its best + expert scores, the top ``group_topk`` groups are kept, and the final top-k + experts are chosen only from those groups. This bounds how many groups a + token's experts span, which limits all-to-all traffic under expert + parallelism (the node-limited routing of DeepSeek-V3). * **Expert bias:** with the sigmoid score function, ``expert_bias`` shifts the selection without changing the returned weights - the bias-adjustment scheme used for auxiliary-loss-free load balancing. From ccdfef4247bacd0e76d42a7850a9491a5824077f Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:26:04 +0200 Subject: [PATCH 17/52] [Docs] MoE: deduplicate expert_bias description Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index fac5ea93f3..08cb589ece 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -107,9 +107,8 @@ architectures: experts are chosen only from those groups. This bounds how many groups a token's experts span, which limits all-to-all traffic under expert parallelism (the node-limited routing of DeepSeek-V3). -* **Expert bias:** with the sigmoid score function, ``expert_bias`` shifts the - selection without changing the returned weights - the bias-adjustment scheme - used for auxiliary-loss-free load balancing. +* **Expert bias:** ``expert_bias`` is added to the scores before the top-k + selection (see :ref:`Load balancing `). * **Scaling:** ``scaling_factor`` rescales the returned probabilities. .. tabs:: @@ -142,10 +141,12 @@ respect to every expert's logit. Those dense scores come from ``fused_compute_score_for_moe_aux_loss`` in PyTorch, or from ``fused_topk_with_score_function(..., compute_aux_scores=True)`` in JAX. -An alternative that needs no auxiliary loss is to bias the selection directly: -with the sigmoid score function, the router's ``expert_bias`` shifts which experts -are selected without changing the returned weights, so it can be adjusted between -steps to steer load towards under-used experts. +An alternative that needs no auxiliary loss is to bias the selection directly. +With the sigmoid score function, the router's ``expert_bias`` is added to the +scores only for the top-k selection, so it changes which experts are picked but +not the returned routing weights. Adjusting it between steps - lowering it for +overloaded experts and raising it for under-used ones - steers the load without +touching the training objective. .. tabs:: From 9aeff33f61e47951d58c8bc3ec4be526e055022d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:30:42 +0200 Subject: [PATCH 18/52] [Docs] MoE: move routing-kernels lead paragraphs above the router subsection Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 41 ++++++++++--------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 08cb589ece..cc1394ede8 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -70,6 +70,27 @@ complete MoE layer. Routing kernels --------------- +The router decides where each token goes; once it has produced a routing map, +the tokens must be moved into the expert-contiguous layout expected by the +grouped GEMM (``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX) and, +afterwards, moved back. Transformer Engine provides differentiable kernels for +all of these steps. This section starts with the router and then focuses on the +two core data-movement operations - token dispatch and token combine - because +they illustrate the layout transformation used by the other variants. + +The token dispatch and combine snippets below show one concrete instance of this +pattern: the mask-map routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` +and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as +``transformer_engine.jax.permutation.token_dispatch`` and +``transformer_engine.jax.permutation.token_combine``. Other routing variants +(for example, index-map routing in PyTorch via ``map_type="index"``) are +available in both frameworks and follow the same pattern; see the +:doc:`PyTorch API reference ` and +:doc:`JAX API reference ` for the complete list and signatures. +The mask-map APIs have different framework-specific wrappers, but lower to the +same shared Triton permutation kernels, and both pairs are differentiable so they +can be used directly inside training graphs. + .. _moe-router: Router @@ -164,26 +185,6 @@ touching the training objective. :start-after: # START_ROUTER_AUX_JAX :end-before: # END_ROUTER_AUX_JAX -Once the router has produced a routing map, the tokens must be moved into the -expert-contiguous layout expected by the grouped GEMM (``GroupedLinear`` in -PyTorch, ``grouped_dense`` in JAX) and, afterwards, moved back. Transformer -Engine provides differentiable kernels for both directions. This section focuses on -the two core operations - token dispatch and token combine - because they -illustrate the layout transformation used by the other variants. - -The snippets below show one concrete instance of this pattern: the mask-map -routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` -and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as -``transformer_engine.jax.permutation.token_dispatch`` and -``transformer_engine.jax.permutation.token_combine``. Other routing variants -(for example, index-map routing in PyTorch via ``map_type="index"``) are -available in both frameworks and follow the same pattern; see the -:doc:`PyTorch API reference ` and -:doc:`JAX API reference ` for the complete list and signatures. -The mask-map APIs have different framework-specific wrappers, but lower to the -same shared Triton permutation kernels, and both pairs are differentiable so they -can be used directly inside training graphs. - Token dispatch ~~~~~~~~~~~~~~ From 9c52362d9396384e52e80f8fd839f7a90b8192fe Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:38:32 +0200 Subject: [PATCH 19/52] [Docs] MoE: describe what the blocks do and how to use them; move framework specifics to snippets Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 383 +++++++----------- 1 file changed, 136 insertions(+), 247 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index cc1394ede8..812520e819 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -17,12 +17,9 @@ Mixture of Experts Introduction ------------ -Mixture of Experts (MoE) layers replace a dense feed-forward network with a set +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. -This keeps the activated parameter count per token small while allowing the -model to scale to many more total parameters. - -A token passes through an MoE layer in the following stages: +A token passes through the layer in the following stages: #. The **router** scores the experts for each token and selects the top-k of them. @@ -45,22 +42,19 @@ the* ``routing_map`` *consumed by token dispatch and the* ``probs`` *used as mer weights in token combine; the all-to-all dispatch and combine are only present when the experts are sharded across ranks.* -Transformer Engine provides an optimized 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: - -* :ref:`Routing kernels `: the router fuses the score - function with the top-k selection, a fused :ref:`load-balancing loss - ` keeps the routing spread evenly across experts, and token - dispatch and combine move tokens between their original order and the - expert-contiguous layout with optimized kernels instead of Python-level - gather / sort / concatenate chains. -* :ref:`Grouped GEMM ` primitives execute the expert linear - layers efficiently once the tokens are laid out in expert-contiguous blocks, - and the :ref:`grouped MLP ` fuses the whole expert MLP into - one kernel. -* :ref:`Expert parallelism ` shards the experts across - devices with all-to-all dispatch and combine. +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: + +* :ref:`Routing kernels `: a fused router (score function + and top-k selection), a fused :ref:`load-balancing loss `, + and token dispatch and combine kernels that move tokens between their original + order and the expert-contiguous layout. +* :ref:`Grouped GEMM `: the expert linear layers as one call + over expert-contiguous blocks; the :ref:`grouped MLP ` fuses + the whole expert MLP into one kernel. +* :ref:`Expert parallelism `: all-to-all dispatch and + combine for experts sharded across devices. The :ref:`example at the end ` wires the blocks into a complete MoE layer. @@ -70,26 +64,12 @@ complete MoE layer. Routing kernels --------------- -The router decides where each token goes; once it has produced a routing map, -the tokens must be moved into the expert-contiguous layout expected by the -grouped GEMM (``GroupedLinear`` in PyTorch, ``grouped_dense`` in JAX) and, -afterwards, moved back. Transformer Engine provides differentiable kernels for -all of these steps. This section starts with the router and then focuses on the -two core data-movement operations - token dispatch and token combine - because -they illustrate the layout transformation used by the other variants. - -The token dispatch and combine snippets below show one concrete instance of this -pattern: the mask-map routing path, exposed in PyTorch as ``transformer_engine.pytorch.moe_permute`` -and ``transformer_engine.pytorch.moe_unpermute``, and in JAX as -``transformer_engine.jax.permutation.token_dispatch`` and -``transformer_engine.jax.permutation.token_combine``. Other routing variants -(for example, index-map routing in PyTorch via ``map_type="index"``) are -available in both frameworks and follow the same pattern; see the -:doc:`PyTorch API reference ` and -:doc:`JAX API reference ` for the complete list and signatures. -The mask-map APIs have different framework-specific wrappers, but lower to the -same shared Triton permutation kernels, and both pairs are differentiable so they -can be used directly inside training graphs. +The router produces a routing map. Token dispatch moves the tokens into the +expert-contiguous layout expected by the grouped GEMM, and token combine moves +the expert outputs back. All of these kernels are differentiable. The snippets +below use the mask-map routing variant; other variants (for example index-map +routing) follow the same pattern, see the :doc:`PyTorch API reference +` and :doc:`JAX API reference `. .. _moe-router: @@ -105,9 +85,9 @@ the two tensors that drive the rest of the layer: * ``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. -Transformer Engine fuses the score function and the top-k selection into a single -differentiable kernel, ``fused_topk_with_score_function``. All internal math runs -in FP32 for numerical stability, regardless of the logits dtype. +``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. .. raw:: html :file: img/moe_router.svg @@ -116,18 +96,14 @@ in FP32 for numerical stability, regardless of the logits dtype. The selected entries populate* ``routing_map`` *(a 0/1 mask) and* ``probs`` *(the routing weights); all other entries are zero.* -The kernel covers the score functions and selection variants used by common MoE -architectures: - -* **Score function:** ``"softmax"`` or ``"sigmoid"`` (the PyTorch API also offers - ``"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 - (for example, one group per node). Each group is scored by the sum of its best - expert scores, the top ``group_topk`` groups are kept, and the final top-k - experts are chosen only from those groups. This bounds how many groups a - token's experts span, which limits all-to-all traffic under expert - parallelism (the node-limited routing of DeepSeek-V3). +Options: + +* **Score function:** softmax or sigmoid. 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 :ref:`Load balancing `). * **Scaling:** ``scaling_factor`` rescales the returned probabilities. @@ -153,21 +129,18 @@ architectures: Load balancing ~~~~~~~~~~~~~~ -Left unconstrained, a router tends to collapse onto a handful of experts. The -usual remedy is an auxiliary load-balancing loss that rewards spreading tokens -evenly across experts. Transformer Engine computes it with ``fused_moe_aux_loss`` -from the per-expert token counts and the *dense* routing scores - one value per -expert rather than only the selected top-k - so the loss has a gradient with -respect to every expert's logit. Those dense scores come from -``fused_compute_score_for_moe_aux_loss`` in PyTorch, or from -``fused_topk_with_score_function(..., compute_aux_scores=True)`` in JAX. - -An alternative that needs no auxiliary loss is to bias the selection directly. -With the sigmoid score function, the router's ``expert_bias`` is added to the -scores only for the top-k selection, so it changes which experts are picked but -not the returned routing weights. Adjusting it between steps - lowering it for -overloaded experts and raising it for under-used ones - steers the load without -touching the training objective. +``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. + +``expert_bias`` balances the load without an extra loss term. With the sigmoid +score function 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 and raise it for under-used +ones. .. tabs:: @@ -188,15 +161,10 @@ touching the training objective. Token dispatch ~~~~~~~~~~~~~~ -Token dispatch is the canonical routing operation: given the original token -tensor and a routing map describing each token's destination expert, it returns -a permuted token buffer in which all rows assigned to the same expert are -stored contiguously. In PyTorch this operation is exposed as ``moe_permute``; -in JAX it is exposed as ``token_dispatch``. This is exactly the layout that -the grouped linear layer consumes via its per-expert token-count argument -(``m_splits`` in PyTorch ``GroupedLinear``, ``group_sizes`` in JAX -``grouped_dense``), so token dispatch followed by the grouped GEMM forms a -typical MoE forward block. +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. .. raw:: html :file: img/moe_permute.svg @@ -223,30 +191,22 @@ A typical call looks like: :start-after: # START_MOE_PERMUTE_JAX :end-before: # END_MOE_PERMUTE_JAX -Both variants return the permuted token buffer of shape -``[num_out_tokens, hidden_size]`` together with a ``row_id_map`` that -carries enough information for token combine to restore the original token -order once the expert computation is done. Token dispatch and token combine -are typically used as a matched pair around the grouped GEMM call. +The 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 routing operation: it takes the expert-contiguous -output produced by the grouped GEMM (or any per-expert computation) and the -``row_id_map`` returned by token dispatch, and returns a single tensor of -shape ``[num_tokens, hidden_size]`` with the rows written back into the -original token order. In PyTorch this operation is exposed as -``moe_unpermute``; in JAX it is exposed as ``token_combine``. - -For top-1 routing each token has exactly one expert contribution, so -``merging_probs`` is omitted. For top-k routing pass the per-token expert -weights as ``merging_probs`` and the kernel computes a weighted sum of the -per-expert contributions in the same fused pass; without it the per-expert -contributions are summed unweighted. In PyTorch, also pass -``restore_shape=(num_tokens, hidden_size)`` whenever the permuted buffer has more -rows than the original tokens (top-k routing); JAX infers the original token -count from the ``row_id_map``. +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. + +For top-k routing 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 it the contributions are summed unweighted; for top-1 +routing it is not needed. .. raw:: html :file: img/moe_unpermute.svg @@ -277,27 +237,20 @@ A typical call looks like: Token probabilities ~~~~~~~~~~~~~~~~~~~ -In top-k routing each token contributes to several experts, and those -contributions are recombined using the routing weights. There are two equivalent -places to apply the weights: +The routing weights can be applied in two equivalent places: -* **At combine (output side).** Pass the routing weights to token combine as - ``merging_probs``; it forms the weighted sum of the per-expert contributions in - the same fused pass. This is the path used in the examples above. -* **At dispatch (input side).** Scale each expert's input by its routing weight - before the grouped GEMM. ``moe_permute_with_probs`` (PyTorch) and the ``probs`` - argument of ``token_dispatch`` (JAX) permute a probability tensor alongside the - tokens, so the weights arrive already aligned with the expert-contiguous - layout. +* **At combine (output side).** Pass them to token combine as ``merging_probs``, + as in the examples above. +* **At dispatch (input side).** Pass them to token dispatch as ``probs``. They + are permuted alongside the tokens into the expert-contiguous layout, so each + expert's input can be scaled before the grouped GEMM. Padding and alignment ~~~~~~~~~~~~~~~~~~~~~ -Grouped GEMM backends are most efficient when each expert's token block starts at -an aligned offset (for example, a multiple of 128 rows). Because the number of -tokens routed to an expert is data dependent, the blocks are generally ragged. -Transformer Engine can pad each block up to a multiple of ``align_size`` as part -of the dispatch kernel, avoiding a separate padding pass. +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. .. raw:: html :file: img/moe_padding.svg @@ -306,11 +259,9 @@ of the dispatch kernel, avoiding a separate padding pass. The per-expert padding offsets are returned so that token combine can drop the padding again.* -In PyTorch this is ``moe_permute_and_pad_with_probs``; in JAX it is the -``align_size`` argument of ``token_dispatch``. Both return the padded token -buffer, the aligned per-expert token counts (used as ``m_splits`` / -``group_sizes`` for the grouped GEMM), and the per-expert ``pad_offsets`` that -token combine needs in order to remove the padding. +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. .. tabs:: @@ -331,28 +282,20 @@ token combine needs in order to remove the padding. Reordering expert chunks ~~~~~~~~~~~~~~~~~~~~~~~~ -When experts are sharded across devices, the per-expert token blocks often have -to be reordered - for example, to regroup tokens by destination rank before an -all-to-all, or to restore the original grouping afterwards. -``moe_sort_chunks_by_index`` (PyTorch) and ``sort_chunks_by_index`` (JAX) permute -contiguous chunks of a token tensor according to a list of chunk sizes and a -permutation of chunk indices, without falling back to Python-level slicing and -concatenation. ``moe_sort_chunks_by_index_with_probs`` reorders an accompanying -probability tensor in the same call. +The sort-chunks-by-index kernels permute contiguous chunks of a token tensor +according to a list of chunk sizes and a permutation of chunk indices, for +example to regroup tokens by destination rank before an all-to-all and to +restore the original grouping afterwards. A ``_with_probs`` variant reorders an +accompanying probability tensor in the same call. See the API reference for the +signatures. .. _moe-grouped-gemm: Grouped GEMM ------------ -The straightforward way to apply per-expert linear layers is to loop over the -experts and call a separate ``Linear`` for each one. This is correct, but it -is not the most efficient way to execute many expert GEMMs. - -Transformer Engine provides a grouped GEMM primitive -(``GroupedLinear`` in PyTorch and ``grouped_dense`` in JAX) - an optimized -replacement that produces the same outputs as the loop while using -implementations that are better suited for MoE workloads. +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: @@ -367,20 +310,18 @@ The full layer output is the concatenation of all expert outputs: Y = \mathrm{concat}(Y_0, Y_1, \ldots, Y_{G-1}) -The grouped GEMM is told how many token rows belong to each expert via a -per-expert token-count argument: ``m_splits`` in PyTorch ``GroupedLinear`` and -``group_sizes`` in JAX ``grouped_dense``. +The number of token rows belonging to each expert is passed as a per-expert +token-count argument. .. raw:: html :file: img/grouped_linear.svg *Figure 6. Both paths produce the same outputs from the same inputs. The -baseline launches one* ``Linear`` *per expert, while the grouped GEMM -(*\ ``GroupedLinear`` *in PyTorch,* ``grouped_dense`` *in JAX) is an optimized -grouped implementation that replaces the loop.* +baseline launches one* ``Linear`` *per expert; the grouped GEMM replaces the +loop with one call.* -The following snippets show how to replace the loop with the grouped GEMM. -They assume the tokens have already been permuted into expert-contiguous order. +The snippets assume the tokens have already been permuted into +expert-contiguous order. .. tabs:: @@ -398,49 +339,27 @@ They assume the tokens have already been permuted into expert-contiguous order. :start-after: # START_GROUPED_LINEAR_JAX :end-before: # END_GROUPED_LINEAR_JAX -The grouped GEMM uses implementations tuned for grouped expert execution: - -* **Optimized backends:** Transformer Engine selects from several grouped GEMM - backends depending on the framework, datatype, and GPU architecture. This - can be, for example, cuBLAS GEMMs launched on multiple CUDA streams or a - single grouped GEMM kernel, among other backend-specific implementations. -* **Recipe compatibility:** the grouped GEMM is integrated with - Transformer Engine's :doc:`low-precision training stack - `, so the recipes available to - regular ``Linear`` layers can also be used for MoE experts. The exact set - depends on the execution path, GPU architecture, and cuBLAS version; see the - ``GroupedLinear`` / ``grouped_dense`` API reference for the current - constraints. -* **Fused quantization:** Low-precision grouped GEMM paths can fuse - quantization-related work such as scale computation, casting, and - cast/transpose steps across experts instead of repeating the same work in a - Python loop. -* **Fused expert MLP:** Through the :doc:`operation-based API - `, the two expert GEMMs and the activation - between them can be fused into a single grouped operation on recent - architectures; see :ref:`Grouped MLP `. - -The PyTorch ``GroupedLinear`` module also supports the features expected of a -Transformer Engine linear layer - tensor and sequence parallelism, gradient -accumulation fusion, and FP8 weight caching - so it can serve as a drop-in expert -layer. See the :doc:`PyTorch API reference ` for the full -signature. +* **Backends:** the grouped GEMM backend is selected based on datatype and GPU + architecture, for example cuBLAS GEMMs on multiple CUDA streams or a single + grouped GEMM kernel. +* **Low-precision recipes:** the grouped GEMM works with the + :doc:`low-precision training recipes ` + available to ``Linear``. The supported set depends on the GPU architecture and + cuBLAS version; see the API reference for the current constraints. +* **Fused quantization:** in low-precision paths the scale computation, casting + and cast/transpose steps are fused across experts. +* **Fused expert MLP:** the two expert GEMMs and the activation between them can + be fused into one operation, see :ref:`Grouped MLP `. .. _moe-grouped-mlp: Grouped MLP ----------- -An expert MLP is two grouped GEMMs with an activation between them: the first -projects into the (gated) feed-forward dimension, the activation is applied, and -the second projects back. Running these as separate kernels writes the large -intermediate activation out to HBM and reads it back for the second GEMM, and -re-quantizes it in a separate pass. - -On Blackwell (SM100) GPUs, Transformer Engine can fuse the whole expert MLP - -both grouped GEMMs and the activation - into a single CuTe DSL kernel. The -intermediate stays on chip and the cross-expert quantization is folded into the -GEMMs, removing the HBM round-trip and the extra kernel launches. +An expert MLP is two grouped GEMMs with an activation between them. On +Blackwell (SM100) GPUs the whole expert MLP can run as a single CuTe DSL kernel: +the intermediate activation stays on chip and its quantization is folded into +the GEMMs. .. raw:: html :file: img/moe_grouped_mlp.svg @@ -449,11 +368,9 @@ GEMMs, removing the HBM round-trip and the extra kernel launches. and the second grouped GEMM with a single fused grouped-MLP kernel that keeps the intermediate on chip.* -The fusion is exposed through the operation-based API and applied automatically -by the :doc:`operation fuser `: when it sees a -grouped linear, a scaled GLU (or SReLU) activation, and another grouped linear in -sequence, it replaces them with one fused grouped-MLP operation. No change to the -forward code is needed to opt in. +The fusion is applied by the :doc:`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. .. tabs:: @@ -470,8 +387,8 @@ forward code is needed to opt in. :start-after: # START_GROUPED_MLP_PYTORCH :end-before: # END_GROUPED_MLP_PYTORCH -The fused path is taken when all of the following hold; otherwise the three ops -run separately and produce identical results: +The fused path is taken when all of the following hold; otherwise the three +operations run separately with identical results: * **Architecture:** Blackwell (SM100) with cuDNN frontend 1.23 or newer. * **Recipe:** a block-scaled low-precision recipe - MXFP8, or NVFP4 with the @@ -491,17 +408,11 @@ Expert parallelism or newer. It is compiled in by default when Transformer Engine is built for these architectures; set ``NVTE_WITH_NCCL_EP=0`` at build time to disable it. -The grouped GEMM keeps all experts on a single device. When the experts no longer -fit there - or to add another dimension of parallelism - they are sharded across -devices, a scheme called expert parallelism (EP). Each device then owns only a -slice of the experts, so a token routed to a non-local expert has to travel to -the device that owns it. - -That data movement is two all-to-all collectives wrapped around the local expert +With expert parallelism (EP) the experts are sharded across devices, and each +device owns a slice of them. Two all-to-all collectives wrap the local expert computation: a **dispatch** all-to-all sends each token to the rank that owns its expert, the local grouped GEMM runs, and a **combine** all-to-all returns the -results to the source rank. It is the distributed counterpart of the -:ref:`token dispatch and token combine kernels `. +results to the source rank. .. raw:: html :file: img/moe_expert_parallel.svg @@ -510,24 +421,12 @@ results to the source rank. It is the distributed counterpart of the token to the rank owning its expert and a combine all-to-all returns the outputs to the source rank.* -Transformer Engine implements dispatch and combine directly on NCCL, using -NCCL symmetric-memory windows for zero-copy transfers. The backend is a common -C API (``nvte_ep_dispatch`` / ``nvte_ep_combine`` and their backward passes, -declared in ``transformer_engine/common/include/transformer_engine/ep.h``) that -both frameworks build on: - -* **PyTorch.** ``transformer_engine.pytorch.ep`` exposes the primitives with - autograd support. ``ep_bootstrap`` initializes EP once per process on an - existing process group, an ``EpBuffer`` holds the per-call state, and - ``ep_dispatch`` / ``ep_combine`` perform the two all-to-alls. The routing - itself comes from the :ref:`router `; the local experts run - on the receive buffer between the two calls. -* **JAX.** ``transformer_engine.jax.moe.moe`` runs the entire layer - router, - dispatch, grouped expert GEMMs, and combine - as a single differentiable call. - ``ep_axis`` names the mesh axis the experts are sharded over, and the dispatch - and combine steps become all-to-all collectives over that axis. The underlying - primitives are also available separately in ``transformer_engine.jax.ep``. - This API is currently experimental. +Dispatch and combine are implemented directly on NCCL, using symmetric-memory +windows for zero-copy transfers. Both are differentiable. The common C API +(``nvte_ep_dispatch`` / ``nvte_ep_combine`` and their backward passes, declared +in ``transformer_engine/common/include/transformer_engine/ep.h``) is exposed in +both frameworks; the snippets show how the dispatch, the local experts and the +combine are wired together. .. tabs:: @@ -560,27 +459,26 @@ both frameworks build on: Sizing the receive buffer ~~~~~~~~~~~~~~~~~~~~~~~~~ -Each rank receives a data-dependent number of tokens per step. Passing -``recv_capacity_per_rank`` fixes the size of the receive buffer up front, so the -step needs no device-to-host synchronization and can be captured in a CUDA graph; -the dropless worst case is ``ep_size * max_tokens_per_rank * top_k``. Omitting it -selects eager mode, which sizes the buffer from the actual receive count each -step at the cost of a host sync. +Each rank receives a data-dependent number of tokens per step. Passing a fixed +receive capacity (``recv_capacity_per_rank``) sizes the receive buffer up front, +so the step needs no device-to-host synchronization and can be captured in a +CUDA graph; the dropless worst case is ``ep_size * max_tokens_per_rank * top_k``. +Without it the buffer is sized from the actual receive count each step, at the +cost of a host sync. -In PyTorch, ``ep_dispatch`` can quantize the tokens on the fly when the -``EpBuffer`` is created with an MXFP8 ``dispatch_fwd_quant_recipe``, so the -all-to-all moves the low-precision payload and the local grouped GEMM consumes -it directly. Complete runnable examples live in ``examples/pytorch/ep/`` and -``examples/jax/ep/`` in the repository. +Dispatch can quantize the tokens to MXFP8 before the all-to-all +(``dispatch_fwd_quant_recipe``), so the communication moves the low-precision +payload and the local grouped GEMM consumes it directly. Complete runnable +examples live in ``examples/pytorch/ep/`` and ``examples/jax/ep/`` in the +repository. .. _moe-putting-it-together: Example: putting it all together -------------------------------- -The building blocks assemble into the four stages from the introduction: route, -dispatch, run the experts, and combine. The example below wires them together -for top-k routing on a single device. +The example below wires the blocks together for top-k routing on a single +device: route, dispatch, run the experts, combine. .. tabs:: @@ -598,25 +496,16 @@ for top-k routing on a single device. :start-after: # START_MOE_LAYER_JAX :end-before: # END_MOE_LAYER_JAX -This 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. The -expert step is built from the :ref:`grouped GEMM `; a full -expert MLP stacks two grouped GEMMs around an activation. Every stage is differentiable, so the assembled layer -trains end to end. - -When the experts are sharded across devices, the same layer gains the two -all-to-all collectives from Figure 1 around the local experts: token dispatch -groups the tokens by destination rank, the all-to-all dispatch moves them to the -ranks owning their experts, the local grouped MLP runs, and the all-to-all -combine returns the outputs before token combine restores the original order and -applies the routing weights. +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. -The routing kernels and expert parallelism complement each other. With a generic +With experts sharded across devices there are two options. With a generic all-to-all, the routing kernels do the reordering on both sides of the communication: tokens are sorted by destination rank before the all-to-all and regrouped by local expert after it (see :ref:`Reordering expert chunks `). With the NCCL-based :ref:`expert parallelism -` primitives, this permutation is folded into the -communication itself: the dispatch delivers an expert-contiguous receive buffer -and the combine writes the results straight back into the original token order, -so no separate permute or unpermute is needed on the local side. +` primitives the permutation is folded into the +communication: the dispatch delivers an expert-contiguous receive buffer and the +combine writes the results straight back into the original token order, so no +separate token dispatch or combine is needed. From 8cae00b49c5180f64bf475667f78dc48e1fd9b4f Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:43:58 +0200 Subject: [PATCH 20/52] [Docs] MoE: describe grouped GEMM execution paths; shorten grouped MLP conditions Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/grouped_linear_jax.py | 5 ++- .../grouped_linear_pytorch.py | 18 +++++++- .../mixture_of_experts/mixture_of_experts.rst | 42 ++++++++++--------- 3 files changed, 43 insertions(+), 22 deletions(-) diff --git a/docs/features/mixture_of_experts/grouped_linear_jax.py b/docs/features/mixture_of_experts/grouped_linear_jax.py index 94042cba21..3b0e89e141 100644 --- a/docs/features/mixture_of_experts/grouped_linear_jax.py +++ b/docs/features/mixture_of_experts/grouped_linear_jax.py @@ -20,7 +20,10 @@ axis=0, ) -# Transformer Engine: one grouped dense call. +# Transformer Engine: one grouped dense call. group_sizes is a device array. +# On Blackwell, BF16 and MXFP8 inputs without bias run as a single grouped GEMM +# with the group sizes kept on the device; other cases launch one GEMM per +# expert and copy group_sizes to the host first. grouped_out = te_dense.grouped_dense( x, kernel, diff --git a/docs/features/mixture_of_experts/grouped_linear_pytorch.py b/docs/features/mixture_of_experts/grouped_linear_pytorch.py index f927e0ac9a..d2320e8eda 100644 --- a/docs/features/mixture_of_experts/grouped_linear_pytorch.py +++ b/docs/features/mixture_of_experts/grouped_linear_pytorch.py @@ -19,7 +19,8 @@ dim=0, ) -# Transformer Engine: one grouped linear call. +# Transformer Engine: one grouped linear call. By default one GEMM per expert +# is launched; m_splits is read on the host. grouped_linear = te.GroupedLinear( num_experts, hidden_size, @@ -28,4 +29,19 @@ params_dtype=torch.bfloat16, ).cuda() grouped_out = grouped_linear(x, m_splits) + +# Single grouped GEMM with the token counts on the device (no host sync, +# CUDA-graph capturable): opt in with use_grouped_tensor=True and pass +# m_splits as a CUDA int64 tensor. Falls back to per-expert GEMMs when the +# recipe / GPU / cuBLAS version does not support it. +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) # END_GROUPED_LINEAR_PYTORCH diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 812520e819..5ed993f84a 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -339,17 +339,25 @@ expert-contiguous order. :start-after: # START_GROUPED_LINEAR_JAX :end-before: # END_GROUPED_LINEAR_JAX -* **Backends:** the grouped GEMM backend is selected based on datatype and GPU - architecture, for example cuBLAS GEMMs on multiple CUDA streams or a single - grouped GEMM kernel. -* **Low-precision recipes:** the grouped GEMM works with the - :doc:`low-precision training recipes ` - available to ``Linear``. The supported set depends on the GPU architecture and - cuBLAS version; see the API reference for the current constraints. -* **Fused quantization:** in low-precision paths the scale computation, casting - and cast/transpose steps are fused across experts. -* **Fused expert MLP:** the two expert GEMMs and the activation between them can - be fused into one operation, see :ref:`Grouped MLP `. +The grouped GEMM works with the :doc:`low-precision training recipes +` available to ``Linear``: the inputs +are quantized per expert and the expert GEMMs run in the recipe's precision. + +There are two execution paths: + +* **Per-expert GEMMs.** The per-expert token counts are read on the host, the + input is split and quantized per expert, and one cuBLAS GEMM per expert is + launched on a pool of CUDA streams (on Hopper, ``NVTE_USE_CUTLASS_GROUPED_GEMM=1`` + switches BF16/FP16 to a CUTLASS grouped GEMM kernel). This path supports all + recipes, but reading the token counts is a device-to-host synchronization, so + it cannot be captured in a CUDA graph. +* **Single grouped GEMM.** The token counts stay on the device and all experts + run as one cuBLASLt grouped GEMM (cuBLAS 13.3 or newer), with the + quantization fused across experts. There is no host synchronization, so the + step is CUDA-graph capturable. Supported for BF16/FP16, and for MXFP8 and NVFP4 + on Blackwell; FP8 current scaling and FP8 block scaling on Hopper need cuBLAS + 13.5 / 13.6. FP8 delayed scaling and custom recipes are not supported on this + path. The snippets show how it is selected. .. _moe-grouped-mlp: @@ -387,15 +395,9 @@ in sequence are replaced with one fused grouped-MLP operation. :start-after: # START_GROUPED_MLP_PYTORCH :end-before: # END_GROUPED_MLP_PYTORCH -The fused path is taken when all of the following hold; otherwise the three -operations run separately with identical results: - -* **Architecture:** Blackwell (SM100) with cuDNN frontend 1.23 or newer. -* **Recipe:** a block-scaled low-precision recipe - MXFP8, or NVFP4 with the - randomized Hadamard transform enabled. -* **Opt-in:** the environment variable ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``. -* **Activation:** a scaled ``SwiGLU`` / ``GeGLU`` (gated) or ``SReLU`` (unary), - with feature dimensions aligned to 64 and the token count to 128. +The fusion is enabled with ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` and requires +Blackwell and a block-scaled recipe (MXFP8 or NVFP4). When the configuration is +not supported, the three operations run separately with identical results. .. _moe-expert-parallelism: From 4960ac367d7435f79d2e859ea56156d38cdfbe09 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:49:44 +0200 Subject: [PATCH 21/52] [Docs] MoE: move grouped GEMM paths above snippet; simplify grouped MLP figure Signed-off-by: Pawel Gadzinski --- .../img/moe_grouped_mlp.svg | 48 ++++++------------- .../mixture_of_experts/mixture_of_experts.rst | 43 ++++++++--------- 2 files changed, 35 insertions(+), 56 deletions(-) diff --git a/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg b/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg index 3f2d614476..cf004c87e6 100644 --- a/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg +++ b/docs/features/mixture_of_experts/img/moe_grouped_mlp.svg @@ -1,14 +1,11 @@ - + @@ -16,37 +13,20 @@ - Fused grouped MLP + + FC1 grouped GEMM + - - unfused: three kernels, with the intermediate written to and re-read from HBM - - FC1 - grouped GEMM + + activation + - - SwiGLU + + FC2 grouped GEMM - - FC2 - grouped GEMM + + operation fuser - - - HBM - - - - HBM - - - - - op fuser - - - - Fused grouped MLP — single CuTe DSL kernel - FC1 + SwiGLU + FC2, intermediate kept on-chip - Blackwell (SM100), MXFP8 or NVFP4; activation also GeGLU or SReLU + + fused grouped MLP (one kernel) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 5ed993f84a..c7744d689f 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -320,25 +320,6 @@ token-count argument. baseline launches one* ``Linear`` *per expert; the grouped GEMM replaces the loop with one call.* -The snippets assume the tokens have already been permuted into -expert-contiguous order. - -.. tabs:: - - .. tab:: PyTorch - - .. literalinclude:: grouped_linear_pytorch.py - :language: python - :start-after: # START_GROUPED_LINEAR_PYTORCH - :end-before: # END_GROUPED_LINEAR_PYTORCH - - .. tab:: JAX - - .. literalinclude:: grouped_linear_jax.py - :language: python - :start-after: # START_GROUPED_LINEAR_JAX - :end-before: # END_GROUPED_LINEAR_JAX - The grouped GEMM works with the :doc:`low-precision training recipes ` available to ``Linear``: the inputs are quantized per expert and the expert GEMMs run in the recipe's precision. @@ -359,6 +340,25 @@ There are two execution paths: 13.5 / 13.6. FP8 delayed scaling and custom recipes are not supported on this path. The snippets show how it is selected. +The snippets assume the tokens have already been permuted into +expert-contiguous order. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: grouped_linear_pytorch.py + :language: python + :start-after: # START_GROUPED_LINEAR_PYTORCH + :end-before: # END_GROUPED_LINEAR_PYTORCH + + .. tab:: JAX + + .. literalinclude:: grouped_linear_jax.py + :language: python + :start-after: # START_GROUPED_LINEAR_JAX + :end-before: # END_GROUPED_LINEAR_JAX + .. _moe-grouped-mlp: Grouped MLP @@ -372,9 +372,8 @@ the GEMMs. .. raw:: html :file: img/moe_grouped_mlp.svg -*Figure 7. The operation fuser replaces the first grouped GEMM, the activation, -and the second grouped GEMM with a single fused grouped-MLP kernel that keeps -the intermediate on chip.* +*Figure 7. The operation fuser replaces the two grouped GEMMs and the activation +between them with a single fused grouped-MLP kernel.* The fusion is applied by the :doc:`operation fuser `: a grouped linear, a scaled GLU (or SReLU) activation and another grouped linear From f8eed3df0240ba435ef9c97271c13a1ece628d0b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:54:00 +0200 Subject: [PATCH 22/52] [Docs] MoE: EP dispatch/combine replace token dispatch/combine; redraw layer figure Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/img/moe_layer_ep.svg | 96 +++++++------------ .../mixture_of_experts/mixture_of_experts.rst | 23 +++-- 2 files changed, 45 insertions(+), 74 deletions(-) diff --git a/docs/features/mixture_of_experts/img/moe_layer_ep.svg b/docs/features/mixture_of_experts/img/moe_layer_ep.svg index e961772b86..3ec9f263da 100644 --- a/docs/features/mixture_of_experts/img/moe_layer_ep.svg +++ b/docs/features/mixture_of_experts/img/moe_layer_ep.svg @@ -1,77 +1,45 @@ - + - - - - Mixture of Experts layer with expert parallelism - - - - Router - - - Token - Dispatch - - - All-to-all - dispatch - - - Grouped MLP - (local experts) - - - All-to-all - combine - - - Token - Combine - - - - tokens - - - routing_map - - - by rank - - - by expert - - - expert out - - - returned - - - output - - - - probs (routing weights) - - - - expert parallelism: tokens cross ranks, experts stay local + + single device + + Router + + + Token dispatch + + + Grouped MLP + + + Token combine + + + expert parallelism (experts sharded across ranks) + + Router + + + All-to-all dispatch + grouped by local expert + + + Grouped MLP + local experts + + + All-to-all combine + back in token order diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index c7744d689f..5df51449b8 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -24,23 +24,23 @@ A token passes through the layer in the following stages: #. The **router** scores the experts for each token and selects the top-k of them. #. **Token dispatch** gathers the tokens into expert-contiguous order. -#. With expert parallelism - experts sharded across devices - an - **all-to-all dispatch** sends each token to the rank that owns its expert. #. The **grouped MLP** (the experts) runs a single batched computation over all - expert blocks local to the device. -#. With expert parallelism, an **all-to-all combine** returns the expert outputs - to the rank the token came from. + expert blocks. #. **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. + .. raw:: html :file: img/moe_layer_ep.svg -*Figure 1. The stages of an MoE layer with expert parallelism. The router produces -the* ``routing_map`` *consumed by token dispatch and the* ``probs`` *used as merging -weights in token combine; the all-to-all dispatch and combine are only present -when the experts are sharded across ranks.* +*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 @@ -413,7 +413,10 @@ With expert parallelism (EP) the experts are sharded across devices, and each device owns a slice of them. Two all-to-all collectives wrap the local expert computation: a **dispatch** all-to-all sends each token to the rank that owns its expert, the local grouped GEMM runs, and a **combine** all-to-all returns the -results to the source rank. +results to the source rank. Dispatch takes the router output (expert indices and +weights) directly and delivers a receive buffer grouped by local expert, and +combine writes the results back in the original token order, so no separate +token dispatch or token combine is needed. .. raw:: html :file: img/moe_expert_parallel.svg From a85211786862ae7771ec065d475b07aa174a6c00 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:58:18 +0200 Subject: [PATCH 23/52] [Docs] MoE: split routing kernels into Router and Token permutation sections Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 37 ++++++++++--------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 5df51449b8..73c74be9d2 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -46,10 +46,11 @@ 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: -* :ref:`Routing kernels `: a fused router (score function - and top-k selection), a fused :ref:`load-balancing loss `, - and token dispatch and combine kernels that move tokens between their original - order and the expert-contiguous layout. +* :ref:`Router `: fused score function and top-k selection, and a + fused :ref:`load-balancing loss `. +* :ref:`Token permutation `: token dispatch and combine + kernels that move tokens between their original order and the + expert-contiguous layout. * :ref:`Grouped GEMM `: the expert linear layers as one call over expert-contiguous blocks; the :ref:`grouped MLP ` fuses the whole expert MLP into one kernel. @@ -59,22 +60,10 @@ dropped into an existing implementation one piece at a time: The :ref:`example at the end ` wires the blocks into a complete MoE layer. -.. _moe-routing-kernels: - -Routing kernels ---------------- - -The router produces a routing map. Token dispatch moves the tokens into the -expert-contiguous layout expected by the grouped GEMM, and token combine moves -the expert outputs back. All of these kernels are differentiable. The snippets -below use the mask-map routing variant; other variants (for example index-map -routing) follow the same pattern, see the :doc:`PyTorch API reference -` and :doc:`JAX API reference `. - .. _moe-router: 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 @@ -158,6 +147,18 @@ ones. :start-after: # START_ROUTER_AUX_JAX :end-before: # END_ROUTER_AUX_JAX +.. _moe-token-permutation: + +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. All of these +kernels are differentiable. The snippets below use the mask-map routing variant; +other variants (for example index-map routing) follow the same pattern, see the +:doc:`PyTorch API reference ` and :doc:`JAX API reference +`. + Token dispatch ~~~~~~~~~~~~~~ @@ -508,7 +509,7 @@ With experts sharded across devices there are two options. With a generic all-to-all, the routing kernels do the reordering on both sides of the communication: tokens are sorted by destination rank before the all-to-all and regrouped by local expert after it (see :ref:`Reordering expert chunks -`). With the NCCL-based :ref:`expert parallelism +`). With the NCCL-based :ref:`expert parallelism ` primitives the permutation is folded into the communication: the dispatch delivers an expert-contiguous receive buffer and the combine writes the results straight back into the original token order, so no From fa068efcd9c91aed956b24da92aaec610ce4ff32 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:06:01 +0200 Subject: [PATCH 24/52] [Docs] MoE: fix EP figure dispatch arrows; animate dispatch and combine Signed-off-by: Pawel Gadzinski --- .../img/moe_expert_parallel.svg | 45 ++++++++++++------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/docs/features/mixture_of_experts/img/moe_expert_parallel.svg b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg index 8952b7ca7a..654a84ee48 100644 --- a/docs/features/mixture_of_experts/img/moe_expert_parallel.svg +++ b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg @@ -14,6 +14,7 @@ .box { fill: #ffffff; stroke: #bdbdbd; stroke-width: 1.3; } .panel { fill: #f5f5f5; stroke: #9e9e9e; stroke-width: 1.2; stroke-dasharray: 5 4; } .arrow { stroke: #424242; stroke-width: 1.6; fill: none; marker-end: url(#ah); } + .chip { stroke-width: 1.6; } @@ -27,58 +28,68 @@ combine all-to-all output - Rank 0 Rank 1 - t0 → E0 + t0 → E0 t1 → E2 t2 → E0 t3 → E1 t4 → E3 t5 → E2 - - + - - + + - - E0: t0 + E0: t0 E0: t2 E1: t3 E2: t1 E2: t5 E3: t4 - - - - + + + Grouped MLPE0, E1 - + + + Grouped MLPE2, E3 - - + - - y0 + y0 y1 y2 y3 y4 y5 + + + + + + + + + + + + + + From 20bc2ac97aca0aaffb5b12d8d0471784879b8190 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:08:49 +0200 Subject: [PATCH 25/52] [Docs] MoE: drop EP figure animation Signed-off-by: Pawel Gadzinski --- .../img/moe_expert_parallel.svg | 23 ++----------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/docs/features/mixture_of_experts/img/moe_expert_parallel.svg b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg index 654a84ee48..b0bcc61680 100644 --- a/docs/features/mixture_of_experts/img/moe_expert_parallel.svg +++ b/docs/features/mixture_of_experts/img/moe_expert_parallel.svg @@ -14,7 +14,6 @@ .box { fill: #ffffff; stroke: #bdbdbd; stroke-width: 1.3; } .panel { fill: #f5f5f5; stroke: #9e9e9e; stroke-width: 1.2; stroke-dasharray: 5 4; } .arrow { stroke: #424242; stroke-width: 1.6; fill: none; marker-end: url(#ah); } - .chip { stroke-width: 1.6; } @@ -56,13 +55,9 @@ E3: t4 - - - + Grouped MLPE0, E1 - - - + Grouped MLPE2, E3 @@ -78,18 +73,4 @@ y3 y4 y5 - - - - - - - - - - - - - - From d9b260e1395a0e510494e33ce8126cecfe3ce2f1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:16:26 +0200 Subject: [PATCH 26/52] [Docs] MoE: restructure expert parallelism (overview, common API, per-framework API) Signed-off-by: Pawel Gadzinski --- .../img/moe_expert_placement.svg | 60 ++++++++ .../mixture_of_experts/mixture_of_experts.rst | 143 +++++++++++------- 2 files changed, 150 insertions(+), 53 deletions(-) create mode 100644 docs/features/mixture_of_experts/img/moe_expert_placement.svg diff --git a/docs/features/mixture_of_experts/img/moe_expert_placement.svg b/docs/features/mixture_of_experts/img/moe_expert_placement.svg new file mode 100644 index 0000000000..74f7128767 --- /dev/null +++ b/docs/features/mixture_of_experts/img/moe_expert_placement.svg @@ -0,0 +1,60 @@ + + + + + + Rank 0 + + tokens + local batch shard + + E0 + expert weights + + E1 + expert weights + + Rank 1 + + tokens + local batch shard + + E2 + expert weights + + E3 + expert weights + + Rank 2 + + tokens + local batch shard + + E4 + expert weights + + E5 + expert weights + + Rank 3 + + tokens + local batch shard + + E6 + expert weights + + E7 + expert weights + 8 experts on 4 ranks: every rank keeps its own tokens and holds 2 of the 8 experts + diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 73c74be9d2..bb3d59d29a 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -410,72 +410,109 @@ Expert parallelism or newer. It is compiled in by default when Transformer Engine is built for these architectures; set ``NVTE_WITH_NCCL_EP=0`` at build time to disable it. -With expert parallelism (EP) the experts are sharded across devices, and each -device owns a slice of them. Two all-to-all collectives wrap the local expert -computation: a **dispatch** all-to-all sends each token to the rank that owns its -expert, the local grouped GEMM runs, and a **combine** all-to-all returns the -results to the source rank. Dispatch takes the router output (expert indices and -weights) directly and delivers a receive buffer grouped by local expert, and -combine writes the results back in the original token order, so no separate -token dispatch or token combine is needed. +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. .. raw:: html - :file: img/moe_expert_parallel.svg + :file: img/moe_expert_placement.svg -*Figure 8. With experts sharded across ranks, a dispatch all-to-all routes each -token to the rank owning its expert and a combine all-to-all returns the -outputs to the source rank.* +*Figure 8. Expert placement: each rank holds its token shard and a subset of the +experts.* -Dispatch and combine are implemented directly on NCCL, using symmetric-memory -windows for zero-copy transfers. Both are differentiable. The common C API -(``nvte_ep_dispatch`` / ``nvte_ep_combine`` and their backward passes, declared -in ``transformer_engine/common/include/transformer_engine/ep.h``) is exposed in -both frameworks; the snippets show how the dispatch, the local experts and the -combine are wired together. +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: a **dispatch** +all-to-all sends each token to the rank that owns its expert, the local grouped +GEMM runs, and a **combine** all-to-all returns the results to the source rank. +Dispatch takes the router output (expert indices and weights) directly and +delivers a receive buffer grouped by local expert, and combine writes the results +back in the original token order, so no separate token dispatch or token combine +is needed. -.. tabs:: +.. raw:: html + :file: img/moe_expert_parallel.svg - .. tab:: PyTorch +*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 implements dispatch and combine directly on NCCL, using +symmetric-memory windows for zero-copy transfers. The implementation is shared by +both frameworks through the common C API in +``transformer_engine/common/include/transformer_engine/ep.h``: + +* ``nvte_ep_initialize`` / ``nvte_ep_shutdown`` set up the EP group on an + existing NCCL communicator, once per process. +* ``nvte_ep_prepare`` seeds the routing for one step from the top-k expert + indices; ``nvte_ep_dispatch`` and ``nvte_ep_combine`` run the two all-to-alls, + with ``nvte_ep_dispatch_bwd`` and ``nvte_ep_combine_bwd`` for the backward + pass. Per-layer state lives in a caller-owned ``handle_mem`` buffer. + +The per-step operations are allocation-free and CUDA-graph capturable when the +receive buffer has a fixed size: ``recv_capacity_per_rank`` bounds the tokens a +rank receives per step, with ``ep_size * max_tokens_per_rank * top_k`` as the +dropless worst case. Without it the buffer is sized from the actual receive count +each step, at the cost of a host synchronization. Dispatch can also quantize the +tokens to MXFP8 before the all-to-all, so the communication moves the +low-precision payload and the local grouped GEMM consumes it directly. Complete +runnable examples live in ``examples/pytorch/ep/`` and ``examples/jax/ep/`` in +the repository. + +PyTorch +~~~~~~~ + +``transformer_engine.pytorch.ep`` exposes the primitives with autograd support: + +* ``ep_bootstrap(ep_group, ...)`` initializes EP once per process on an existing + process group and fixes the group-wide sizes (number of experts, maximum tokens + per rank, hidden size, top-k, receive capacity). +* ``EpBuffer`` holds the per-call state (routing handle and per-expert token + counts). Use one buffer per layer call that is in flight at the same time, for + example one per pipeline microbatch. Its ``dispatch_fwd_quant_recipe`` enables + the MXFP8 quantization in dispatch. +* ``ep_dispatch(buffer, tokens, topk_idx, topk_weights)`` returns the receive + buffer with one fixed slot range per local expert, the routing weights of the + received tokens, and the number of valid tokens per local expert. +* ``ep_combine(buffer, expert_out)`` returns the summed expert outputs in the + original token order. The routing weights are applied by the caller before the + combine. - .. raw:: html +.. raw:: html -
- Requires SM90 (Hopper) or later -
+
+ Requires SM90 (Hopper) or later +
- .. literalinclude:: moe_expert_parallel_pytorch.py - :language: python - :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH - :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH +.. literalinclude:: moe_expert_parallel_pytorch.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH + :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH - .. tab:: JAX +JAX +~~~ - .. raw:: html +JAX offers two levels of API, both experimental: -
- Requires SM90 (Hopper) or later -
+* ``transformer_engine.jax.moe.moe`` runs the whole MoE block (router, dispatch, + expert MLPs, combine) as a single differentiable call. It is executed inside a + ``Mesh``; ``ep_axis`` names the mesh axis the experts are sharded over and the + dispatch and combine become all-to-all collectives over that axis. It also + returns the load-balancing loss when ``aux_loss_coeff`` is non-zero. +* ``transformer_engine.jax.ep`` exposes the primitives separately: + ``ep_bootstrap`` initializes EP once per process from the mesh, ``ep_dispatch`` + scatters tokens and weights to the expert ranks and returns the receive buffer + together with the routing handle and token counts, and ``ep_combine`` sums the + expert outputs back on the source ranks. - .. literalinclude:: moe_expert_parallel_jax.py - :language: python - :start-after: # START_MOE_EXPERT_PARALLEL_JAX - :end-before: # END_MOE_EXPERT_PARALLEL_JAX - -Sizing the receive buffer -~~~~~~~~~~~~~~~~~~~~~~~~~ - -Each rank receives a data-dependent number of tokens per step. Passing a fixed -receive capacity (``recv_capacity_per_rank``) sizes the receive buffer up front, -so the step needs no device-to-host synchronization and can be captured in a -CUDA graph; the dropless worst case is ``ep_size * max_tokens_per_rank * top_k``. -Without it the buffer is sized from the actual receive count each step, at the -cost of a host sync. - -Dispatch can quantize the tokens to MXFP8 before the all-to-all -(``dispatch_fwd_quant_recipe``), so the communication moves the low-precision -payload and the local grouped GEMM consumes it directly. Complete runnable -examples live in ``examples/pytorch/ep/`` and ``examples/jax/ep/`` in the -repository. +.. raw:: html + +
+ Requires SM90 (Hopper) or later +
+ +.. literalinclude:: moe_expert_parallel_jax.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_JAX + :end-before: # END_MOE_EXPERT_PARALLEL_JAX .. _moe-putting-it-together: From 2b66b28b2ab5a517cbdc95a6fced09f155b16261 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:19:49 +0200 Subject: [PATCH 27/52] [Docs] MoE: EP back in tabs without C API; describe sort_chunks_by_index Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 133 +++++++++--------- 1 file changed, 68 insertions(+), 65 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index bb3d59d29a..932dedcc3f 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -283,12 +283,21 @@ that token combine needs to remove the padding. Reordering expert chunks ~~~~~~~~~~~~~~~~~~~~~~~~ -The sort-chunks-by-index kernels permute contiguous chunks of a token tensor -according to a list of chunk sizes and a permutation of chunk indices, for -example to regroup tokens by destination rank before an all-to-all and to -restore the original grouping afterwards. A ``_with_probs`` variant reorders an -accompanying probability tensor in the same call. See the API reference for the -signatures. +``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``, and the chunks are concatenated again in the order +given by ``sorted_indices``: output chunk ``i`` is input chunk +``sorted_indices[i]``. The 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. The buffer a +rank receives 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 the received chunks are ``(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``. After the experts have run, the inverse +permutation restores the rank-major order for the combine all-to-all. .. _moe-grouped-gemm: @@ -436,18 +445,8 @@ experts run on the receive buffer, and combine returns the outputs to the source rank.* Transformer Engine implements dispatch and combine directly on NCCL, using -symmetric-memory windows for zero-copy transfers. The implementation is shared by -both frameworks through the common C API in -``transformer_engine/common/include/transformer_engine/ep.h``: - -* ``nvte_ep_initialize`` / ``nvte_ep_shutdown`` set up the EP group on an - existing NCCL communicator, once per process. -* ``nvte_ep_prepare`` seeds the routing for one step from the top-k expert - indices; ``nvte_ep_dispatch`` and ``nvte_ep_combine`` run the two all-to-alls, - with ``nvte_ep_dispatch_bwd`` and ``nvte_ep_combine_bwd`` for the backward - pass. Per-layer state lives in a caller-owned ``handle_mem`` buffer. - -The per-step operations are allocation-free and CUDA-graph capturable when the +symmetric-memory windows for zero-copy transfers. Both operations are +differentiable. They are allocation-free and CUDA-graph capturable when the receive buffer has a fixed size: ``recv_capacity_per_rank`` bounds the tokens a rank receives per step, with ``ep_size * max_tokens_per_rank * top_k`` as the dropless worst case. Without it the buffer is sized from the actual receive count @@ -457,62 +456,66 @@ low-precision payload and the local grouped GEMM consumes it directly. Complete runnable examples live in ``examples/pytorch/ep/`` and ``examples/jax/ep/`` in the repository. -PyTorch -~~~~~~~ - -``transformer_engine.pytorch.ep`` exposes the primitives with autograd support: - -* ``ep_bootstrap(ep_group, ...)`` initializes EP once per process on an existing - process group and fixes the group-wide sizes (number of experts, maximum tokens - per rank, hidden size, top-k, receive capacity). -* ``EpBuffer`` holds the per-call state (routing handle and per-expert token - counts). Use one buffer per layer call that is in flight at the same time, for - example one per pipeline microbatch. Its ``dispatch_fwd_quant_recipe`` enables - the MXFP8 quantization in dispatch. -* ``ep_dispatch(buffer, tokens, topk_idx, topk_weights)`` returns the receive - buffer with one fixed slot range per local expert, the routing weights of the - received tokens, and the number of valid tokens per local expert. -* ``ep_combine(buffer, expert_out)`` returns the summed expert outputs in the - original token order. The routing weights are applied by the caller before the - combine. +.. tabs:: -.. raw:: html + .. tab:: PyTorch -
- Requires SM90 (Hopper) or later -
+ ``transformer_engine.pytorch.ep`` exposes the primitives with autograd + support: + + * ``ep_bootstrap(ep_group, ...)`` initializes EP once per process on an + existing process group and fixes the group-wide sizes (number of experts, + maximum tokens per rank, hidden size, top-k, receive capacity). + * ``EpBuffer`` holds the per-call state (routing handle and per-expert + token counts). Use one buffer per layer call that is in flight at the + same time, for example one per pipeline microbatch. Its + ``dispatch_fwd_quant_recipe`` enables the MXFP8 quantization in dispatch. + * ``ep_dispatch(buffer, tokens, topk_idx, topk_weights)`` returns the + receive buffer with one fixed slot range per local expert, the routing + weights of the received tokens, and the number of valid tokens per local + expert. + * ``ep_combine(buffer, expert_out)`` returns the summed expert outputs in + the original token order. The routing weights are applied by the caller + before the combine. -.. literalinclude:: moe_expert_parallel_pytorch.py - :language: python - :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH - :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH + .. raw:: html + +
+ Requires SM90 (Hopper) or later +
-JAX -~~~ + .. literalinclude:: moe_expert_parallel_pytorch.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_PYTORCH + :end-before: # END_MOE_EXPERT_PARALLEL_PYTORCH -JAX offers two levels of API, both experimental: + .. tab:: JAX -* ``transformer_engine.jax.moe.moe`` runs the whole MoE block (router, dispatch, - expert MLPs, combine) as a single differentiable call. It is executed inside a - ``Mesh``; ``ep_axis`` names the mesh axis the experts are sharded over and the - dispatch and combine become all-to-all collectives over that axis. It also - returns the load-balancing loss when ``aux_loss_coeff`` is non-zero. -* ``transformer_engine.jax.ep`` exposes the primitives separately: - ``ep_bootstrap`` initializes EP once per process from the mesh, ``ep_dispatch`` - scatters tokens and weights to the expert ranks and returns the receive buffer - together with the routing handle and token counts, and ``ep_combine`` sums the - expert outputs back on the source ranks. + JAX offers two levels of API, both experimental: + + * ``transformer_engine.jax.moe.moe`` runs the whole MoE block (router, + dispatch, expert MLPs, combine) as a single differentiable call. It is + executed inside a ``Mesh``; ``ep_axis`` names the mesh axis the experts + are sharded over and the dispatch and combine become all-to-all + collectives over that axis. It also returns the load-balancing loss when + ``aux_loss_coeff`` is non-zero. + * ``transformer_engine.jax.ep`` exposes the primitives separately: + ``ep_bootstrap`` initializes EP once per process from the mesh, + ``ep_dispatch`` scatters tokens and weights to the expert ranks and + returns the receive buffer together with the routing handle and token + counts, and ``ep_combine`` sums the expert outputs back on the source + ranks. -.. raw:: html + .. raw:: html -
- Requires SM90 (Hopper) or later -
+
+ Requires SM90 (Hopper) or later +
-.. literalinclude:: moe_expert_parallel_jax.py - :language: python - :start-after: # START_MOE_EXPERT_PARALLEL_JAX - :end-before: # END_MOE_EXPERT_PARALLEL_JAX + .. literalinclude:: moe_expert_parallel_jax.py + :language: python + :start-after: # START_MOE_EXPERT_PARALLEL_JAX + :end-before: # END_MOE_EXPERT_PARALLEL_JAX .. _moe-putting-it-together: From 86afa764e6ba9d677bad8901b47d8f311fa39545 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:25:55 +0200 Subject: [PATCH 28/52] [Docs] MoE: EP placement figure matches flow figure; clarify capacity, MXFP8 dispatch, shared experts Signed-off-by: Pawel Gadzinski --- .../img/moe_expert_placement.svg | 76 ++++++++----------- .../mixture_of_experts/mixture_of_experts.rst | 28 ++++--- 2 files changed, 46 insertions(+), 58 deletions(-) diff --git a/docs/features/mixture_of_experts/img/moe_expert_placement.svg b/docs/features/mixture_of_experts/img/moe_expert_placement.svg index 74f7128767..10bcb732f6 100644 --- a/docs/features/mixture_of_experts/img/moe_expert_placement.svg +++ b/docs/features/mixture_of_experts/img/moe_expert_placement.svg @@ -3,58 +3,42 @@ - - Rank 0 - - tokens - local batch shard - - E0 - expert weights - - E1 - expert weights - - Rank 1 - - tokens - local batch shard - - E2 - expert weights - - E3 - expert weights - - Rank 2 - - tokens - local batch shard - - E4 - expert weights - - E5 - expert weights - - Rank 3 - - tokens - local batch shard - - E6 - expert weights - - E7 - expert weights - 8 experts on 4 ranks: every rank keeps its own tokens and holds 2 of the 8 experts + + Rank 0 + tokens + + t0 → E0 + + t1 → E2 + + t2 → E0 + local experts + + E0 + + E1 + + Rank 1 + tokens + + t3 → E1 + + t4 → E3 + + t5 → E2 + local experts + + E2 + + E3 + 4 experts on 2 ranks: each rank holds its own tokens and two of the experts; token colors mark the selected expert diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 932dedcc3f..dc47a56d83 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -426,7 +426,7 @@ keeps its own shard of the tokens and holds only a slice of the experts. :file: img/moe_expert_placement.svg *Figure 8. Expert placement: each rank holds its token shard and a subset of the -experts.* +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: a **dispatch** @@ -435,7 +435,8 @@ GEMM runs, and a **combine** all-to-all returns the results to the source rank. Dispatch takes the router output (expert indices and weights) directly and delivers a receive buffer grouped by local expert, and combine writes the results back in the original token order, so no separate token dispatch or token combine -is needed. +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. .. raw:: html :file: img/moe_expert_parallel.svg @@ -447,14 +448,15 @@ rank.* Transformer Engine implements dispatch and combine directly on NCCL, using symmetric-memory windows for zero-copy transfers. Both operations are differentiable. They are allocation-free and CUDA-graph capturable when the -receive buffer has a fixed size: ``recv_capacity_per_rank`` bounds the tokens a -rank receives per step, with ``ep_size * max_tokens_per_rank * top_k`` as the -dropless worst case. Without it the buffer is sized from the actual receive count -each step, at the cost of a host synchronization. Dispatch can also quantize the -tokens to MXFP8 before the all-to-all, so the communication moves the -low-precision payload and the local grouped GEMM consumes it directly. Complete -runnable examples live in ``examples/pytorch/ep/`` and ``examples/jax/ep/`` in -the repository. +receive buffer has a fixed size: ``recv_capacity_per_rank`` is the maximum +number of tokens (rows of ``hidden_size``) a rank receives per step, with +``ep_size * max_tokens_per_rank * top_k`` as the dropless worst case. Without it +the buffer is sized from the actual receive count each step, at the cost of a +host synchronization. Dispatch can also quantize the tokens before the +all-to-all, so the communication moves the low-precision payload and the local +grouped GEMM consumes it directly; currently only the MXFP8 recipe is supported +there. Complete runnable examples live in ``examples/pytorch/ep/`` and +``examples/jax/ep/`` in the repository. .. tabs:: @@ -468,8 +470,10 @@ the repository. maximum tokens per rank, hidden size, top-k, receive capacity). * ``EpBuffer`` holds the per-call state (routing handle and per-expert token counts). Use one buffer per layer call that is in flight at the - same time, for example one per pipeline microbatch. Its - ``dispatch_fwd_quant_recipe`` enables the MXFP8 quantization in dispatch. + same time, for example one per pipeline microbatch. Passing + ``dispatch_fwd_quant_recipe=MXFP8BlockScaling()`` makes dispatch return + the receive buffer as an MXFP8 grouped tensor (see + ``tests/pytorch/distributed/run_ep.py`` for a complete example). * ``ep_dispatch(buffer, tokens, topk_idx, topk_weights)`` returns the receive buffer with one fixed slot range per local expert, the routing weights of the received tokens, and the number of valid tokens per local From dfc3a19e43d3830ee3dcf7a3a5fde5c48e4008ce Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:29:59 +0200 Subject: [PATCH 29/52] [Docs] MoE: break long paragraphs into lists Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 99 +++++++++++-------- 1 file changed, 58 insertions(+), 41 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index dc47a56d83..585b78dbdf 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -125,11 +125,13 @@ 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. -``expert_bias`` balances the load without an extra loss term. With the sigmoid -score function 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 and raise it for under-used -ones. +``expert_bias`` balances the load without an extra loss term: + +* with the sigmoid score function 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. .. tabs:: @@ -153,11 +155,12 @@ 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. All of these -kernels are differentiable. The snippets below use the mask-map routing variant; -other variants (for example index-map routing) follow the same pattern, see the -:doc:`PyTorch API reference ` and :doc:`JAX API reference -`. +the grouped GEMM, and token combine moves the expert outputs back. + +* All of these kernels are differentiable. +* The snippets below use the mask-map routing variant. Other variants (for + example index-map routing) follow the same pattern, see the :doc:`PyTorch API + reference ` and :doc:`JAX API reference `. Token dispatch ~~~~~~~~~~~~~~ @@ -283,21 +286,27 @@ that token combine needs to remove the padding. 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``, and the chunks are concatenated again in the order -given by ``sorted_indices``: output chunk ``i`` is input chunk -``sorted_indices[i]``. The 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. +``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. The buffer a -rank receives is ordered by source rank and then by expert, while the grouped +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 the received chunks are ``(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``. After the experts have run, the inverse -permutation restores the rank-major order for the combine all-to-all. +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. .. _moe-grouped-gemm: @@ -429,14 +438,18 @@ keeps its own shard of the tokens and holds only a slice of the experts. 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: a **dispatch** -all-to-all sends each token to the rank that owns its expert, the local grouped -GEMM runs, and a **combine** all-to-all returns the results to the source rank. -Dispatch takes the router output (expert indices and weights) directly and -delivers a receive buffer grouped by local expert, and combine writes the results -back in the original token order, so 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. +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. .. raw:: html :file: img/moe_expert_parallel.svg @@ -447,16 +460,20 @@ rank.* Transformer Engine implements dispatch and combine directly on NCCL, using symmetric-memory windows for zero-copy transfers. Both operations are -differentiable. They are allocation-free and CUDA-graph capturable when the -receive buffer has a fixed size: ``recv_capacity_per_rank`` is the maximum -number of tokens (rows of ``hidden_size``) a rank receives per step, with -``ep_size * max_tokens_per_rank * top_k`` as the dropless worst case. Without it -the buffer is sized from the actual receive count each step, at the cost of a -host synchronization. Dispatch can also quantize the tokens before the -all-to-all, so the communication moves the low-precision payload and the local -grouped GEMM consumes it directly; currently only the MXFP8 recipe is supported -there. Complete runnable examples live in ``examples/pytorch/ep/`` and -``examples/jax/ep/`` in the repository. +differentiable. + +* **Receive buffer size.** ``recv_capacity_per_rank`` is the maximum number of + tokens (rows of ``hidden_size``) a rank receives per step; the dropless worst + case is ``ep_size * max_tokens_per_rank * top_k``. With a fixed capacity the + step is allocation-free and CUDA-graph capturable. Without it the buffer is + sized from the actual receive count each step, at the cost of a host + synchronization. +* **Quantized dispatch.** Dispatch can quantize the tokens before the + all-to-all, so the communication moves the low-precision payload and the local + grouped GEMM consumes it directly. Currently only the MXFP8 recipe is + supported there. +* **Examples.** Complete runnable examples live in ``examples/pytorch/ep/`` and + ``examples/jax/ep/`` in the repository. .. tabs:: From a0eec9bf987dcea18d8bf0c99dd75580f8994507 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:32:45 +0200 Subject: [PATCH 30/52] [Docs] MoE: explain NCCL EP and zero-copy mode; link examples Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 585b78dbdf..bb5e7a26de 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -458,10 +458,19 @@ dense MLP on the local tokens on every rank. experts run on the receive buffer, and combine returns the outputs to the source rank.* -Transformer Engine implements dispatch and combine directly on NCCL, using -symmetric-memory windows for zero-copy transfers. Both operations are +Dispatch and combine are built on the NCCL EP library (``libnccl_ep``, loaded +at runtime), not on generic all-to-all collectives. Both operations are differentiable. +* **Communication.** The NCCL EP kernels move each token straight to the slot + of its expert on the owning rank, so the routing and the communication happen + in one step. +* **Zero-copy mode.** Optionally, the token and receive buffers are allocated as + NCCL symmetric memory (``symm_mem_alloc``): the same buffer is registered on + every rank as a window, so the kernels write directly into the peer's buffer + instead of staging the payload in internal NCCL buffers. Without it the + library copies through its own staging buffers. + * **Receive buffer size.** ``recv_capacity_per_rank`` is the maximum number of tokens (rows of ``hidden_size``) a rank receives per step; the dropless worst case is ``ep_size * max_tokens_per_rank * top_k``. With a fixed capacity the @@ -472,8 +481,9 @@ differentiable. all-to-all, so the communication moves the low-precision payload and the local grouped GEMM consumes it directly. Currently only the MXFP8 recipe is supported there. -* **Examples.** Complete runnable examples live in ``examples/pytorch/ep/`` and - ``examples/jax/ep/`` in the repository. +* **Examples.** Complete runnable examples: + `examples/pytorch/ep `_ + and `examples/jax/ep `_. .. tabs:: @@ -490,7 +500,8 @@ differentiable. same time, for example one per pipeline microbatch. Passing ``dispatch_fwd_quant_recipe=MXFP8BlockScaling()`` makes dispatch return the receive buffer as an MXFP8 grouped tensor (see - ``tests/pytorch/distributed/run_ep.py`` for a complete example). + `tests/pytorch/distributed/run_ep.py `_ + for a complete example). * ``ep_dispatch(buffer, tokens, topk_idx, topk_weights)`` returns the receive buffer with one fixed slot range per local expert, the routing weights of the received tokens, and the number of valid tokens per local From 2ea91cb475ca9fda2d80987e67b28caf7f4cb234 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:38:55 +0200 Subject: [PATCH 31/52] [Docs] MoE: EP wording (optimized ops, worst-case capacity, MXFP8 status) Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index bb5e7a26de..5239e2713c 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -458,9 +458,10 @@ dense MLP on the local tokens on every rank. experts run on the receive buffer, and combine returns the outputs to the source rank.* -Dispatch and combine are built on the NCCL EP library (``libnccl_ep``, loaded -at runtime), not on generic all-to-all collectives. Both operations are -differentiable. +Transformer Engine provides optimized implementations of both operations, +including their backward passes, so the layer does not have to assemble them +from generic collectives and permutation kernels. They are built on the NCCL EP +library (``libnccl_ep``, loaded at runtime) and are differentiable. * **Communication.** The NCCL EP kernels move each token straight to the slot of its expert on the owning rank, so the routing and the communication happen @@ -472,15 +473,18 @@ differentiable. library copies through its own staging buffers. * **Receive buffer size.** ``recv_capacity_per_rank`` is the maximum number of - tokens (rows of ``hidden_size``) a rank receives per step; the dropless worst - case is ``ep_size * max_tokens_per_rank * top_k``. With a fixed capacity the - step is allocation-free and CUDA-graph capturable. Without it the buffer is - sized from the actual receive count each step, at the cost of a host - synchronization. + tokens (rows of ``hidden_size``) a rank receives per step. Every rank sends at + most ``max_tokens_per_rank`` tokens to ``top_k`` experts each, and in the + worst case all of them are routed to experts on the same rank, so + ``ep_size * max_tokens_per_rank * top_k`` never drops a token; a smaller + capacity saves memory but can overflow when the routing is skewed (see + ``drop_on_overflow``). With a fixed capacity the step is allocation-free and + CUDA-graph capturable. Without it the buffer is sized from the actual receive + count each step, at the cost of a host synchronization. * **Quantized dispatch.** Dispatch can quantize the tokens before the all-to-all, so the communication moves the low-precision payload and the local - grouped GEMM consumes it directly. Currently only the MXFP8 recipe is - supported there. + grouped GEMM consumes it directly. MXFP8 is supported today; support for + further recipes is in progress. * **Examples.** Complete runnable examples: `examples/pytorch/ep `_ and `examples/jax/ep `_. From 6b144c92e70562528860a62d6da3d2912ae36ee0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 17:42:45 +0200 Subject: [PATCH 32/52] [Docs] MoE: describe EP buffer allocation and data flow (PyTorch) Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 5239e2713c..8209680daf 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -514,6 +514,23 @@ library (``libnccl_ep``, loaded at runtime) and are differentiable. the original token order. The routing weights are applied by the caller before the combine. + Data flow between the calls: + + * ``EpBuffer`` itself allocates only the routing state (a small + ``handle_mem`` byte buffer and the per-expert token counts). + * ``ep_dispatch`` allocates the receive buffer + ``[recv_capacity_per_rank, hidden_size]`` and the received weights on + every call, or writes into caller-owned buffers passed as + ``recv_tokens`` / ``recv_topk_weights`` (needed for CUDA graphs and + zero-copy). The tokens land directly in their expert's slot range. + * The local experts read the receive buffer as their input and produce a + new ``expert_out`` tensor of the same shape; padded slots must be zero. + * ``ep_combine`` reads ``expert_out`` in place and writes the result into a + newly allocated ``[num_tokens, hidden_size]`` tensor. In zero-copy mode + ``expert_out`` is transferred straight from that tensor when it is + symmetric-memory backed; otherwise it goes through the library's + staging buffers. + .. raw:: html
From 4ffa15f32c80b6f4324db50ca302b340b137d660 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 18:05:19 +0200 Subject: [PATCH 33/52] [Docs] MoE: clarify what EpBuffer holds and when to reuse it Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 8209680daf..f544faafae 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -499,13 +499,17 @@ library (``libnccl_ep``, loaded at runtime) and are differentiable. * ``ep_bootstrap(ep_group, ...)`` initializes EP once per process on an existing process group and fixes the group-wide sizes (number of experts, maximum tokens per rank, hidden size, top-k, receive capacity). - * ``EpBuffer`` holds the per-call state (routing handle and per-expert - token counts). Use one buffer per layer call that is in flight at the - same time, for example one per pipeline microbatch. Passing - ``dispatch_fwd_quant_recipe=MXFP8BlockScaling()`` makes dispatch return - the receive buffer as an MXFP8 grouped tensor (see - `tests/pytorch/distributed/run_ep.py `_ - for a complete example). + * ``EpBuffer`` holds the routing state of one dispatch/combine pair: where + each token was sent, how many tokens each local expert received, and the + metadata the combine and both backward passes need to undo the dispatch. + Dispatch writes this state and combine and backward read it, so a buffer + must not be reused until the backward of that call has run. Use one + buffer per MoE layer, and one per microbatch when several microbatches + are in flight (pipeline parallelism). ``dispatch_fwd_quant_recipe`` + selects the quantization applied by dispatch; with + ``MXFP8BlockScaling()`` the receive buffer is returned as an MXFP8 + grouped tensor that the grouped GEMM consumes directly (see + `tests/pytorch/distributed/run_ep.py `_). * ``ep_dispatch(buffer, tokens, topk_idx, topk_weights)`` returns the receive buffer with one fixed slot range per local expert, the routing weights of the received tokens, and the number of valid tokens per local From 3bda8aa426222f9789d3820905771b7e3b45e9c2 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 18:05:58 +0200 Subject: [PATCH 34/52] [Docs] MoE: spell out the JAX EP primitives Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index f544faafae..6af4fa34ad 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -556,12 +556,29 @@ library (``libnccl_ep``, loaded at runtime) and are differentiable. are sharded over and the dispatch and combine become all-to-all collectives over that axis. It also returns the load-balancing loss when ``aux_loss_coeff`` is non-zero. - * ``transformer_engine.jax.ep`` exposes the primitives separately: - ``ep_bootstrap`` initializes EP once per process from the mesh, - ``ep_dispatch`` scatters tokens and weights to the expert ranks and - returns the receive buffer together with the routing handle and token - counts, and ``ep_combine`` sums the expert outputs back on the source - ranks. + * ``transformer_engine.jax.ep`` exposes the primitives separately. Unlike + the PyTorch ``EpBuffer``, the routing state is not kept in an object: + dispatch returns it as arrays and the caller passes them on to combine. + + * ``ep_bootstrap(world_size, rank, num_experts, max_tokens_per_rank, + recv_capacity_per_rank, hidden_dim, ...)`` initializes the EP group + once per process. It runs inside the active ``Mesh`` and reads the EP + axis (and the data-parallel axes) from ``MeshResource``; one process + per device is required. + * ``EpLayerConfig(top_k, ...)`` is a small per-layer configuration that + every per-step call takes as its first argument. + * ``ep_dispatch(cfg, topk_idx, tokens, topk_weights, + recv_capacity_per_rank)`` scatters the tokens to the expert ranks and + returns ``(recv_tokens, recv_topk_weights, handle_mem, token_counts, + total_recv_tokens)``: the receive buffer grouped by local expert, the + weights of the received tokens, the routing handle and per-expert + token counts needed by combine, and the pre-drop receive total that can + be used to detect overflow. + * ``ep_combine(cfg, handle_mem, token_counts, expert_out, + num_local_tokens)`` sums the expert outputs back on the source ranks in + the original token order. It is unweighted: multiply ``expert_out`` by + ``recv_topk_weights`` (and zero the padded slots) before calling it. + ``num_local_tokens`` must be static because it fixes the output shape. .. raw:: html From 11dec2ed88f3068921c9420711992cd0d1799999 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 18:09:49 +0200 Subject: [PATCH 35/52] [Docs] MoE: precise wording for MXFP8 dispatch output Signed-off-by: Pawel Gadzinski --- docs/features/mixture_of_experts/mixture_of_experts.rst | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index 6af4fa34ad..c3a9219a5e 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -506,9 +506,10 @@ library (``libnccl_ep``, loaded at runtime) and are differentiable. must not be reused until the backward of that call has run. Use one buffer per MoE layer, and one per microbatch when several microbatches are in flight (pipeline parallelism). ``dispatch_fwd_quant_recipe`` - selects the quantization applied by dispatch; with - ``MXFP8BlockScaling()`` the receive buffer is returned as an MXFP8 - grouped tensor that the grouped GEMM consumes directly (see + makes dispatch quantize the tokens before sending them; with + ``MXFP8BlockScaling()`` the receive buffer comes back as an MXFP8 + ``GroupedTensor`` (FP8 rows with per-block scales, one group per local + expert, ``alignment=128`` required) instead of BF16 rows (see `tests/pytorch/distributed/run_ep.py `_). * ``ep_dispatch(buffer, tokens, topk_idx, topk_weights)`` returns the receive buffer with one fixed slot range per local expert, the routing From 7edde444f3b5d0312185e6543bc4df5dc14b3052 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 18:22:58 +0200 Subject: [PATCH 36/52] [Docs] MoE: restructure EP overview along one step (buffer, payload, transfer path) Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 77 ++++++++++++------- 1 file changed, 48 insertions(+), 29 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index c3a9219a5e..a7e974e5c3 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -458,36 +458,51 @@ dense MLP on the local tokens on every rank. experts run on the receive buffer, and combine returns the outputs to the source rank.* -Transformer Engine provides optimized implementations of both operations, -including their backward passes, so the layer does not have to assemble them -from generic collectives and permutation kernels. They are built on the NCCL EP -library (``libnccl_ep``, loaded at runtime) and are differentiable. - -* **Communication.** The NCCL EP kernels move each token straight to the slot - of its expert on the owning rank, so the routing and the communication happen - in one step. -* **Zero-copy mode.** Optionally, the token and receive buffers are allocated as - NCCL symmetric memory (``symm_mem_alloc``): the same buffer is registered on - every rank as a window, so the kernels write directly into the peer's buffer - instead of staging the payload in internal NCCL buffers. Without it the - library copies through its own staging buffers. - -* **Receive buffer size.** ``recv_capacity_per_rank`` is the maximum number of - tokens (rows of ``hidden_size``) a rank receives per step. Every rank sends at - most ``max_tokens_per_rank`` tokens to ``top_k`` experts each, and in the - worst case all of them are routed to experts on the same rank, so - ``ep_size * max_tokens_per_rank * top_k`` never drops a token; a smaller - capacity saves memory but can overflow when the routing is skewed (see - ``drop_on_overflow``). With a fixed capacity the step is allocation-free and - CUDA-graph capturable. Without it the buffer is sized from the actual receive +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. + +**One step.** Dispatch reads the top-k expert indices of the local tokens and +moves the tokens in a single pass: + +* it counts how many tokens every rank and every local expert will receive; +* it writes each token straight into the slot range of its expert in the + receive buffer on the owning rank, so the receive buffer is already grouped by + local expert; +* combine reverses this: it returns each expert output to the source rank and + sums the contributions into the original token order. + +**Receive buffer.** Because every expert owns a fixed slot range, the receive +buffer has a fixed layout and has to be sized up front: + +* ``recv_capacity_per_rank`` is the maximum number of tokens (rows of + ``hidden_size``) a rank receives per step. Every rank sends at most + ``max_tokens_per_rank`` tokens to ``top_k`` experts each, and in the worst + case all of them go to one rank, so ``ep_size * max_tokens_per_rank * top_k`` + never drops a token. A smaller capacity saves memory but can overflow when the + routing is skewed (see ``drop_on_overflow``). +* With a fixed capacity the step allocates nothing and needs no host + synchronization, so it can be captured in a CUDA graph. +* Without a capacity (eager mode) the buffer is sized from the actual receive count each step, at the cost of a host synchronization. -* **Quantized dispatch.** Dispatch can quantize the tokens before the - all-to-all, so the communication moves the low-precision payload and the local - grouped GEMM consumes it directly. MXFP8 is supported today; support for - further recipes is in progress. -* **Examples.** Complete runnable examples: - `examples/pytorch/ep `_ - and `examples/jax/ep `_. + +**Payload.** By default the tokens travel as BF16 rows. Dispatch can quantize +them first: + +* the communication then moves the low-precision payload, and the receive + buffer comes back as a quantized ``GroupedTensor`` (one group per local + expert) that the fused grouped MLP accepts without quantizing it again; +* MXFP8 is supported today; support for further recipes is in progress. + +**Transfer path.** By default the library copies the payload through its own +staging buffers. Zero-copy mode removes these copies: + +* the token and receive buffers are allocated as NCCL symmetric memory + (``symm_mem_alloc``), so the same buffer is registered on every rank as a + window and the kernels write directly into the peer's buffer; +* the buffers have to be persistent, which also makes them the buffers to pass + in when capturing a CUDA graph. .. tabs:: @@ -592,6 +607,10 @@ library (``libnccl_ep``, loaded at runtime) and are differentiable. :start-after: # START_MOE_EXPERT_PARALLEL_JAX :end-before: # END_MOE_EXPERT_PARALLEL_JAX +Complete runnable examples: +`examples/pytorch/ep `_ +and `examples/jax/ep `_. + .. _moe-putting-it-together: Example: putting it all together From 490e922e6704d24eb18494b6cd6ef63125fe8c77 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 18:27:28 +0200 Subject: [PATCH 37/52] [Docs] MoE: EP overview headings and terser bullets Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/mixture_of_experts.rst | 72 +++++++++---------- 1 file changed, 34 insertions(+), 38 deletions(-) diff --git a/docs/features/mixture_of_experts/mixture_of_experts.rst b/docs/features/mixture_of_experts/mixture_of_experts.rst index a7e974e5c3..830d9a64d7 100644 --- a/docs/features/mixture_of_experts/mixture_of_experts.rst +++ b/docs/features/mixture_of_experts/mixture_of_experts.rst @@ -463,46 +463,42 @@ 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. -**One step.** Dispatch reads the top-k expert indices of the local tokens and -moves the tokens in a single pass: - -* it counts how many tokens every rank and every local expert will receive; -* it writes each token straight into the slot range of its expert in the - receive buffer on the owning rank, so the receive buffer is already grouped by - local expert; -* combine reverses this: it returns each expert output to the source rank and - sums the contributions into the original token order. - -**Receive buffer.** Because every expert owns a fixed slot range, the receive -buffer has a fixed layout and has to be sized up front: - -* ``recv_capacity_per_rank`` is the maximum number of tokens (rows of - ``hidden_size``) a rank receives per step. Every rank sends at most - ``max_tokens_per_rank`` tokens to ``top_k`` experts each, and in the worst - case all of them go to one rank, so ``ep_size * max_tokens_per_rank * top_k`` - never drops a token. A smaller capacity saves memory but can overflow when the - routing is skewed (see ``drop_on_overflow``). -* With a fixed capacity the step allocates nothing and needs no host - synchronization, so it can be captured in a CUDA graph. -* Without a capacity (eager mode) the buffer is sized from the actual receive - count each step, at the cost of a host synchronization. - -**Payload.** By default the tokens travel as BF16 rows. Dispatch can quantize -them first: - -* the communication then moves the low-precision payload, and the receive - buffer comes back as a quantized ``GroupedTensor`` (one group per local - expert) that the fused grouped MLP accepts without quantizing it again; -* MXFP8 is supported today; support for further recipes is in progress. - -**Transfer path.** By default the library copies the payload through its own -staging buffers. Zero-copy mode removes these copies: +**Communication and permutation in one step** + +Dispatch reads the top-k expert indices and moves each token straight into the +slot range of its expert on the owning rank: + +* the receive buffer is already grouped by local expert; +* combine reverses it and sums the contributions into the original token order. + +**Receive buffer** + +Every expert owns a fixed slot range, so the buffer is sized up front: + +* ``recv_capacity_per_rank`` is the maximum number of tokens a rank receives per + step; ``ep_size * max_tokens_per_rank * top_k`` never drops a token, a smaller + value can overflow on skewed routing (see ``drop_on_overflow``); +* with a fixed capacity the step allocates nothing, needs no host + synchronization and is CUDA-graph capturable; +* without it (eager mode) the buffer is sized per step, with a host + synchronization. + +**Low precision** + +Dispatch can quantize the tokens before sending them: + +* the receive buffer comes back as a quantized ``GroupedTensor`` (one group per + local expert) that the fused grouped MLP accepts as is; +* MXFP8 is supported today; further recipes are in progress. + +**Zero-copy mode** + +By default the payload is copied through the library's staging buffers. +Optionally: * the token and receive buffers are allocated as NCCL symmetric memory - (``symm_mem_alloc``), so the same buffer is registered on every rank as a - window and the kernels write directly into the peer's buffer; -* the buffers have to be persistent, which also makes them the buffers to pass - in when capturing a CUDA graph. + (``symm_mem_alloc``) and the kernels write directly into the peer's buffer; +* the buffers must be persistent, which also makes them CUDA-graph friendly. .. tabs:: From df2781998f81144747457d6a7993d8552b7ce594 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 18:31:42 +0200 Subject: [PATCH 38/52] [Docs] MoE figures: larger fonts, fix label anchors and clipping, drop placement footer Signed-off-by: Pawel Gadzinski --- .../mixture_of_experts/img/grouped_linear.svg | 19 +++++++++---------- .../img/moe_expert_parallel.svg | 10 +++++----- .../img/moe_expert_placement.svg | 11 +++++------ .../img/moe_grouped_mlp.svg | 4 ++-- .../mixture_of_experts/img/moe_layer_ep.svg | 4 ++-- .../mixture_of_experts/img/moe_padding.svg | 14 +++++++------- .../mixture_of_experts/img/moe_permute.svg | 16 ++++++++-------- .../mixture_of_experts/img/moe_router.svg | 12 ++++++------ .../mixture_of_experts/img/moe_unpermute.svg | 10 +++++----- 9 files changed, 49 insertions(+), 51 deletions(-) diff --git a/docs/features/mixture_of_experts/img/grouped_linear.svg b/docs/features/mixture_of_experts/img/grouped_linear.svg index e438d10ba3..a5aa5dfb79 100644 --- a/docs/features/mixture_of_experts/img/grouped_linear.svg +++ b/docs/features/mixture_of_experts/img/grouped_linear.svg @@ -1,12 +1,12 @@ - +