[CI/Build][NVIDIA] Build Rubin images on the public nvidia/cuda base image (#59288)
Signed-off-by: Shang Wang <samshang.wang@mail.utoronto.ca> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
This commit is contained in:
co-authored by
Claude Opus 5.5
Tyler Michael Smith
Shengqi Chen
parent
8bdd8d847d
commit
8e443871f0
@@ -339,7 +339,8 @@ steps:
|
||||
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
|
||||
--build-arg RUN_WHEEL_CHECK=false \
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.4 \
|
||||
--build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \
|
||||
--build-arg FINAL_BASE_IMAGE=nvidia/cuda:13.4.1-base-ubuntu24.04 \
|
||||
--build-arg NCCL_VERSION=2.32.3 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
@@ -369,7 +370,8 @@ steps:
|
||||
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
|
||||
--build-arg RUN_WHEEL_CHECK=false \
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.4 \
|
||||
--build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \
|
||||
--build-arg FINAL_BASE_IMAGE=nvidia/cuda:13.4.1-base-ubuntu24.04 \
|
||||
--build-arg NCCL_VERSION=2.32.3 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
|
||||
+9
-109
@@ -251,33 +251,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
fi
|
||||
|
||||
# CUDA 13.4 Rubin images need a compatibility link for the prerelease
|
||||
# PyTorch wheel's libcurand.so.11 dependency. CUDA 13.5 must use its native
|
||||
# libcurand.so.11 instead.
|
||||
RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
|
||||
CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \
|
||||
&& if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \
|
||||
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \
|
||||
-print -quit) \
|
||||
&& test -n "${CURAND_LIBRARY}" \
|
||||
&& ln -sf "${CURAND_LIBRARY}" \
|
||||
/usr/local/cuda/lib64/libcurand.so.11; \
|
||||
elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \
|
||||
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \
|
||||
-print -quit) \
|
||||
&& test -n "${CURAND_LIBRARY}" \
|
||||
&& ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \
|
||||
else \
|
||||
echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
|
||||
exit 1; \
|
||||
fi \
|
||||
&& printf '%s\n' \
|
||||
/opt/venv/lib \
|
||||
"/opt/venv/lib/python${PYTHON_VERSION}/site-packages/nvpl/lib" \
|
||||
> /etc/ld.so.conf.d/vllm-rubin-python.conf \
|
||||
&& ldconfig; \
|
||||
fi
|
||||
|
||||
# Track PyTorch lib versions used during build and match in downstream instances.
|
||||
# We do this for both nightly and release so we can strip dependencies/*.txt as needed.
|
||||
# Otherwise library dependencies can upgrade/downgrade torch incorrectly.
|
||||
@@ -518,20 +491,6 @@ ARG NVSHMEM_VER
|
||||
ARG NCCL_VERSION
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
mkdir -p /tmp/ep_kernels_workspace/dist && \
|
||||
if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
|
||||
# The Rubin CUDA builder provides NCCL in its architecture-specific \
|
||||
# toolkit directory rather than as a Python package. Give DeepEP a \
|
||||
# stable prefix that also exists in the Ubuntu runtime image. \
|
||||
NCCL_HEADER=$(find /usr/local/cuda/targets -name nccl.h -print -quit) \
|
||||
&& test -n "${NCCL_HEADER}" \
|
||||
&& NCCL_LIBRARY=$(find -H /usr/local/cuda/targets \
|
||||
-name libnccl.so.2 -print -quit) \
|
||||
&& test -n "${NCCL_LIBRARY}" \
|
||||
&& ln -sf "${NCCL_HEADER}" /usr/include/nccl.h \
|
||||
&& ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2 \
|
||||
&& ln -sf libnccl.so.2 /usr/lib/libnccl.so \
|
||||
&& export EP_NCCL_ROOT_DIR=/usr; \
|
||||
fi && \
|
||||
CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
|
||||
if [ -n "$NCCL_VERSION" ]; then \
|
||||
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
|
||||
@@ -659,10 +618,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "${PYTORCH_NIGHTLY}" = "1" ] || \
|
||||
[ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
|
||||
python3 tools/use_existing_torch.py --prefix; \
|
||||
fi \
|
||||
&& if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
|
||||
sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \
|
||||
requirements/cuda.txt; \
|
||||
fi
|
||||
|
||||
# Build the vLLM wheel
|
||||
@@ -764,11 +719,6 @@ ARG DEADSNAKES_MIRROR_URL
|
||||
ARG DEADSNAKES_GPGKEY_URL
|
||||
ARG GET_PIP_URL
|
||||
|
||||
# Keep the output image's CUDA metadata aligned with the selected toolkit.
|
||||
# This also normalizes prerelease base images with four-component metadata for
|
||||
# OCI runtimes that evaluate NVIDIA's CUDA driver requirement.
|
||||
ENV CUDA_VERSION=${CUDA_VERSION}
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
WORKDIR /vllm-workspace
|
||||
|
||||
@@ -841,41 +791,6 @@ ENV CUDA_HOME=/usr/local/cuda
|
||||
ENV PATH="${CUDA_HOME}/bin:${PATH}"
|
||||
ENV LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}"
|
||||
|
||||
# Apply the CUDA 13.4-only libcurand compatibility link described in the build
|
||||
# stage. CUDA 13.5 uses its native libcurand.so.11. DeepEP expects NCCL directly
|
||||
# under its configured /usr prefix, while Ubuntu installs it in a multiarch
|
||||
# subdirectory - but that symlink is only possible when NCCL is already
|
||||
# present (e.g. nvcr.io/nvidia/cuda-dl-base bundles it; a bare-CUDA image
|
||||
# doesn't, and gets it later via the standard nvidia-nccl-cu* pip package,
|
||||
# UV_OVERRIDE below). Check reality (whether libnccl.so.2 is actually found)
|
||||
# and tolerate not finding it instead of failing the build.
|
||||
RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
|
||||
CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \
|
||||
&& if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \
|
||||
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \
|
||||
-print -quit) \
|
||||
&& test -n "${CURAND_LIBRARY}" \
|
||||
&& ln -sf "${CURAND_LIBRARY}" \
|
||||
/usr/local/cuda/lib64/libcurand.so.11; \
|
||||
elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \
|
||||
CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \
|
||||
-print -quit) \
|
||||
&& test -n "${CURAND_LIBRARY}" \
|
||||
&& ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \
|
||||
else \
|
||||
echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
|
||||
exit 1; \
|
||||
fi \
|
||||
&& if NCCL_LIBRARY=$(find /usr/lib -name libnccl.so.2 -print -quit) \
|
||||
&& [ -n "${NCCL_LIBRARY}" ]; then \
|
||||
if [ "${NCCL_LIBRARY}" != "/usr/lib/libnccl.so.2" ]; then \
|
||||
ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2; \
|
||||
fi \
|
||||
&& ln -sf libnccl.so.2 /usr/lib/libnccl.so; \
|
||||
fi \
|
||||
&& ldconfig; \
|
||||
fi
|
||||
|
||||
# Install the initialized-snapshot runtime on the architecture supported by
|
||||
# `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned
|
||||
# source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's
|
||||
@@ -990,11 +905,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
&& uv pip install --system -r /tmp/requirements-cuda.txt \
|
||||
&& uv pip install --system \
|
||||
--prerelease=allow --index-strategy unsafe-first-match \
|
||||
-r /tmp/vllm-requirements/rubin-prerelease.txt \
|
||||
&& printf '%s\n' \
|
||||
"/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvpl/lib" \
|
||||
> /etc/ld.so.conf.d/vllm-rubin-python.conf \
|
||||
&& ldconfig; \
|
||||
-r /tmp/vllm-requirements/rubin-prerelease.txt; \
|
||||
else \
|
||||
uv pip install --system \
|
||||
-r /tmp/requirements-cuda.txt \
|
||||
@@ -1007,10 +918,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# From versions.json: .flashinfer.version
|
||||
ARG FLASHINFER_VERSION=0.7.0.post1
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
fi
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
# ============================================================
|
||||
# OPENAI API SERVER DEPENDENCIES
|
||||
@@ -1023,22 +932,13 @@ ARG GDRCOPY_CUDA_VERSION=12.8
|
||||
ARG GDRCOPY_OS_VERSION=Ubuntu24_04
|
||||
ARG TARGETPLATFORM
|
||||
COPY tools/install_gdrcopy.sh /tmp/install_gdrcopy.sh
|
||||
# Some Rubin FINAL_BASE_IMAGE choices (internal Rubin CUDA images,
|
||||
# nvcr.io/nvidia/cuda-dl-base) already provide GDRCopy; others don't, so it
|
||||
# falls through to the normal install below. Check reality (ldconfig) rather
|
||||
# than inferring from which FINAL_BASE_IMAGE flag is set.
|
||||
RUN set -eux; \
|
||||
if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ] && ldconfig -p | grep -q 'libgdrapi\.so'; then \
|
||||
true; \
|
||||
else \
|
||||
case "${TARGETPLATFORM}" in \
|
||||
linux/arm64) UUARCH="aarch64" ;; \
|
||||
linux/amd64) UUARCH="x64" ;; \
|
||||
*) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
|
||||
esac; \
|
||||
/tmp/install_gdrcopy.sh \
|
||||
"${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"; \
|
||||
fi; \
|
||||
case "${TARGETPLATFORM}" in \
|
||||
linux/arm64) UUARCH="aarch64" ;; \
|
||||
linux/amd64) UUARCH="x64" ;; \
|
||||
*) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
|
||||
esac; \
|
||||
/tmp/install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}" && \
|
||||
rm /tmp/install_gdrcopy.sh
|
||||
|
||||
# Install vllm-openai dependencies (saves ~2.6s per build)
|
||||
|
||||
@@ -440,7 +440,8 @@ ref changes. Use `--no-cache-filter extensions-build` to refresh an empty,
|
||||
branch, or tag revision.
|
||||
|
||||
For `FINAL_BASE_IMAGE`, use the public, multi-arch
|
||||
`nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04` image.
|
||||
`nvidia/cuda:13.4.1-base-ubuntu24.04` image. Set `NCCL_VERSION` to 2.32.3 or
|
||||
newer, the first NCCL release with Rubin (SM107) support.
|
||||
For `BUILD_BASE_IMAGE`, use:
|
||||
|
||||
- `pytorch/manylinux2_28-builder:cuda13.4` for x86_64 CPUs.
|
||||
@@ -462,7 +463,8 @@ For `BUILD_BASE_IMAGE`, use:
|
||||
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
|
||||
--build-arg CUDA_VERSION=13.4 \
|
||||
--build-arg BUILD_BASE_IMAGE="pytorch/manylinuxaarch64-builder:cuda13.4" \
|
||||
--build-arg FINAL_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04" \
|
||||
--build-arg FINAL_BASE_IMAGE="nvidia/cuda:13.4.1-base-ubuntu24.04" \
|
||||
--build-arg NCCL_VERSION=2.32.3 \
|
||||
.
|
||||
```
|
||||
|
||||
@@ -482,7 +484,8 @@ For `BUILD_BASE_IMAGE`, use:
|
||||
--build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
|
||||
--build-arg CUDA_VERSION=13.4 \
|
||||
--build-arg BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.4" \
|
||||
--build-arg FINAL_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04" \
|
||||
--build-arg FINAL_BASE_IMAGE="nvidia/cuda:13.4.1-base-ubuntu24.04" \
|
||||
--build-arg NCCL_VERSION=2.32.3 \
|
||||
.
|
||||
```
|
||||
|
||||
@@ -490,9 +493,6 @@ For `BUILD_BASE_IMAGE`, use:
|
||||
Keep the default explicit `torch_cuda_arch_list`. GPU-less BuildKit builds
|
||||
cannot inspect the host GPU. R100 and VR200 report compute capability 10.7,
|
||||
for which the generic `10.0` target provides family-compatible kernels.
|
||||
The Ubuntu `devel` final image is also required: the corresponding `base`
|
||||
image lacks the CUDA runtime/JIT package closure used by vLLM and the
|
||||
prerelease PyTorch wheel.
|
||||
|
||||
`RUN_WHEEL_CHECK=false` disables only the PyPI publication-size guard for
|
||||
this private staging image.
|
||||
|
||||
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.3
|
||||
instanttensor >= 0.1.9
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl[cu13]==4.7.1
|
||||
nvidia-cutlass-dsl[cu13]==4.8.0
|
||||
quack-kernels==0.6.5 # Required by MSA
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
|
||||
@@ -3,24 +3,13 @@
|
||||
# TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels
|
||||
# (vllm-project/vllm#49735).
|
||||
--extra-index-url https://download.pytorch.org/whl/nightly/cu134
|
||||
--extra-index-url https://flashinfer.ai/whl/cu134
|
||||
--extra-index-url https://flashinfer.ai/whl
|
||||
|
||||
# Match each architecture's torch pin to the latest torchvision wheel.
|
||||
torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64"
|
||||
torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64"
|
||||
torchvision==0.30.0.dev20260909+cu134
|
||||
torchaudio==2.11.0.dev20260911+cu134
|
||||
|
||||
# docker/Dockerfile strips these packages from requirements/cuda.txt whenever
|
||||
# INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to
|
||||
# still get installed.
|
||||
nvidia-cutlass-dsl[cu13]==4.8.0.dev0
|
||||
flashinfer-python==0.7.0.post1
|
||||
quack-kernels==0.6.5
|
||||
torch==2.15.0.dev20260928+cu134; platform_machine == "x86_64"
|
||||
torch==2.15.0.dev20260926+cu134; platform_machine == "aarch64"
|
||||
torchvision==0.30.0.dev20260928+cu134
|
||||
torchaudio==2.11.0.dev20260928+cu134
|
||||
|
||||
cuda-python==13.4.1
|
||||
cuda-bindings==13.4.1
|
||||
flashinfer-jit-cache==0.7.0.post1+cu134
|
||||
flashinfer-cubin==0.7.0.post1
|
||||
nixl==1.4.1
|
||||
|
||||
Reference in New Issue
Block a user