[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:
Shang Wang
2026-10-03 00:12:33 +00:00
committed by GitHub
co-authored by Claude Opus 5.5 Tyler Michael Smith Shengqi Chen
parent 8bdd8d847d
commit 8e443871f0
5 changed files with 24 additions and 133 deletions
+4 -2
View File
@@ -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
View File
@@ -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.
+1 -1
View File
@@ -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
+4 -15
View File
@@ -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