[PyTorch] [torch.compile] torch.compile support for LayerNormLinear and LayerNormMLP

.github/workflows/minimum_versions.yml

@@ -6,7 +6,12 @@
 name: 'Minimum supported versions'
 on:
   pull_request:
+  push:
+    branches:
+      - main
   workflow_dispatch:
+permissions:
+  contents: read
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true

AGENTS.md

@@ -374,6 +374,33 @@ requirements, or performance-relevant behavior.
 - Standalone examples live under `examples/pytorch/` and `examples/jax/`.
 - Tutorials and notebooks live under `docs/examples/`.
 
+## Code Quality
+
+It is not sufficient for code to be correct and performant: it must be
+human-readable as well. Human review is the bottleneck in TE
+development, and PRs that are unpleasant to read are more likely to be
+neglected by reviewers. These guidelines are defaults, not substitutes
+for judgment.
+
+- Public APIs, including the C API, must be backward compatible. Design
+  new ones to be clean, general, and extensible. Within reason, generalize
+  existing APIs rather than adding new ones.
+- Limit changes to the task's scope; do not refactor unrelated code.
+- Within that scope, restructure abstractions when needed to keep the code
+  coherent rather than adding workarounds that minimize the diff. Avoid
+  overengineering and unnecessary indirection.
+- Write code that is understandable without its comments, and add
+  comments where they aid understanding. Keep them concise, since
+  humans skim long comments. Describe the code as it is for a reader
+  without context, not its edit history. Treat an incorrect comment
+  like incorrect code.
+- Use descriptive names that read naturally. Follow the conventions of the
+  surrounding code and domain, and avoid inventing idiosyncratic
+  terminology.
+- Extend existing infrastructure rather than building parallel systems.
+  For example, add cases to an existing test suite instead of creating
+  a new test file for each change.
+
 ## Contribution Requirements
 
 Follow `CONTRIBUTING.rst` and `.github/PULL_REQUEST_TEMPLATE.md`.

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,27 @@ 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.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

@@ -99,16 +99,55 @@ 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>`.
+
+``MoeDispatch`` and ``MoeCombine`` can wrap the local expert operations in a
+``transformer_engine.pytorch.ops.Sequential``. They share an ``EpConfig`` and
+an ``EpBuffer``; the guide shows how to connect the per-expert counts and
+routing weights through the sequence's extra-tensor channels.
+
+.. autoapiclass:: transformer_engine.pytorch.ops.MoeDispatch
+
+.. autoapiclass:: transformer_engine.pytorch.ops.MoeCombine
+
+.. autoapiclass:: transformer_engine.pytorch.ep.EpConfig
+
+.. 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/envvars.rst

@@ -351,12 +351,6 @@ Kernel Configuration
    :Default: ``0``
    :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. NVFP4 4over6 quantization, fused softmax, and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_NVFP4``, ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX``, or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise no static fallback is available.
 
-.. envvar:: NVTE_MXFP8_PER_TILE_SWIZZLE
-
-   :Type: ``int`` (0 or 1)
-   :Default: ``0``
-   :Description: Swizzle MXFP8 scaling factors for GEMM with kernels that give each 128-row tile of scaling factors (row-wise) or each 128-column by 32-row block (column-wise) its own thread block, instead of batching 32 tiles in one block. This lowers the latency of swizzling the scaling factors of tensors with few tiles, such as activations.
-
 .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM
 
    :Type: ``int`` (0 or 1)

docs/features/mixture_of_experts/grouped_linear_jax.py

@@ -0,0 +1,31 @@
+# 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] int32 array; group_sizes[i] is the number of routed
+#              tokens for expert i
+group_sizes = jnp.asarray(group_sizes, dtype=jnp.int32)
+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.
+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,43 @@
+# 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)
+
+# Keep token counts on the device with use_grouped_tensor=True.
+grouped_linear = te.GroupedLinear(
+    num_experts,
+    hidden_size,
+    ffn_hidden_size,
+    bias=True,
+    params_dtype=torch.bfloat16,
+    use_grouped_tensor=True,
+).cuda()
+m_splits_dev = torch.tensor(m_splits, dtype=torch.int64, device="cuda")
+grouped_out = grouped_linear(x, m_splits_dev)
+# END_GROUPED_LINEAR_PYTORCH

docs/features/mixture_of_experts/grouped_mlp_pytorch.py

@@ -0,0 +1,25 @@
+# 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
+
+# FC1 produces gate and value features interleaved in blocks of 32.
+expert_mlp = te.ops.Sequential(
+    te.ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size, split_alignment=256),
+    te.ops.ScaledSwiGLU(glu_interleave_size=32),
+    te.ops.GroupedLinear(
+        num_experts, ffn_hidden_size, hidden_size, scale_bias=True, split_alignment=256
+    ),
+)
+
+# Inputs from padded dispatch with align_size=256.
+expert_out = expert_mlp(
+    padded,
+    padded_tokens_per_expert,
+    permuted_probs,
+    padded_tokens_per_expert,
+    permuted_probs,
+)
+# END_GROUPED_MLP_PYTORCH