Add a PyTorch EP=1 backend for MoE dispatch and combine

.github/workflows/docs.yml

@@ -7,6 +7,13 @@ name: 'Documentation'
 on:
   pull_request:
   workflow_dispatch:
+    inputs:
+      enable_public_features:
+        # Disabling this blocks the cookie banner from appearing on localhost and other domains
+        description: 'Enable public features (required for public publishing)'
+        default: false
+        required: true
+        type: boolean
   workflow_call:
 concurrency:
   # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes)
@@ -27,7 +34,7 @@ jobs:
           export GIT_SHA=$(git show-ref --hash HEAD)
       - name: 'Build docs'
         env:
-          NVTE_DOCS_DISABLE_PUBLIC_FEATURES: '1'
+          NVTE_DOCS_DISABLE_PUBLIC_FEATURES: ${{ inputs.enable_public_features && '0' || '1' }}
           NVTE_DOCS_SWITCHER_JSON_URL: 'versions1.json'
         run: | # SPHINXOPTS="-W" errors out on warnings
           doxygen docs/Doxyfile

.github/workflows/lint.yml

@@ -33,11 +33,15 @@ jobs:
     steps:
       - name: 'Checkout'
         uses: actions/checkout@v3
+      - name: 'Set up Python'
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
       - name: 'Lint'
         run: |
           sudo apt-get update
           sudo apt-get install pip -y
-          pip install torch numpy
+          pip install torch numpy wheel
           export PYTHON_ONLY=1
           export TE_PATH=.
           bash ./qa/L0_pytorch_lint/test.sh

.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`.

README.rst

@@ -38,7 +38,7 @@ News
 ====
 
 * **[09/2026]** `Transformer Engine v2.19 <https://github.com/NVIDIA/TransformerEngine/releases/tag/v2.19>`_ adds Rubin support, hybrid quantization, MXFP8 EP communication, and expanded FP8 attention support.
-* **[09/2026]** `Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine <https://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transform-engine/>`_ describes optimized JAX MoE training on GB200 and GB300 systems.
+* **[09/2026]** `Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine <https://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine/>`_ describes optimized JAX MoE training on GB200 and GB300 systems.
 * **[08/2026]** `Transformer Engine v2.18 <https://github.com/NVIDIA/TransformerEngine/releases/tag/v2.18>`_ adds FP8 block scaling in PyTorch, zero-copy EP, and CUDA Graph support for THD attention.
 * **[07/2026]** `Transformer Engine v2.17 <https://github.com/NVIDIA/TransformerEngine/releases/tag/v2.17>`_ introduces NCCL-backed EP for PyTorch and JAX, faster MoE routing, and a JAX Flax MoE block.
 * **[06/2026]** `Boosting MoE Training Throughput with Advanced Fusion Kernels <https://developer.nvidia.com/blog/boosting-moe-training-throughput-with-advanced-fusion-kernels/>`_ shows how fused Transformer Engine operations accelerate MoE training.

benchmarks/linear/benchmark_graph_safe_grouped_mlp.py

@@ -35,7 +35,6 @@
 
 os.environ.setdefault("CUDA_DEVICE_MAX_CONNECTIONS", "1")
 os.environ.setdefault("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")
-os.environ.setdefault("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "1")
 os.environ.setdefault("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1")
 
 import argparse
@@ -91,6 +90,7 @@ def build_grouped_mlp(
             dtype=dtype,
             single_grouped_weight=single_grouped_weight,
             accumulate_into_main_grad=accumulate_into_main_grad,
+            split_alignment=256,
         )
         fc2 = te_ops.GroupedLinear(
             num_groups,
@@ -101,6 +101,7 @@ def build_grouped_mlp(
             dtype=dtype,
             single_grouped_weight=single_grouped_weight,
             accumulate_into_main_grad=accumulate_into_main_grad,
+            split_alignment=256,
         )
         return te_ops.Sequential(
             fc1,
@@ -318,7 +319,6 @@ def main() -> None:
     for name in (
         "CUDA_DEVICE_MAX_CONNECTIONS",
         "NVTE_ALLOW_NONDETERMINISTIC_ALGO",
-        "NVTE_CUTEDSL_FUSED_GROUPED_MLP",
         "CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL",
     ):
         print(f"  {name}={os.environ.get(name)}")

build_tools/VERSION.txt

@@ -1 +1 @@
-2.21.0.dev0
+2.22.0.dev0

build_tools/jax.py

@@ -16,6 +16,7 @@
     all_files_in_dir,
     cudnn_frontend_include_path,
     debug_build_enabled,
+    get_bolt_build_flags,
     setup_mpi_flags,
     nccl_include_path,
     nccl_lib_path,
@@ -26,7 +27,7 @@
 
 def install_requirements() -> List[str]:
     """Install dependencies for TE/JAX extensions."""
-    return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"]
+    return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0", "triton"]
 
 
 def test_requirements() -> List[str]:
@@ -122,6 +123,9 @@ def setup_jax_extension(
     else:
         cxx_flags.append("-g0")
 
+    bolt_cxx_flags, linker_flags = get_bolt_build_flags()
+    cxx_flags.extend(bolt_cxx_flags)
+
     setup_mpi_flags(include_dirs, cxx_flags)
 
     if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))):
@@ -144,5 +148,6 @@ def setup_jax_extension(
         sources=[str(path) for path in sources],
         include_dirs=[str(path) for path in include_dirs],
         extra_compile_args=cxx_flags,
+        extra_link_args=linker_flags,
         **kwargs,
     )

build_tools/pytorch.py

@@ -15,6 +15,7 @@
     cuda_version,
     get_cuda_include_dirs,
     debug_build_enabled,
+    get_bolt_build_flags,
     nccl_ep_enabled,
     setup_mpi_flags,
 )
@@ -76,6 +77,9 @@ def setup_pytorch_extension(
     else:
         cxx_flags.append("-g0")
 
+    bolt_cxx_flags, linker_flags = get_bolt_build_flags()
+    cxx_flags.extend(bolt_cxx_flags)
+
     # Version-dependent CUDA options
     try:
         version = cuda_version()
@@ -127,6 +131,7 @@ def setup_pytorch_extension(
         sources=[str(src) for src in sources],
         include_dirs=[str(inc) for inc in include_dirs],
         extra_compile_args={"cxx": cxx_flags},
+        extra_link_args=linker_flags,
         libraries=[str(lib) for lib in libraries],
         library_dirs=[str(lib_dir) for lib_dir in library_dirs],
     )

build_tools/utils.py

@@ -9,6 +9,7 @@
 import importlib
 import os
 import re
+import shlex
 import shutil
 import subprocess
 import sys
@@ -43,6 +44,72 @@ def debug_build_enabled() -> bool:
     return bool(int(os.getenv("NVTE_BUILD_DEBUG", "0")))
 
 
+@functools.lru_cache(maxsize=None)
+def _cxx_compiler_predefined_macros() -> Tuple[str, ...]:
+    """Predefined macros from the configured C++ compiler."""
+    cxx = shlex.split(os.getenv("CXX", "c++"))
+    if not cxx:
+        return ()
+    try:
+        result = subprocess.run(
+            [*cxx, "-dM", "-E", "-x", "c++", "-"],
+            input="",
+            capture_output=True,
+            text=True,
+        )
+    except OSError:
+        return ()
+    if result.returncode != 0:
+        return ()
+    return tuple(
+        fields[1]
+        for line in result.stdout.splitlines()
+        if len(fields := line.split(maxsplit=2)) >= 2 and fields[0] == "#define"
+    )
+
+
+@functools.lru_cache(maxsize=None)
+def cxx_compiler_supports_flag(flag: str) -> bool:
+    """Whether the configured C++ compiler accepts a command-line flag."""
+    cxx = shlex.split(os.getenv("CXX", "c++"))
+    if not cxx:
+        return False
+    try:
+        result = subprocess.run(
+            [*cxx, flag, "-x", "c++", "-fsyntax-only", "-"],
+            input="",
+            capture_output=True,
+            text=True,
+        )
+    except OSError:
+        return False
+    return result.returncode == 0
+
+
+def cxx_compiler_targets_arm64() -> bool:
+    """Whether the configured C++ compiler targets Arm64."""
+    macros = _cxx_compiler_predefined_macros()
+    return any(macro in macros for macro in ("__aarch64__", "__arm64__", "_M_ARM64"))
+
+
+def get_bolt_build_flags() -> Tuple[List[str], List[str]]:
+    """BOLT-compatible host compiler and linker flags."""
+    compiler_flags = ["-fno-jump-tables"]
+    if cxx_compiler_supports_flag("-fno-reorder-blocks-and-partition"):
+        compiler_flags.append("-fno-reorder-blocks-and-partition")
+    linker_flags = ["-Wl,--emit-relocs", "-Wl,-z,now"]
+    if cxx_compiler_targets_arm64():
+        compiler_flags.extend(
+            [
+                "-mno-fix-cortex-a53-835769",
+                "-mno-fix-cortex-a53-843419",
+            ]
+        )
+        # The Cortex-A53 843419 workaround is applied by the linker.
+        linker_flags.append("-mno-fix-cortex-a53-843419")
+    return compiler_flags, linker_flags
+
+
 @functools.lru_cache(maxsize=None)
 def get_max_jobs_for_parallel_build() -> int:
     """Number of parallel jobs for Nina build"""

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