Merge/dev to main 20260807

.github/actions/build-pytorch-wheel/Dockerfile

@@ -41,9 +41,5 @@ RUN CUDA_MAJOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'}) && \
 # Install PyTorch
 RUN export MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'}) && \
     export MATRIX_TORCH_VERSION=$(echo $TORCH_VERSION | awk -F \. {'print $1 "." $2'}) && \
-    export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \ 
-    minv = {'2.5': 118, '2.6': 118, '2.7': 118, '2.8': 126, '2.9': 126}[env['MATRIX_TORCH_VERSION']]; \
-    maxv = {'2.5': 124, '2.6': 126, '2.7': 128, '2.8': 129, '2.9': 130}[env['MATRIX_TORCH_VERSION']]; \
-    print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)" \
-    ) && \
-    pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
+    export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; versions = {'2.5': (118, 124), '2.6': (118, 126), '2.7': (118, 128), '2.8': (126, 129), '2.9': (126, 130)}; minv, maxv = versions[env['MATRIX_TORCH_VERSION']]; print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)") && \
+    pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}

.github/workflows/build.yml

@@ -1,66 +1,160 @@
-# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
 
 # A workflow to trigger TE build on GitHub
-
 name: 'Build'
 on:
   pull_request:
   workflow_dispatch:
+concurrency:
+  # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes)
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
 jobs:
+  core:
+    name: 'Core'
+    runs-on: ubuntu-latest
+    container:
+      image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04
+      options: --user root
+    steps:
+      - name: 'Dependencies'
+        run: |
+          apt-get update
+          apt-get install -y git python3.9 pip cudnn9-cuda-12
+          pip install cmake==3.21.0 pybind11[global] ninja
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+      - name: ccache
+        uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad
+      - name: 'Build'
+        run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v
+        env:
+          NVTE_FRAMEWORK: none
+          MAX_JOBS: 1
+          SCCACHE_GHA_ENABLED: "true"
+      - name: 'Sanity check'
+        run: python3 -c "import transformer_engine"
+        working-directory: /
   pytorch:
     name: 'PyTorch'
-    runs-on: [ nv-8g-cicd-te ]
-    defaults:
-      run:
-        shell: bash
-    container:
-      image: harbor.baai.ac.cn/flagscale/cuda12.8.1-torch2.7.1-python3.10-te2.9:20260209
-      ports:
-        - 80:80
-      options: >-
-        --gpus all 
-        --shm-size=500g 
-        --privileged 
-        --ipc=host 
-        --ulimit memlock=-1 
-        --ulimit stack=67108864 
-        --ulimit nofile=65535:65535 
-        --user root
-        --pull never
+    runs-on: ubuntu-latest
     steps:
-      - name: Configure Git Safe Directory on Cuda
-        run: /usr/bin/git config --global safe.directory '*'
+      - name: Move /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker"
+
+      - name: Maximize build space
+        uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794
+        with:
+          root-reserve-mb: 5120
+          temp-reserve-mb: 32
+          swap-size-mb: 10240
+          remove-dotnet: 'true'
+          remove-android: 'true'
+          remove-haskell: 'true'
+          remove-codeql: 'true'
+          build-mount-path: '/var/lib/docker/'
+
+      - name: Restore /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"
 
       - name: 'Checkout'
-        uses: actions/checkout@v4
+        uses: actions/checkout@v3
         with:
-          fetch-depth: 0
           submodules: recursive
-          set-safe-directory: true
 
-      - name: 'Setup Environment'
+      - name: Start named container
         run: |
-          source /opt/miniconda3/etc/profile.d/conda.sh
-          conda activate flagscale-train
-          echo "PATH=$PATH" >> $GITHUB_ENV
+          docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 sleep infinity
 
-      - name: 'Build'
+      - name: 'Dependencies'
         run: |
-          pip uninstall transformer_engine transformer_engine_torch -y || true
-          echo "GITHUB_WORKSPACE=$GITHUB_WORKSPACE"
-          cd $GITHUB_WORKSPACE
-          pip install nvdlfw-inspect
-          pip install expecttest
-          pip install . -v --no-deps --no-build-isolation
+          docker exec builder bash -c '\
+            apt-get update && \
+            apt-get install -y git python3.9 pip cudnn9-cuda-12 && \
+            pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript && \
+            apt-get clean \
+          '
+
+      - name: 'Build'
+        run: docker exec builder bash -c 'pip install --no-build-isolation . -v --no-deps'
         env:
           NVTE_FRAMEWORK: pytorch
-          TE_WITH_NCCL: '1'
-          NVTE_WITH_CUDA: '1'
-          CUDA_HOME: /usr/local/cuda-12.8
-          NVCC: /usr/local/cuda-12.8/bin/nvcc
+          MAX_JOBS: 1
+      - name: 'Sanity check'
+        run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py'
+  jax:
+    name: 'JAX'
+    runs-on: ubuntu-latest
+    container:
+      image: ghcr.io/nvidia/jax:jax
+      options: --user root
+    steps:
+      - name: 'Dependencies'
+        run: pip install cmake==3.21.0 pybind11[global]
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+      - name: ccache
+        uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad
+      - name: 'Build'
+        run: |
+          NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v
+        env:
+          NVTE_FRAMEWORK: jax
+          MAX_JOBS: 1
+          SCCACHE_GHA_ENABLED: "true"
+      - name: 'Sanity check'
+        run: python3 tests/jax/test_sanity_import.py
+  all:
+    name: 'All'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Move /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker"
+
+      - name: Maximize build space
+        uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794
+        with:
+          root-reserve-mb: 5120
+          temp-reserve-mb: 32
+          swap-size-mb: 10240
+          remove-dotnet: 'true'
+          remove-android: 'true'
+          remove-haskell: 'true'
+          remove-codeql: 'true'
+          build-mount-path: '/var/lib/docker/'
 
+      - name: Restore /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"
+
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+
+      - name: Start named container
+        run: |
+          docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity
+
+      - name: 'Dependencies'
+        run: |
+          docker exec builder bash -c '\
+            pip install cmake==3.21.0 pybind11[global] einops onnxscript && \
+            pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130
+          '
+      - name: 'Build'
+        run: docker exec builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps'
+        env:
+          NVTE_FRAMEWORK: all
+          MAX_JOBS: 1
       - name: 'Sanity check'
-        run: 
-          python3 tests/pytorch/test_sanity_import.py
+        run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py'

.github/workflows/community_label.yml

@@ -0,0 +1,63 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+# A workflow to automatically label the contributions as community/org
+name: Label community contributions
+
+on:
+  pull_request_target:
+    types: [opened, reopened, ready_for_review, synchronize]
+
+permissions:
+  contents: read
+  issues: write
+  pull-requests: write
+
+jobs:
+  label:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3
+        with:
+          script: |
+            const pr = context.payload.pull_request;
+            const user = pr.user.login;
+            const association = pr.author_association;
+
+            const communityLabel = "community-contribution";
+            const orgLabel = "org-contribution";
+
+            let targetLabel = null;
+
+            const isOrgMember =
+              association === "MEMBER" || association === "OWNER";
+
+            let permission = "none";
+
+            try {
+              const res = await github.rest.repos.getCollaboratorPermissionLevel({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                username: user,
+              });
+              permission = res.data.permission;
+            } catch (e) {
+              if (e.status !== 404) throw e;
+            }
+
+            const isCore = permission === "write" || permission === "admin";
+            if (!isOrgMember) {
+              targetLabel = communityLabel;
+            } else {
+              targetLabel = orgLabel;
+            }
+
+            if (!isCore) {
+              await github.rest.issues.addLabels({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                issue_number: pr.number,
+                labels: [targetLabel],
+              });
+            }

.github/workflows/lint.yml

@@ -11,6 +11,8 @@ concurrency:
   # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes)
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true
+permissions:
+  contents: read
 jobs:
   pytorch_cpplint:
     name: 'PyTorch C++'

.gitignore

@@ -48,6 +48,7 @@ artifacts/
 transformer_engine/plugin/core/_build_config.py
 # Mac OS
 .DS_Store
+.claude/
 # Integration test outputs
 qa/L1_pytorch_mcore_integration/output/
 *.distcp

.gitmodules

@@ -7,3 +7,6 @@
 [submodule "3rdparty/cutlass"]
 	path = 3rdparty/cutlass
 	url = https://github.com/NVIDIA/cutlass.git
+[submodule "3rdparty/nccl"]
+	path = 3rdparty/nccl
+	url = https://github.com/NVIDIA/nccl.git

.pre-commit-config.yaml

@@ -40,7 +40,7 @@ repos:
         files: ^transformer_engine.*\.(c|cc|cxx|cpp|cu|cuh|h|hpp)$
 
   - repo: https://github.com/netromdk/vermin
-    rev: c75aca72f4e85c6e47252139e8695f1c8b5f9ae3
+    rev: b70ff9611a01a2bf2f702aa537d14e71e330edba
     hooks:
       - id: vermin
         args: ['-t=3.10-', '--violations']

CODEOWNERS

@@ -0,0 +1,24 @@
+# IMPORTANT:
+# This file is ONLY used to subscribe for notifications for PRs
+# related to a specific file path. Approvals from people in this
+# file are not required for merges.
+
+# C API
+/transformer_engine/common/include/ @ptrendx
+
+# TE/JAX
+/transformer_engine/jax/ @jberchtold-nvidia
+
+# TE/PyTorch
+/transformer_engine/pytorch/ @ksivaman
+
+# te.ops API
+/transformer_engine/pytorch/ops/ @timmoon10
+
+# Quantization kernels
+/transformer_engine/common/cast/ @Oleg-Goncharov
+
+# Attention
+/transformer_engine/pytorch/attention/ @cyanguwa
+/transformer_engine/common/fused_attn/ @cyanguwa
+/transformer_engine/jax/cpp_extensions/attention.py @KshitijLakhani