From 29b1dc3f12a081d3be60d6131dad9258a901eb54 Mon Sep 17 00:00:00 2001
From: ksivamani <ksivamani@nvidia.com>
Date: Sat, 16 May 2026 17:30:27 -0700
Subject: [PATCH] Add optional core features to wheel build

Signed-off-by: ksivamani <ksivamani@nvidia.com>
---
 build_tools/wheel_utils/Dockerfile.aarch | 12 ++++++++--
 build_tools/wheel_utils/Dockerfile.x86   | 12 ++++++++--
 build_tools/wheel_utils/build_wheels.sh  | 29 ++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/build_tools/wheel_utils/Dockerfile.aarch b/build_tools/wheel_utils/Dockerfile.aarch
index c040dadcdb..29af4a0d08 100644
--- a/build_tools/wheel_utils/Dockerfile.aarch
+++ b/build_tools/wheel_utils/Dockerfile.aarch
@@ -35,12 +35,20 @@ RUN dnf clean all
 RUN dnf -y install glog.aarch64 glog-devel.aarch64
 RUN dnf -y install libnccl libnccl-devel libnccl-static
 
-ENV PATH="/usr/local/cuda/bin:${PATH}"
-ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
+RUN dnf -y install openmpi openmpi-devel && dnf clean all
+RUN mkdir -p /opt/mpi && \
+    ln -s /usr/lib64/openmpi/bin /opt/mpi/bin && \
+    ln -s /usr/lib64/openmpi/lib /opt/mpi/lib && \
+    ln -s /usr/include/openmpi-aarch64 /opt/mpi/include && \
+    echo "/usr/lib64/openmpi/lib" >> /etc/ld.so.conf.d/openmpi-aarch64.conf
+
+ENV PATH="/usr/local/cuda/bin:/opt/mpi/bin:${PATH}"
+ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:/opt/mpi/lib:${LD_LIBRARY_PATH}"
 ENV CUDA_HOME=/usr/local/cuda
 ENV CUDA_ROOT=/usr/local/cuda
 ENV CUDA_PATH=/usr/local/cuda
 ENV CUDADIR=/usr/local/cuda
+ENV MPI_HOME=/opt/mpi
 ENV NVTE_RELEASE_BUILD=1
 
 CMD ["/bin/bash", "-c", "bash /TransformerEngine/build_tools/wheel_utils/build_wheels.sh manylinux_2_28_aarch64 $BUILD_METAPACKAGE $BUILD_COMMON $BUILD_PYTORCH $BUILD_JAX $CUDA_MAJOR"]
diff --git a/build_tools/wheel_utils/Dockerfile.x86 b/build_tools/wheel_utils/Dockerfile.x86
index 2728b6b7c1..609b2af093 100644
--- a/build_tools/wheel_utils/Dockerfile.x86
+++ b/build_tools/wheel_utils/Dockerfile.x86
@@ -35,12 +35,20 @@ RUN dnf clean all
 RUN dnf -y install glog.x86_64 glog-devel.x86_64
 RUN dnf -y install libnccl libnccl-devel libnccl-static
 
-ENV PATH="/usr/local/cuda/bin:${PATH}"
-ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
+RUN dnf -y install openmpi openmpi-devel && dnf clean all
+RUN mkdir -p /opt/mpi && \
+    ln -s /usr/lib64/openmpi/bin /opt/mpi/bin && \
+    ln -s /usr/lib64/openmpi/lib /opt/mpi/lib && \
+    ln -s /usr/include/openmpi-x86_64 /opt/mpi/include && \
+    echo "/usr/lib64/openmpi/lib" >> /etc/ld.so.conf.d/openmpi-x86_64.conf
+
+ENV PATH="/usr/local/cuda/bin:/opt/mpi/bin:${PATH}"
+ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:/opt/mpi/lib:${LD_LIBRARY_PATH}"
 ENV CUDA_HOME=/usr/local/cuda
 ENV CUDA_ROOT=/usr/local/cuda
 ENV CUDA_PATH=/usr/local/cuda
 ENV CUDADIR=/usr/local/cuda
+ENV MPI_HOME=/opt/mpi
 ENV NVTE_RELEASE_BUILD=1
 
 CMD ["/bin/bash", "-c", "bash /TransformerEngine/build_tools/wheel_utils/build_wheels.sh manylinux_2_28_x86_64 $BUILD_METAPACKAGE $BUILD_COMMON $BUILD_PYTORCH $BUILD_JAX $CUDA_MAJOR"]
diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh
index e9ec854dba..2ae64bf52b 100644
--- a/build_tools/wheel_utils/build_wheels.sh
+++ b/build_tools/wheel_utils/build_wheels.sh
@@ -24,6 +24,35 @@ git submodule update --init --recursive
 
 # Install deps
 /opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel
+/opt/python/cp310-cp310/bin/pip install \
+        "nvidia-cublasmp-cu${CUDA_MAJOR}" \
+        "nvidia-cusolvermp-cu${CUDA_MAJOR}" \
+        "nvidia-nvshmem-cu${CUDA_MAJOR}"
+
+SITE_PACKAGES=$(/opt/python/cp310-cp310/bin/python -c "import sysconfig; print(sysconfig.get_paths()['purelib'])")
+export CUBLASMP_HOME="${SITE_PACKAGES}/nvidia/cublasmp/cu${CUDA_MAJOR}"
+export CUSOLVERMP_HOME="${SITE_PACKAGES}/nvidia/cu${CUDA_MAJOR}"
+export NVSHMEM_HOME="${SITE_PACKAGES}/nvidia/nvshmem"
+
+# nvidia-cuda-python package compatibility.
+for lib_dir in "${CUBLASMP_HOME}/lib" "${CUSOLVERMP_HOME}/lib" "${NVSHMEM_HOME}/lib" ; do
+        [ -d "$lib_dir" ] || continue
+        for so in "$lib_dir"/lib*.so.* ; do
+                [ -e "$so" ] || continue
+                base=$(basename "$so")
+                unversioned="${base%%.so.*}.so"
+                ln -sf "$base" "${lib_dir}/${unversioned}"
+        done
+done
+
+# Enable optional build features and expose the runtime libs to the linker.
+export NVTE_WITH_CUBLASMP=1
+export NVTE_WITH_CUSOLVERMP=1
+export NVTE_ENABLE_NVSHMEM=1
+export NVTE_UB_WITH_MPI=1
+export MPI_HOME="${MPI_HOME:-/opt/mpi}"
+export LD_LIBRARY_PATH="${NVSHMEM_HOME}/lib:${CUBLASMP_HOME}/lib:${CUSOLVERMP_HOME}/lib:${MPI_HOME}/lib:${LD_LIBRARY_PATH}"
+export PATH="${MPI_HOME}/bin:${PATH}"
 
 if $BUILD_METAPACKAGE ; then
         cd /TransformerEngine