feat(vllm-tensorizer): Add DeepEP build with NVSHMEM support

Author: kptkin

vllm-tensorizer/Dockerfile

@@ -178,6 +178,40 @@ RUN --mount=type=bind,from=deepgemm-downloader,source=/git/DeepGEMM,target=/work
     . /opt/arch_flags.sh && \
     /opt/build.sh
 
+FROM alpine/git:2.36.3 AS deepep-downloader
+WORKDIR /git
+ARG DEEPEP_COMMIT='73b6ea4'
+RUN git clone --filter=tree:0 --no-single-branch --no-checkout \
+      https://github.com/deepseek-ai/DeepEP && \
+    cd DeepEP && \
+    git checkout "${DEEPEP_COMMIT}" && \
+    git submodule update --init --recursive --jobs 8 \
+      --depth 1 --filter=tree:0
+
+FROM builder-base AS deepep-builder
+ARG TARGETPLATFORM
+ARG NVSHMEM_VER='3.3.24'
+# DeepEP only supports x86_64 (Hopper/Blackwell GPUs)
+RUN mkdir -p /wheels /opt/nvshmem/lib && \
+    if [ "$TARGETPLATFORM" != "linux/amd64" ]; then \
+      echo "Skipping DeepEP build for $TARGETPLATFORM (x86_64 only)" && \
+      touch /wheels/.deepep_skipped; \
+      exit 0; \
+    fi && \
+    CUDA_MAJOR=$(echo ${CUDA_VERSION} | cut -d. -f1) && \
+    NVSHMEM_FILE="libnvshmem-linux-x86_64-${NVSHMEM_VER}_cuda${CUDA_MAJOR}-archive.tar.xz" && \
+    curl -fSL "https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/linux-x86_64/${NVSHMEM_FILE}" \
+      -o /tmp/nvshmem.tar.xz && \
+    tar -xf /tmp/nvshmem.tar.xz -C /opt/nvshmem --strip-components=1 && \
+    rm /tmp/nvshmem.tar.xz && \
+    rm -rf /opt/nvshmem/lib/bin /opt/nvshmem/lib/share
+RUN --mount=type=bind,from=deepep-downloader,source=/git/DeepEP,target=/workspace,rw \
+    if [ -f /wheels/.deepep_skipped ]; then exit 0; fi && \
+    . /opt/arch_flags.sh && \
+    export NVSHMEM_DIR=/opt/nvshmem && \
+    export CMAKE_PREFIX_PATH=/opt/nvshmem/lib/cmake:$CMAKE_PREFIX_PATH && \
+    /opt/build.sh
+
 FROM builder-base AS nixl-builder
 RUN apt-get -qq update && \
     apt-get -q install --no-install-recommends --no-upgrade -y \
@@ -235,6 +269,16 @@ RUN --mount=type=bind,from=lmcache-builder,source=/wheels,target=/tmp/wheels \
 RUN --mount=type=bind,from=deepgemm-builder,source=/wheels,target=/tmp/wheels \
     python3 -m pip install --no-cache-dir /tmp/wheels/*.whl -c /tmp/constraints.txt
 
+RUN --mount=type=bind,from=deepep-builder,source=/wheels,target=/tmp/wheels \
+    if ls /tmp/wheels/*.whl >/dev/null 2>&1; then \
+      python3 -m pip install --no-cache-dir /tmp/wheels/*.whl -c /tmp/constraints.txt; \
+    else \
+      echo "No DeepEP wheels found, skipping install"; \
+    fi
+
+COPY --link --from=deepep-builder /opt/nvshmem/lib /usr/local/lib/nvshmem
+ENV LD_LIBRARY_PATH=/usr/local/lib/nvshmem:$LD_LIBRARY_PATH
+
 COPY --link --from=nixl-builder /opt/nixl /opt/nixl
 COPY --link --from=nixl-builder /usr/lib/python3/dist-packages/nixl.pth /usr/lib/python3/dist-packages/nixl.pth