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 12, 12:15 UTC
[CI/Build] allow user modify pplx and deepep ref by ENV or command line (#29131)
Signed-off-by: alec-flowers <aflowers@nvidia.com>
Alec committed
Nov 26, 2025 at 09:41 UTC
c4c0354eec86f3486285f121fa184dd6d9cacb9d
2 files changed
+63
-11
docker/Dockerfile
+7
-1
@@ -244,9 +244,15 @@ RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
244
245
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
246
# Install EP kernels(pplx-kernels and DeepEP)
247
+ARG PPLX_COMMIT_HASH
248
+ARG DEEPEP_COMMIT_HASH
249
RUN --mount=type=cache,target=/root/.cache/uv \
250
export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
249
- /tmp/install_python_libraries.sh /tmp/ep_kernels_workspace wheel && \
251
+ /tmp/install_python_libraries.sh \
252
+ --workspace /tmp/ep_kernels_workspace \
253
+ --mode wheel \
254
+ ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \
255
+ ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} && \
256
find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
257
258
# Check the size of the wheel if RUN_WHEEL_CHECK is true
tools/ep_kernels/install_python_libraries.sh
+56
-10
@@ -1,22 +1,68 @@
1
#!/usr/bin/env bash
2
set -ex
3
4
-# usage: ./build.sh [workspace_dir] [mode]
5
-# mode: "install" (default) → install directly into current Python env
6
-# "wheel" → build wheels into WORKSPACE/dist
4
+# usage: ./install_python_libraries.sh [options]
5
+# --workspace <dir> workspace directory (default: ./ep_kernels_workspace)
6
+# --mode <mode> "install" (default) or "wheel"
7
+# --pplx-ref <commit> pplx-kernels commit hash
8
+# --deepep-ref <commit> DeepEP commit hash
9
+
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
14
+WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
15
+MODE=${MODE:-install}
16
+
17
+# Parse arguments
18
+while [[ $# -gt 0 ]]; do
19
+ case $1 in
20
+ --workspace)
21
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
22
+ echo "Error: --workspace requires an argument." >&2
23
+ exit 1
24
+ fi
25
+ WORKSPACE="$2"
26
+ shift 2
27
+ ;;
28
+ --mode)
29
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
30
+ echo "Error: --mode requires an argument." >&2
31
+ exit 1
32
+ fi
33
+ MODE="$2"
34
+ shift 2
35
+ ;;
36
+ --pplx-ref)
37
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
38
+ echo "Error: --pplx-ref requires an argument." >&2
39
+ exit 1
40
+ fi
41
+ PPLX_COMMIT_HASH="$2"
42
+ shift 2
43
+ ;;
44
+ --deepep-ref)
45
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
46
+ echo "Error: --deepep-ref requires an argument." >&2
47
+ exit 1
48
+ fi
49
+ DEEPEP_COMMIT_HASH="$2"
50
+ shift 2
51
+ ;;
52
+ *)
53
+ echo "Error: Unknown argument '$1'" >&2
54
+ exit 1
55
+ ;;
56
+ esac
57
+done
58
8
-WORKSPACE=${1:-$(pwd)/ep_kernels_workspace}
9
-MODE=${2:-install}
59
mkdir -p "$WORKSPACE"
60
61
WHEEL_DIR="$WORKSPACE/dist"
62
mkdir -p "$WHEEL_DIR"
14
-NVSHMEM_VER=3.3.9
63
64
pushd "$WORKSPACE"
65
18
-CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
19
-
66
# install dependencies if not installed
67
if [ -z "$VIRTUAL_ENV" ]; then
68
uv pip install --system cmake torch ninja
@@ -133,7 +179,7 @@ do_build \
179
"https://github.com/ppl-ai/pplx-kernels" \
180
"pplx-kernels" \
181
"setup.py" \
136
- "12cecfd" \
182
+ "$PPLX_COMMIT_HASH" \
183
""
184
185
# build DeepEP
@@ -141,7 +187,7 @@ do_build \
187
"https://github.com/deepseek-ai/DeepEP" \
188
"DeepEP" \
189
"setup.py" \
144
- "73b6ea4" \
190
+ "$DEEPEP_COMMIT_HASH" \
191
"export NVSHMEM_DIR=$WORKSPACE/nvshmem; "
192
193
if [ "$MODE" = "wheel" ]; then