Split out Linux CUDA Python package builds into separate stages

tools/ci_build/github/azure-pipelines/build-perf-test-binaries-pipeline.yml

@@ -33,3 +33,7 @@ stages:
         cmake_build_type: Release
         cuda_version: 12.8
         docker_base_image: onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda12_x64_almalinux8_gcc14:20251017.1
+        python_version: '3.14'
+        docker_python_exe_path: '/opt/python/cp314-cp314/bin/python3.14'
+        stage_name: Linux_py_GPU_Wheels_x86_64_3_14
+        wheel_artifact_name: onnxruntime_gpu_3_14

tools/ci_build/github/azure-pipelines/stages/py-gpu-packaging-stage.yml

@@ -37,6 +37,23 @@ parameters:
     - '3.13'
     - '3.14'
 
+- name: LinuxPythonConfigurations
+  type: object
+  displayName: 'Linux Python build configurations'
+  default:
+  - python_version: '3.11'
+    docker_python_exe_path: '/opt/python/cp311-cp311/bin/python3.11'
+  - python_version: '3.12'
+    docker_python_exe_path: '/opt/python/cp312-cp312/bin/python3.12'
+  - python_version: '3.13'
+    docker_python_exe_path: '/opt/python/cp313-cp313/bin/python3.13'
+  - python_version: '3.13t'
+    docker_python_exe_path: '/opt/python/cp313-cp313t/bin/python3.13'
+  - python_version: '3.14'
+    docker_python_exe_path: '/opt/python/cp314-cp314/bin/python3.14'
+  - python_version: '3.14t'
+    docker_python_exe_path: '/opt/python/cp314-cp314t/bin/python3.14'
+
 - name: cmake_cuda_archs
   type: string
   displayName: 'CMAKE_CUDA_ARCHITECTURES for Windows'
@@ -67,11 +84,49 @@ stages:
             EP_BUILD_FLAGS: --enable_lto --use_cuda --cuda_home=$(Agent.TempDirectory)\v${{ parameters.cuda_version }} --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=${{ parameters.cmake_cuda_archs }}"
             use_tensorrt: True
 
-    - template: py-linux-gpu-stage.yml
-      parameters:
+    # Linux: one parallel stage per Python version
+    - ${{ each config in parameters.LinuxPythonConfigurations }}:
+      - template: py-linux-gpu-stage.yml
+        parameters:
+          stage_name: Linux_py_GPU_Wheels_x86_64_${{ replace(config.python_version, '.', '_') }}
           arch: 'x86_64'
           machine_pool: 'onnxruntime-Ubuntu2404-AMD-CPU'
           extra_build_arg: ${{ parameters.build_py_parameters }}
           cmake_build_type: ${{ parameters.cmake_build_type }}
           cuda_version: ${{ parameters.cuda_version }}
           docker_base_image: ${{ parameters.docker_base_image }}
+          python_version: ${{ config.python_version }}
+          docker_python_exe_path: ${{ config.docker_python_exe_path }}
+          wheel_artifact_name: onnxruntime_gpu_${{ replace(config.python_version, '.', '_') }}
+          # Only publish the large build intermediates artifact for Python 3.12. That's the only version we test now.
+          ${{ if eq(config.python_version, '3.12') }}:
+            build_intermediates_artifact_name: linux_gpu_wheel_x86_64
+
+    # Merge per-version Linux wheel artifacts into a single combined artifact for downstream consumers
+    - stage: Linux_py_GPU_Wheels_Merge_Artifacts
+      dependsOn:
+      - ${{ each config in parameters.LinuxPythonConfigurations }}:
+        - Linux_py_GPU_Wheels_x86_64_${{ replace(config.python_version, '.', '_') }}
+      jobs:
+      - job: Linux_py_GPU_Wheels_Merge_Artifacts
+        workspace:
+          clean: all
+        pool:
+          name: 'onnxruntime-Ubuntu2404-AMD-CPU'
+          os: linux
+        templateContext:
+          outputs:
+          - output: pipelineArtifact
+            targetPath: $(Build.ArtifactStagingDirectory)/onnxruntime_gpu
+            artifactName: onnxruntime_gpu
+        steps:
+        - checkout: self
+          clean: true
+          submodules: none
+
+        - ${{ each config in parameters.LinuxPythonConfigurations }}:
+          - task: DownloadPipelineArtifact@2
+            displayName: 'Download wheel - Python ${{ config.python_version }}'
+            inputs:
+              artifact: onnxruntime_gpu_${{ replace(config.python_version, '.', '_') }}
+              targetPath: $(Build.ArtifactStagingDirectory)/onnxruntime_gpu

tools/ci_build/github/azure-pipelines/stages/py-linux-gpu-stage.yml

@@ -28,11 +28,35 @@ parameters:
    - 12.8
    - 13.0
 
+- name: python_version
+  type: string
+  displayName: 'Python version label (e.g. 3.12, 3.13t)'
+
+- name: docker_python_exe_path
+  type: string
+  displayName: 'Full path to the Python executable inside the Docker image'
+
+- name: stage_name
+  type: string
+  displayName: 'Unique stage name for this build'
+
+- name: wheel_artifact_name
+  type: string
+  displayName: 'Name for the published wheel artifact'
+
+- name: build_intermediates_artifact_name
+  type: string
+  default: ''
+  displayName: |
+    Name for the published build intermediates artifact.
+    If the name is empty, no build intermediates artifact will be published.
+    The build intermediates can be used for testing.
+
 stages:
-- stage: Linux_py_GPU_Wheels_${{ parameters.arch }}
+- stage: ${{ parameters.stage_name }}
   dependsOn: []
   jobs:
-  - job: Linux_py_GPU_Wheels_${{ parameters.arch }}
+  - job: ${{ parameters.stage_name }}
     timeoutInMinutes: 360
     workspace:
       clean: all
@@ -43,15 +67,16 @@ stages:
       outputs:
       - output: pipelineArtifact
         targetPath: $(Build.ArtifactStagingDirectory)/dist
-        artifactName: onnxruntime_gpu
-      - output: pipelineArtifact
-        targetPath: $(Build.ArtifactStagingDirectory)/${{ parameters.cmake_build_type }}
-        artifactName: linux_gpu_wheel_${{ parameters.arch }}
+        artifactName: ${{ parameters.wheel_artifact_name }}
+      - ${{ if ne(parameters.build_intermediates_artifact_name, '') }}:
+        - output: pipelineArtifact
+          targetPath: $(Build.ArtifactStagingDirectory)/${{ parameters.cmake_build_type }}
+          artifactName: ${{ parameters.build_intermediates_artifact_name }}
     variables:
       - name: extra_build_args
         ${{ if ne(parameters.extra_build_arg, '') }}:
           value: '-x "${{ parameters.extra_build_arg }}"'
-        ${{ if eq(parameters.extra_build_arg, '') }}:
+        ${{ else }}:
           value: ''
       - template: ../templates/common-variables.yml
       - name: trt_version
@@ -75,24 +100,28 @@ stages:
 
 
       - task: Bash@3
-        displayName: 'Build Python Wheel'
+        displayName: 'Build Python Wheel - ${{ parameters.python_version }}'
         inputs:
           targetType: filePath
           filePath: tools/ci_build/github/linux/run_python_dockerbuild.sh
-          arguments: -i onnxruntimecuda${{ replace(parameters.cuda_version, '.', '') }}xtrt86build${{ parameters.arch }} -d "GPU" -c ${{ parameters.cmake_build_type }} $(extra_build_args)
+          arguments: >-
+            -i onnxruntimecuda${{ replace(parameters.cuda_version, '.', '') }}xtrt86build${{ parameters.arch }}
+            -d "GPU"
+            -c ${{ parameters.cmake_build_type }}
+            -p "${{ parameters.docker_python_exe_path }}"
+            $(extra_build_args)
         env:
           CUDA_VERSION: ${{ parameters.cuda_version }}
 
       - script: |
           set -e -x
-          mv $(Build.BinariesDirectory)/${{ parameters.cmake_build_type }} ./${{ parameters.cmake_build_type }}
           mv $(Build.BinariesDirectory)/dist ./dist
-          pushd dist
-          find . -name \*.whl -exec unzip -qq -o {} \;
-          rm -r onnxruntime
-          popd
-          pushd ${{ parameters.cmake_build_type }}
-          find . -name \*.whl -exec unzip -qq -o {} \;
-          popd
         workingDirectory: '$(Build.ArtifactStagingDirectory)'
-        displayName: 'Move files'
+        displayName: 'Move wheel files'
+
+      - ${{ if ne(parameters.build_intermediates_artifact_name, '') }}:
+        - script: |
+            set -e -x
+            mv $(Build.BinariesDirectory)/${{ parameters.cmake_build_type }} ./${{ parameters.cmake_build_type }}
+          workingDirectory: '$(Build.ArtifactStagingDirectory)'
+          displayName: 'Move build intermediates'