A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 20:00 UTC
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>
Dmitry Tokarev committed
Dec 8, 2025 at 15:24 UTC
799804d140fc99ce3964648ba91aaa810cf28fef
1 file changed
+6
-12
tools/ep_kernels/install_python_libraries.sh
+6
-12
@@ -10,9 +10,10 @@ set -ex
10
CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
11
PPLX_COMMIT_HASH=${PPLX_COMMIT_HASH:-"12cecfd"}
12
DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
13
-NVSHMEM_VER=3.3.9
13
+NVSHMEM_VER=3.3.24 # Suppports both CUDA 12 and 13
14
WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
15
MODE=${MODE:-install}
16
+CUDA_VERSION_MAJOR=$(${CUDA_HOME}/bin/nvcc --version | egrep -o "release [0-9]+" | cut -d ' ' -f 2)
17
18
# Parse arguments
19
while [[ $# -gt 0 ]]; do
@@ -75,11 +76,9 @@ ARCH=$(uname -m)
76
case "${ARCH,,}" in
77
x86_64|amd64)
78
NVSHMEM_SUBDIR="linux-x86_64"
78
- NVSHMEM_FILE="libnvshmem-linux-x86_64-${NVSHMEM_VER}_cuda12-archive.tar.xz"
79
;;
80
aarch64|arm64)
81
NVSHMEM_SUBDIR="linux-sbsa"
82
- NVSHMEM_FILE="libnvshmem-linux-sbsa-${NVSHMEM_VER}_cuda12-archive.tar.xz"
82
;;
83
*)
84
echo "Unsupported architecture: ${ARCH}" >&2
@@ -87,6 +86,7 @@ case "${ARCH,,}" in
86
;;
87
esac
88
89
+NVSHMEM_FILE="libnvshmem-${NVSHMEM_SUBDIR}-${NVSHMEM_VER}_cuda${CUDA_VERSION_MAJOR}-archive.tar.xz"
90
NVSHMEM_URL="https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/${NVSHMEM_SUBDIR}/${NVSHMEM_FILE}"
91
92
pushd "$WORKSPACE"
@@ -142,13 +142,6 @@ clone_repo() {
142
fi
143
}
144
145
-deepep_cuda13_patch() {
146
- cuda_version_major=$(${CUDA_HOME}/bin/nvcc --version | egrep -o "release [0-9]+" | cut -d ' ' -f 2)
147
- if [ ${cuda_version_major} -ge 13 ]; then
148
- sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py"
149
- fi
150
-}
151
-
145
do_build() {
146
local repo=$1
147
local name=$2
@@ -160,8 +153,9 @@ do_build() {
153
clone_repo "$repo" "$name" "$key" "$commit"
154
cd "$name"
155
163
- if [ "$name" == "DeepEP" ]; then
164
- deepep_cuda13_patch
156
+ # DeepEP CUDA 13 patch
157
+ if [[ "$name" == "DeepEP" && "${CUDA_VERSION_MAJOR}" -ge 13 ]]; then
158
+ sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py"
159
fi
160
161
if [ "$MODE" = "install" ]; then