diff --git a/docker/Dockerfile b/docker/Dockerfile
index 84a1802dbe03..aa3aad21d6c0 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -244,9 +244,15 @@ RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
# Install EP kernels(pplx-kernels and DeepEP)
+ARG PPLX_COMMIT_HASH
+ARG DEEPEP_COMMIT_HASH
RUN --mount=type=cache,target=/root/.cache/uv \
export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
- /tmp/install_python_libraries.sh /tmp/ep_kernels_workspace wheel && \
+ /tmp/install_python_libraries.sh \
+ --workspace /tmp/ep_kernels_workspace \
+ --mode wheel \
+ ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \
+ ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} && \
find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
# Check the size of the wheel if RUN_WHEEL_CHECK is true
diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh
index 1cea1bef8dbc..88be5cd778ff 100755
--- a/tools/ep_kernels/install_python_libraries.sh
+++ b/tools/ep_kernels/install_python_libraries.sh
@@ -1,22 +1,68 @@
#!/usr/bin/env bash
set -ex
-# usage: ./build.sh [workspace_dir] [mode]
-# mode: "install" (default) → install directly into current Python env
-# "wheel" → build wheels into WORKSPACE/dist
+# usage: ./install_python_libraries.sh [options]
+# --workspace
workspace directory (default: ./ep_kernels_workspace)
+# --mode "install" (default) or "wheel"
+# --pplx-ref pplx-kernels commit hash
+# --deepep-ref DeepEP commit hash
+
+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
+WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
+MODE=${MODE:-install}
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --workspace)
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
+ echo "Error: --workspace requires an argument." >&2
+ exit 1
+ fi
+ WORKSPACE="$2"
+ shift 2
+ ;;
+ --mode)
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
+ echo "Error: --mode requires an argument." >&2
+ exit 1
+ fi
+ MODE="$2"
+ shift 2
+ ;;
+ --pplx-ref)
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
+ echo "Error: --pplx-ref requires an argument." >&2
+ exit 1
+ fi
+ PPLX_COMMIT_HASH="$2"
+ shift 2
+ ;;
+ --deepep-ref)
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
+ echo "Error: --deepep-ref requires an argument." >&2
+ exit 1
+ fi
+ DEEPEP_COMMIT_HASH="$2"
+ shift 2
+ ;;
+ *)
+ echo "Error: Unknown argument '$1'" >&2
+ exit 1
+ ;;
+ esac
+done
-WORKSPACE=${1:-$(pwd)/ep_kernels_workspace}
-MODE=${2:-install}
mkdir -p "$WORKSPACE"
WHEEL_DIR="$WORKSPACE/dist"
mkdir -p "$WHEEL_DIR"
-NVSHMEM_VER=3.3.9
pushd "$WORKSPACE"
-CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
-
# install dependencies if not installed
if [ -z "$VIRTUAL_ENV" ]; then
uv pip install --system cmake torch ninja
@@ -133,7 +179,7 @@ do_build \
"https://github.com/ppl-ai/pplx-kernels" \
"pplx-kernels" \
"setup.py" \
- "12cecfd" \
+ "$PPLX_COMMIT_HASH" \
""
# build DeepEP
@@ -141,7 +187,7 @@ do_build \
"https://github.com/deepseek-ai/DeepEP" \
"DeepEP" \
"setup.py" \
- "73b6ea4" \
+ "$DEEPEP_COMMIT_HASH" \
"export NVSHMEM_DIR=$WORKSPACE/nvshmem; "
if [ "$MODE" = "wheel" ]; then