NVIDIA · bvolpato · Aug 4, 2026 · Aug 4, 2026
diff --git a/.github/actions/build-pytorch-wheel/Dockerfile b/.github/actions/build-pytorch-wheel/Dockerfile
@@ -36,10 +36,18 @@ RUN CUDA_MAJOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'}) && \
     dpkg -i cuda-keyring_1.1-1_all.deb && \
     rm cuda-keyring_1.1-1_all.deb && \
     apt-get update && \
-    apt-get install -y cuda-toolkit-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} cudnn-cuda-$CUDA_MAJOR_VERSION libcudnn$CUDNN_MAJOR_VERSION-cuda-$CUDA_MAJOR_VERSION libnccl2 libnccl-dev cmake
+    apt-get install -y \
+        cuda-compiler-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
+        cuda-libraries-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
+        cuda-nvtx-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
+        cudnn-cuda-$CUDA_MAJOR_VERSION \
+        libcudnn$CUDNN_MAJOR_VERSION-cuda-$CUDA_MAJOR_VERSION \
+        libnccl2 \
+        libnccl-dev \
+        cmake
 
 # 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; 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)") && \
+    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), '2.11': (130, 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}
diff --git a/.github/workflows/attach-wheels-to-release.yml b/.github/workflows/attach-wheels-to-release.yml
@@ -78,7 +78,27 @@ jobs:
               "torch-version": ["2.8.0"], 
               "cuda-version": ["12.9.1"], 
               "cudnn-version": ["9"], 
-              "cxx11_abi": ["TRUE"]
+              "cxx11_abi": ["TRUE"],
+              "include": [
+                {
+                  "os": "ubuntu-22.04",
+                  "release-version": "${{ github.event.release.tag_name }}",
+                  "python-version": "3.12",
+                  "torch-version": "2.11.0",
+                  "cuda-version": "13.0.0",
+                  "cudnn-version": "9",
+                  "cxx11_abi": "TRUE"
+                },
+                {
+                  "os": "ubuntu-22.04-arm",
+                  "release-version": "${{ github.event.release.tag_name }}",
+                  "python-version": "3.12",
+                  "torch-version": "2.11.0",
+                  "cuda-version": "13.0.0",
+                  "cudnn-version": "9",
+                  "cxx11_abi": "TRUE"
+                }
+              ]
             }' | jq -rc)
           else
             MATRIX=$(echo '{
@@ -166,6 +186,71 @@ jobs:
           asset_name: ${{ steps.build-pytorch-wheel.outputs.wheel_name }}
           asset_content_type: application/*
 
+  build_core_wheel:
+    name: Build Core Wheel (ubuntu-22.04-arm, CUDA 13)
+    runs-on: ubuntu-22.04-arm
+    needs: pre-flight
+    if: ${{ github.event_name == 'release' || inputs.ngc-image == '' }}
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+        with:
+          ref: ${{ github.event_name == 'release' && github.event.release.tag_name || inputs.release-version }}
+          submodules: recursive
+
+      - name: Move /var/lib/docker/
+        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/
+        run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"
+
+      - name: Build CUDA 13 core wheel
+        id: build-core-wheel
+        run: |
+          docker build \
+            --build-arg CUDA_MAJOR=13 \
+            --build-arg CUDA_MINOR=0 \
+            --build-arg BUILD_METAPACKAGE=false \
+            --build-arg BUILD_COMMON=true \
+            --build-arg BUILD_PYTORCH=false \
+            --build-arg BUILD_JAX=false \
+            --tag transformer-engine-core-wheel \
+            --file build_tools/wheel_utils/Dockerfile.aarch \
+            .
+
+          mkdir -p dist
+          docker run --rm \
+            --volume "$(pwd)/dist:/wheelhouse" \
+            transformer-engine-core-wheel
+
+          wheel_path=$(find dist -maxdepth 1 -type f \
+            -name 'transformer_engine_cu13-*-manylinux_2_28_aarch64.whl' \
+            -print -quit)
+          test -n "$wheel_path"
+          echo "wheel_name=$(basename "$wheel_path")" | tee -a "$GITHUB_OUTPUT"
+
+      - name: Upload release asset
+        uses: actions/upload-release-asset@e8f9f06c4b078e705bd2ea027f0926603fc9b4d5 # v1.0.2
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          upload_url: ${{ needs.pre-flight.outputs.release-assets-url }}
+          asset_path: ./dist/${{ steps.build-core-wheel.outputs.wheel_name }}
+          asset_name: ${{ steps.build-core-wheel.outputs.wheel_name }}
+          asset_content_type: application/*
+
   build_wheels_for_ngc:
     name: Build Wheels for NGC PyTorch images
     runs-on: ${{ matrix.os }}

diff --git a/build_tools/wheel_utils/Dockerfile.aarch b/build_tools/wheel_utils/Dockerfile.aarch
@@ -25,13 +25,12 @@ ENV CUDA_MAJOR=${CUDA_MAJOR}
 RUN dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo
 RUN dnf -y install cuda-compiler-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \
                    cuda-libraries-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \
-                   cuda-libraries-devel-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64
+                   cuda-libraries-devel-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \
+                   cuda-nvtx-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64
 RUN dnf -y install --allowerasing cudnn9-cuda-${CUDA_MAJOR}
 RUN dnf clean all
 RUN rm -rf /var/cache/dnf/*
 RUN echo "/usr/local/cuda/lib64" >> /etc/ld.so.conf.d/999_nvidia_cuda.conf
-RUN dnf -y install cuda-toolkit-${CUDA_MAJOR}
-RUN dnf clean all
 RUN dnf -y install glog.aarch64 glog-devel.aarch64
 RUN dnf -y install libnccl libnccl-devel libnccl-static