[Docs] Add Mixture of Experts guide

docs/_static/css/diagram-colors.css

@@ -279,6 +279,38 @@ 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"] .moe-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; }
 
 /* fine_grained_quantization diagrams */
 html[data-theme="dark"] .fmt-mxfp8 { fill: #10375c; stroke: #64b5f6; }

docs/api/jax.rst

@@ -59,3 +59,31 @@ Modules
   :members: __call__
 
 .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules
+
+
+Mixture of Experts
+------------------
+Router, routing kernels, and grouped dense for MoE layers. See
+:doc:`Mixture of Experts <../features/mixture_of_experts/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
+
+.. 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

docs/api/pytorch.rst

@@ -90,16 +90,44 @@ Recipe availability
 Mixture of Experts (MoE) functions
 ----------------------------------
 
+See :doc:`Mixture of Experts <../features/mixture_of_experts/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
+
+Expert parallelism
+~~~~~~~~~~~~~~~~~~
+
+NCCL-based dispatch and combine for experts sharded across ranks. See
+:doc:`Expert parallelism <../features/mixture_of_experts/mixture_of_experts>`.
+
+.. 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
 ---------------------------------

docs/features/mixture_of_experts/grouped_linear_jax.py

@@ -0,0 +1,33 @@
+# 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. 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,
+    group_sizes=group_sizes,
+    bias=bias,
+)
+# END_GROUPED_LINEAR_JAX

docs/features/mixture_of_experts/grouped_linear_pytorch.py

@@ -0,0 +1,47 @@
+# 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. By default one GEMM per expert
+# is launched; m_splits is read on the host.
+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)
+
+# 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

docs/features/mixture_of_experts/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