Bump nvshmem to 3.3.24 and fix CUDA 13 installation (#30149)

Signed-off-by: Dmitry Tokarev <dtokarev@nvidia.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Dmitry Tokarev 2025-12-08 15:24:34 -05:00 committed by GitHub
parent 0d402d2600
commit 799804d140
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -10,9 +10,10 @@ set -ex
CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
PPLX_COMMIT_HASH=${PPLX_COMMIT_HASH:-"12cecfd"}
DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
NVSHMEM_VER=3.3.9
NVSHMEM_VER=3.3.24 # Suppports both CUDA 12 and 13
WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
MODE=${MODE:-install}
CUDA_VERSION_MAJOR=$(${CUDA_HOME}/bin/nvcc --version | egrep -o "release [0-9]+" | cut -d ' ' -f 2)
# Parse arguments
while [[ $# -gt 0 ]]; do
@ -75,11 +76,9 @@ ARCH=$(uname -m)
case "${ARCH,,}" in
x86_64|amd64)
NVSHMEM_SUBDIR="linux-x86_64"
NVSHMEM_FILE="libnvshmem-linux-x86_64-${NVSHMEM_VER}_cuda12-archive.tar.xz"
;;
aarch64|arm64)
NVSHMEM_SUBDIR="linux-sbsa"
NVSHMEM_FILE="libnvshmem-linux-sbsa-${NVSHMEM_VER}_cuda12-archive.tar.xz"
;;
*)
echo "Unsupported architecture: ${ARCH}" >&2
@ -87,6 +86,7 @@ case "${ARCH,,}" in
;;
esac
NVSHMEM_FILE="libnvshmem-${NVSHMEM_SUBDIR}-${NVSHMEM_VER}_cuda${CUDA_VERSION_MAJOR}-archive.tar.xz"
NVSHMEM_URL="https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/${NVSHMEM_SUBDIR}/${NVSHMEM_FILE}"
pushd "$WORKSPACE"
@ -142,13 +142,6 @@ clone_repo() {
fi
}
deepep_cuda13_patch() {
cuda_version_major=$(${CUDA_HOME}/bin/nvcc --version | egrep -o "release [0-9]+" | cut -d ' ' -f 2)
if [ ${cuda_version_major} -ge 13 ]; then
sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py"
fi
}
do_build() {
local repo=$1
local name=$2
@ -160,8 +153,9 @@ do_build() {
clone_repo "$repo" "$name" "$key" "$commit"
cd "$name"
if [ "$name" == "DeepEP" ]; then
deepep_cuda13_patch
# DeepEP CUDA 13 patch
if [[ "$name" == "DeepEP" && "${CUDA_VERSION_MAJOR}" -ge 13 ]]; then
sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py"
fi
if [ "$MODE" = "install" ]; then