From 00c8a430ed9e0d89832645655cd0a0f2b299145c Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Wed, 19 Aug 2026 06:02:28 +0000
Subject: [PATCH 1/2] update aarch64 parallel to 8

---
 .../stages/plugin-linux-cuda-stage.yml                |  1 +
 .../github/linux/build_cuda_plugin_package.sh         | 11 ++++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml
index 4eeb7c490af26..70114ccfcc52a 100644
--- a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml
+++ b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml
@@ -130,6 +130,7 @@ stages:
           -p ${{ parameters.docker_python_exe_path }}
           -v ${{ parameters.cuda_version }}
           -a "${{ parameters.cmake_cuda_archs }}"
+          -h "${{ parameters.arch }}"
         workingDirectory: $(Build.SourcesDirectory)
         displayName: 'Build CUDA Plugin (Python ${{ parameters.python_version }}, ${{ parameters.arch }}, CUDA ${{ parameters.cuda_version }})'
         env:
diff --git a/tools/ci_build/github/linux/build_cuda_plugin_package.sh b/tools/ci_build/github/linux/build_cuda_plugin_package.sh
index f12811ba087c4..cdd557a25230c 100755
--- a/tools/ci_build/github/linux/build_cuda_plugin_package.sh
+++ b/tools/ci_build/github/linux/build_cuda_plugin_package.sh
@@ -9,8 +9,9 @@ DOCKER_IMAGE="onnxruntimecuda128pluginbuild"
 PYTHON_EXE="/opt/python/cp312-cp312/bin/python3.12"
 CUDA_VERSION="12.8"
 CMAKE_CUDA_ARCHS="86"
+arch="x64"
 
-while getopts "i:c:p:v:a:" parameter_Option
+while getopts "i:c:p:v:a:h:" parameter_Option
 do case "${parameter_Option}"
 in
 i) DOCKER_IMAGE=${OPTARG};;
@@ -18,11 +19,15 @@ c) BUILD_CONFIG=${OPTARG};;
 p) PYTHON_EXE=${OPTARG};;
 v) CUDA_VERSION=${OPTARG};;
 a) CMAKE_CUDA_ARCHS=${OPTARG};;
-*) echo "Usage: $0 -i <docker_image> [-c <build_config>] [-p <python_exe_path>] [-v <cuda_version>] [-a <cuda_archs>]"
+h) arch=${OPTARG};;
+*) echo "Usage: $0 -i <docker_image> [-c <build_config>] [-p <python_exe_path>] [-v <cuda_version>] [-a <cuda_archs>] [-h <arch>]"
    exit 1;;
 esac
 done
 
+# Update parallel to the minimum of the number of processors and 8 for aarch64 to avoid OOM errors during build.
+PARALLEL=$([ "$arch" = "aarch64" ] && echo "$(( $(nproc) < 8 ? $(nproc) : 8 ))")
+
 PYTHON_BIN_DIR=$(dirname "${PYTHON_EXE}")
 
 # Derive SHORT_CUDA_VERSION (e.g., 12.8 from 12.8, 13.0 from 13.0)
@@ -54,7 +59,7 @@ docker run --rm \
         --build_dir /build \
         --config ${BUILD_CONFIG} \
         --skip_submodule_sync \
-        --parallel \
+        --parallel ${PARALLEL} \
         --nvcc_threads 1 \
         --flash_nvcc_threads 1 \
         --use_binskim_compliant_compile_flags \

From 198a6de72ddd689c4d3c7b4de3913d05d1e990d0 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Wed, 19 Aug 2026 16:59:59 +0000
Subject: [PATCH 2/2] fix x64

---
 tools/ci_build/github/linux/build_cuda_plugin_package.sh | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tools/ci_build/github/linux/build_cuda_plugin_package.sh b/tools/ci_build/github/linux/build_cuda_plugin_package.sh
index cdd557a25230c..a8bc55df45aca 100755
--- a/tools/ci_build/github/linux/build_cuda_plugin_package.sh
+++ b/tools/ci_build/github/linux/build_cuda_plugin_package.sh
@@ -26,7 +26,11 @@ esac
 done
 
 # Update parallel to the minimum of the number of processors and 8 for aarch64 to avoid OOM errors during build.
-PARALLEL=$([ "$arch" = "aarch64" ] && echo "$(( $(nproc) < 8 ? $(nproc) : 8 ))")
+if [ "$arch" = "aarch64" ]; then
+    PARALLEL=$(( $(nproc) < 8 ? $(nproc) : 8 ))
+else
+    PARALLEL=""
+fi
 
 PYTHON_BIN_DIR=$(dirname "${PYTHON_EXE}")