[CI] Publish PyTorch 2.11 CUDA 13 release wheels

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

@@ -9,7 +9,7 @@ ENV DEBIAN_FRONTEND=noninteractive
 ENV CUDA_HOME=/usr/local/cuda
 ENV PATH=$PATH:$CUDA_HOME/bin
 ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
-ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;9.0"
+ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;9.0;10.0"
 
 ARG PYTHON_VERSION=3.12
 ARG TORCH_VERSION=2.9.1
@@ -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}

.github/actions/build-pytorch-wheel/action.yml

@@ -30,6 +30,10 @@ inputs:
   aarch:
     description: 'The architecture to use for the build'
     required: true
+  smoke-test:
+    description: 'Install and import the wheel after building it'
+    required: false
+    default: 'false'
 outputs:
   wheel_name:
     description: 'The name of the built wheel'
@@ -93,24 +97,35 @@ runs:
       shell: bash -euxo pipefail {0}
       id: build_wheel
       env:
+        AARCH: ${{ inputs.aarch }}
         CXX11_ABI: ${{ inputs.cxx11_abi }}
+        PYTHON_VERSION: ${{ inputs.python-version }}
+        RELEASE_VERSION: ${{ inputs.release-version }}
+        SMOKE_TEST: ${{ inputs.smoke-test }}
       run: |
         echo ::group::Build wheel
 
-        EXIT_CODE=$(docker run \
+        set +e
+        docker run \
             --rm \
             --shm-size=64g \
             --workdir /workspace/transformer_engine/pytorch \
-            --volume $(pwd):/workspace \
-            --volume $GITHUB_OUTPUT:$GITHUB_OUTPUT \
+            --volume "$(pwd):/workspace" \
+            --volume "$GITHUB_OUTPUT:$GITHUB_OUTPUT" \
+            -e AARCH="$AARCH" \
             -e PIP_CONSTRAINT= \
-            -e CXX11_ABI=$CXX11_ABI \
-            -e GITHUB_OUTPUT=$GITHUB_OUTPUT \
-            transformer-engine-build bash /workspace/build-tools/.github/actions/build-pytorch-wheel/build.sh | tail -n 1)
-
-        # Do not fail the job if timeout killed the build
-        exit $EXIT_CODE
+            -e CXX11_ABI="$CXX11_ABI" \
+            -e GITHUB_OUTPUT="$GITHUB_OUTPUT" \
+            -e MAX_JOBS \
+            -e NVTE_FRAMEWORK \
+            -e PYTHON_VERSION="$PYTHON_VERSION" \
+            -e RELEASE_VERSION="$RELEASE_VERSION" \
+            -e SMOKE_TEST="$SMOKE_TEST" \
+            transformer-engine-build bash /workspace/build-tools/.github/actions/build-pytorch-wheel/build.sh
+        exit_code=$?
+        set -e
         echo ::endgroup::
+        exit "$exit_code"
 
     - name: Log Built Wheels
       shell: bash -euxo pipefail {0}

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

@@ -14,13 +14,68 @@ export PIP_CONSTRAINT=
 pip install wheel packaging nvidia-mathdx ninja pybind11
 
 # 5h timeout since GH allows max 6h and we want some buffer
-EXIT_CODE=0
-timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$?
+timeout 5h python setup.py bdist_wheel --dist-dir=dist
+
+mapfile -t wheel_paths < <(find dist -maxdepth 1 -type f -name '*.whl' -print)
+if [[ ${#wheel_paths[@]} -ne 1 ]]; then
+    echo "Expected one built wheel, found ${#wheel_paths[@]}" >&2
+    exit 1
+fi
+built_wheel=${wheel_paths[0]}
+
+if [[ -n "${AARCH:-}" ]]; then
+    case "$AARCH" in
+        x86_64) platform_tag=linux_x86_64 ;;
+        sbsa) platform_tag=linux_aarch64 ;;
+        *) echo "Unsupported wheel architecture: $AARCH" >&2; exit 1 ;;
+    esac
+
+    python_tag="cp${PYTHON_VERSION//./}"
+    expected_tag="${python_tag}-${python_tag}-${platform_tag}"
+    python - "$built_wheel" "$expected_tag" <<'PY'
+import sys
+import zipfile
+
+wheel_path, expected_tag = sys.argv[1:]
+with zipfile.ZipFile(wheel_path) as wheel:
+    wheel_metadata = next(name for name in wheel.namelist() if name.endswith(".dist-info/WHEEL"))
+    tags = {
+        line.removeprefix("Tag:").strip()
+        for line in wheel.read(wheel_metadata).decode().splitlines()
+        if line.startswith("Tag:")
+    }
+if expected_tag not in tags:
+    raise SystemExit(f"Expected wheel tag {expected_tag}, found {sorted(tags)}")
+PY
 
-if [ $EXIT_CODE -eq 0 ]; then
     wheel_name=$(python -c "import setup; print(setup.get_wheel_url()[1])" | tail -n 1)
-    ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
-    echo "wheel_name=${wheel_name}" | tee -a "$GITHUB_OUTPUT"
+    expected_suffix="cxx11abi${CXX11_ABI}-${expected_tag}.whl"
+    if [[ "$wheel_name" != *"$expected_suffix" ]]; then
+        echo "Expected wheel filename suffix $expected_suffix, found $wheel_name" >&2
+        exit 1
+    fi
+else
+    wheel_name=$(python -c "import setup; print(setup.get_wheel_url()[1])" | tail -n 1)
+fi
+
+if [[ "${SMOKE_TEST:-false}" == "true" ]]; then
+    release_version=${RELEASE_VERSION:-}
+    release_version=${release_version#v}
+    if [[ -z "$release_version" ]]; then
+        echo "Release version is required for the install smoke test" >&2
+        exit 1
+    fi
+    pip install --no-cache-dir "transformer-engine==${release_version}" "$built_wheel"
+
+    cuda_stub_dir=$(mktemp -d)
+    test -f /usr/local/cuda/lib64/stubs/libcuda.so
+    ln -s /usr/local/cuda/lib64/stubs/libcuda.so "$cuda_stub_dir/libcuda.so.1"
+    (
+        cd /tmp
+        LD_LIBRARY_PATH="$cuda_stub_dir:${LD_LIBRARY_PATH:-}" python -c \
+            "import transformer_engine.pytorch; print('Transformer Engine import OK')"
+    )
 fi
 
-echo $EXIT_CODE
+mv "$built_wheel" "dist/$wheel_name"
+echo "wheel_name=${wheel_name}" | tee -a "$GITHUB_OUTPUT"

.github/workflows/attach-wheels-to-release.yml

@@ -78,7 +78,30 @@ jobs:
               "torch-version": ["2.8.0"], 
               "cuda-version": ["12.9.1"], 
               "cudnn-version": ["9"], 
-              "cxx11_abi": ["TRUE"]
+              "cxx11_abi": ["TRUE"],
+              "smoke-test": ["false"],
+              "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",
+                  "smoke-test": "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",
+                  "smoke-test": "true"
+                }
+              ]
             }' | jq -rc)
           else
             MATRIX=$(echo '{
@@ -151,6 +174,7 @@ jobs:
           torch-version: ${{ matrix.torch-version }}
           cxx11_abi: ${{ matrix.cxx11_abi }}
           aarch: ${{ matrix.os == 'ubuntu-22.04' && 'x86_64' || 'sbsa' }}
+          smoke-test: ${{ matrix.smoke-test || 'false' }}
         env:
           NVTE_FRAMEWORK: pytorch
           MAX_JOBS: 1