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 10, 17:45 UTC
[CI/Build] Allow user to configure NVSHMEM version via ENV or command line (#30732)
Signed-off-by: Seiji Eicher <seiji@anyscale.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Seiji Eicher committed
Jan 5, 2026 at 15:56 UTC
3c98c2d21b38dc2037dd39b77d8f4d52eda572cf
2 files changed
+24
-2
docker/Dockerfile
+3
-1
@@ -273,6 +273,7 @@ RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
273
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
274
ARG PPLX_COMMIT_HASH
275
ARG DEEPEP_COMMIT_HASH
276
+ARG NVSHMEM_VER
277
RUN --mount=type=cache,target=/root/.cache/uv \
278
mkdir -p /tmp/ep_kernels_workspace/dist && \
279
export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
@@ -280,7 +281,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
281
--workspace /tmp/ep_kernels_workspace \
282
--mode wheel \
283
${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \
283
- ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} && \
284
+ ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} \
285
+ ${NVSHMEM_VER:+--nvshmem-ver "$NVSHMEM_VER"} && \
286
find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
287
#################### EXTENSIONS BUILD IMAGE ####################
288
tools/ep_kernels/install_python_libraries.sh
+21
-1
@@ -6,11 +6,12 @@ set -ex
6
# --mode <mode> "install" (default) or "wheel"
7
# --pplx-ref <commit> pplx-kernels commit hash
8
# --deepep-ref <commit> DeepEP commit hash
9
+# --nvshmem-ver <ver> NVSHMEM version
10
11
CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
12
PPLX_COMMIT_HASH=${PPLX_COMMIT_HASH:-"12cecfd"}
13
DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
13
-NVSHMEM_VER=3.3.24 # Suppports both CUDA 12 and 13
14
+NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"} # Default supports both CUDA 12 and 13
15
WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
16
MODE=${MODE:-install}
17
CUDA_VERSION_MAJOR=$(${CUDA_HOME}/bin/nvcc --version | egrep -o "release [0-9]+" | cut -d ' ' -f 2)
@@ -50,6 +51,18 @@ while [[ $# -gt 0 ]]; do
51
DEEPEP_COMMIT_HASH="$2"
52
shift 2
53
;;
54
+ --nvshmem-ver)
55
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
56
+ echo "Error: --nvshmem-ver requires an argument." >&2
57
+ exit 1
58
+ fi
59
+ if [[ "$2" =~ / ]]; then
60
+ echo "Error: NVSHMEM version should not contain slashes." >&2
61
+ exit 1
62
+ fi
63
+ NVSHMEM_VER="$2"
64
+ shift 2
65
+ ;;
66
*)
67
echo "Error: Unknown argument '$1'" >&2
68
exit 1
@@ -57,6 +70,13 @@ while [[ $# -gt 0 ]]; do
70
esac
71
done
72
73
+# Validate NVSHMEM_VER to prevent path traversal attacks
74
+# Only allow alphanumeric characters, dots, and hyphens (typical version string chars)
75
+if [[ ! "$NVSHMEM_VER" =~ ^[a-zA-Z0-9.-]+$ ]]; then
76
+ echo "Error: NVSHMEM_VER contains invalid characters. Only alphanumeric, dots, and hyphens are allowed." >&2
77
+ exit 1
78
+fi
79
+
80
mkdir -p "$WORKSPACE"
81
82
WHEEL_DIR="$WORKSPACE/dist"