@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9650dc73e2

[CI/Build][Hardware][NVIDIA] Add opt-in Rubin Docker builds (#53443)

Signed-off-by: Shang Wang <samshang.wang@mail.utoronto.ca> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Shang Wang committed Aug 27, 2026 at 12:06 UTC 9650dc73e249688d9f4dac1f24a3f6f09e8c04a1
7 files changed +429 -24
docker/Dockerfile
+255 -22
@@ -25,6 +25,9 @@
25 ARG CUDA_VERSION=13.0.3
26 ARG PYTHON_VERSION=3.12
27 ARG UBUNTU_VERSION=24.04
28 +ARG INSTALL_RUBIN_PRERELEASE=false
29 +ARG TRITON_INSTALL_FROM_SOURCE_REPO=
30 +ARG TRITON_INSTALL_FROM_SOURCE_REVISION=
31 # DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
32 # This version is used for CUDA 12+ builds to keep the DeepEP build and
33 # runtime NCCL versions compatible.
@@ -94,6 +97,7 @@ ARG INSTALL_KV_CONNECTORS=false
97 FROM ${BUILD_BASE_IMAGE} AS base
98
99 ARG TARGETPLATFORM
100 +ARG INSTALL_RUBIN_PRERELEASE
101 ARG CUDA_VERSION
102 ARG PYTHON_VERSION
103 ARG USE_SCCACHE
@@ -205,7 +209,13 @@ COPY requirements/cuda.txt requirements/cuda.txt
209 COPY use_existing_torch.py use_existing_torch.py
210 COPY pyproject.toml pyproject.toml
211 RUN --mount=type=cache,target=/opt/uv/cache \
208 - if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
212 + --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
213 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ] && \
214 + [ ! -f /tmp/vllm-requirements/rubin-prerelease.txt ]; then \
215 + echo "INSTALL_RUBIN_PRERELEASE=true requires requirements/rubin-prerelease.txt" >&2; \
216 + exit 1; \
217 + fi \
218 + && if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
219 sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
220 sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' requirements/cuda.txt; \
221 fi \
@@ -216,7 +226,16 @@ RUN --mount=type=cache,target=/opt/uv/cache \
226 printf 'cython<3.3\n' > /tmp/build-constraints.txt \
227 && export UV_BUILD_CONSTRAINT=/tmp/build-constraints.txt; \
228 fi \
219 - && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
229 + && if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
230 + echo "Installing Rubin prerelease requirements..." \
231 + && /opt/venv/bin/python3 use_existing_torch.py --prefix \
232 + && uv pip install --python /opt/venv/bin/python3 \
233 + -r /tmp/vllm-requirements/rubin-prerelease.txt \
234 + && uv pip install --python /opt/venv/bin/python3 \
235 + -r requirements/cuda.txt \
236 + && uv pip install --python /opt/venv/bin/python3 \
237 + -r /tmp/vllm-requirements/rubin-prerelease.txt; \
238 + elif [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
239 echo "Installing torch nightly..." \
240 && uv pip install --python /opt/venv/bin/python3 torch torchaudio torchvision --pre \
241 --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
@@ -229,6 +248,33 @@ RUN --mount=type=cache,target=/opt/uv/cache \
248 --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
249 fi
250
251 +# CUDA 13.4 Rubin images need a compatibility link for the prerelease
252 +# PyTorch wheel's libcurand.so.11 dependency. CUDA 13.5 must use its native
253 +# libcurand.so.11 instead.
254 +RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
255 + CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \
256 + && if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \
257 + CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \
258 + -print -quit) \
259 + && test -n "${CURAND_LIBRARY}" \
260 + && ln -sf "${CURAND_LIBRARY}" \
261 + /usr/local/cuda/lib64/libcurand.so.11; \
262 + elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \
263 + CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \
264 + -print -quit) \
265 + && test -n "${CURAND_LIBRARY}" \
266 + && ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \
267 + else \
268 + echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
269 + exit 1; \
270 + fi \
271 + && printf '%s\n' \
272 + /opt/venv/lib \
273 + "/opt/venv/lib/python${PYTHON_VERSION}/site-packages/nvpl/lib" \
274 + > /etc/ld.so.conf.d/vllm-rubin-python.conf \
275 + && ldconfig; \
276 + fi
277 +
278 # Track PyTorch lib versions used during build and match in downstream instances.
279 # We do this for both nightly and release so we can strip dependencies/*.txt as needed.
280 # Otherwise library dependencies can upgrade/downgrade torch incorrectly.
@@ -237,9 +283,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
283 && TORCH_LIB_VERSIONS=$(cat torch_lib_versions.txt | xargs) \
284 && echo "Installed torch libs: ${TORCH_LIB_VERSIONS}"
285
240 -# CUDA arch list used by torch
241 -# Explicitly set the list to avoid issues with torch 2.2
242 -# See https://github.com/pytorch/pytorch/pull/123243
286 +# CUDA arch list used by torch. Keep this explicit so GPU-less container builds
287 +# do not depend on PyTorch's changing common-architecture fallback.
288 # From versions.json: .torch.cuda_arch_list
289 # Do not add +PTX here: vLLM filters torch's top-level PTX flag when it
290 # converts global gencode flags into per-kernel arch lists. If a specific
@@ -258,7 +303,7 @@ ARG USE_SCCACHE
303 ARG SCCACHE_ENDPOINT
304
305 # Install native tools needed only for Rust builds.
261 -RUN dnf install -y --setopt=install_weak_deps=False make \
306 +RUN dnf install -y --setopt=install_weak_deps=False make perl-core \
307 && dnf clean all \
308 && rm -rf /var/cache/dnf
309
@@ -310,6 +355,7 @@ RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
355 #################### CSRC BUILD IMAGE ####################
356 FROM base AS csrc-build
357
358 +ARG INSTALL_RUBIN_PRERELEASE
359 ARG PIP_INDEX_URL UV_INDEX_URL
360 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
361 ARG PYTORCH_CUDA_INDEX_BASE_URL
@@ -330,7 +376,11 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
376 ENV UV_LINK_MODE=copy
377
378 RUN --mount=type=cache,target=/opt/uv/cache \
333 - if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
379 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
380 + echo "Installing build requirements with Rubin prerelease PyTorch..." \
381 + && python3 use_existing_torch.py --prefix \
382 + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt; \
383 + elif [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
384 echo "Installing build requirements without torch..." \
385 && python3 use_existing_torch.py --prefix \
386 && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \
@@ -370,9 +420,10 @@ ARG VLLM_MAIN_CUDA_VERSION=""
420 # Use dummy version for csrc-build wheel (only .so files are extracted, version doesn't matter)
421 ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+csrc.build"
422
373 -# Use existing torch for nightly builds
423 +# Use existing torch for nightly and Rubin prerelease builds
424 RUN --mount=type=cache,target=/opt/uv/cache \
375 - if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
425 + if [ "${PYTORCH_NIGHTLY}" = "1" ] || \
426 + [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
427 python3 use_existing_torch.py --prefix; \
428 fi
429
@@ -424,7 +475,12 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
475 # Build DeepEP - runs in PARALLEL with csrc-build
476 # This stage is independent and doesn't affect csrc cache
477 FROM base AS extensions-build
478 +ARG TARGETARCH
479 ARG CUDA_VERSION
480 +ARG INSTALL_RUBIN_PRERELEASE
481 +ARG TRITON_INSTALL_FROM_SOURCE_REPO
482 +ARG TRITON_INSTALL_FROM_SOURCE_REVISION
483 +ARG max_jobs
484
485 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
486 ENV UV_HTTP_TIMEOUT=500
@@ -441,6 +497,20 @@ ARG NVSHMEM_VER
497 ARG NCCL_VERSION
498 RUN --mount=type=cache,target=/opt/uv/cache \
499 mkdir -p /tmp/ep_kernels_workspace/dist && \
500 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
501 + # The Rubin CUDA builder provides NCCL in its architecture-specific \
502 + # toolkit directory rather than as a Python package. Give DeepEP a \
503 + # stable prefix that also exists in the Ubuntu runtime image. \
504 + NCCL_HEADER=$(find /usr/local/cuda/targets -name nccl.h -print -quit) \
505 + && test -n "${NCCL_HEADER}" \
506 + && NCCL_LIBRARY=$(find -H /usr/local/cuda/targets \
507 + -name libnccl.so.2 -print -quit) \
508 + && test -n "${NCCL_LIBRARY}" \
509 + && ln -sf "${NCCL_HEADER}" /usr/include/nccl.h \
510 + && ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2 \
511 + && ln -sf libnccl.so.2 /usr/lib/libnccl.so \
512 + && export EP_NCCL_ROOT_DIR=/usr; \
513 + fi && \
514 CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
515 if [ -n "$NCCL_VERSION" ]; then \
516 echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
@@ -454,11 +524,56 @@ RUN --mount=type=cache,target=/opt/uv/cache \
524 ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} \
525 ${NVSHMEM_VER:+--nvshmem-ver "$NVSHMEM_VER"} && \
526 find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
527 +
528 +# Build Triton from source when a repository is configured. Triton 3.8 pins
529 +# ptxas-blackwell 13.3, which cannot assemble SM107. Rubin CUDA 13.4 and 13.5
530 +# builds use their SM107-capable system ptxas during the wheel build and at
531 +# runtime.
532 +COPY tools/build_triton_from_source.sh /tmp/build_triton_from_source.sh
533 +RUN --mount=type=cache,target=/opt/uv/cache \
534 + --mount=type=cache,id=triton-source-${TARGETARCH},target=/root/.triton,sharing=locked \
535 + --mount=type=cache,id=triton-source-ccache-${TARGETARCH},target=/root/.cache/ccache,sharing=locked \
536 + set -eux; \
537 + TRITON_WORKSPACE=/tmp/triton_from_source_workspace; \
538 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ] \
539 + && [ -z "${TRITON_INSTALL_FROM_SOURCE_REPO}" ]; then \
540 + echo "Rubin prerelease builds require TRITON_INSTALL_FROM_SOURCE_REPO" >&2; \
541 + exit 1; \
542 + elif [ -z "${TRITON_INSTALL_FROM_SOURCE_REPO}" ] \
543 + && [ -n "${TRITON_INSTALL_FROM_SOURCE_REVISION}" ]; then \
544 + echo "TRITON_INSTALL_FROM_SOURCE_REVISION requires TRITON_INSTALL_FROM_SOURCE_REPO" >&2; \
545 + exit 1; \
546 + elif [ -n "${TRITON_INSTALL_FROM_SOURCE_REPO}" ]; then \
547 + dnf install -y --setopt=install_weak_deps=False zlib-devel; \
548 + dnf clean all; \
549 + rm -rf /var/cache/dnf; \
550 + CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2); \
551 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
552 + case "${CUDA_MAJOR_MINOR}" in \
553 + 13.4|13.5) \
554 + export TRITON_PTXAS_BLACKWELL_PATH=/usr/local/cuda/bin/ptxas \
555 + ;; \
556 + *) \
557 + echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
558 + exit 1 \
559 + ;; \
560 + esac; \
561 + fi; \
562 + TRITON_WORKSPACE="${TRITON_WORKSPACE}" \
563 + MAX_JOBS="${max_jobs:-}" \
564 + bash /tmp/build_triton_from_source.sh; \
565 + else \
566 + mkdir -p "${TRITON_WORKSPACE}/dist"; \
567 + printf '%s\n' disabled > "${TRITON_WORKSPACE}/dist/build-state"; \
568 + (cd "${TRITON_WORKSPACE}/dist" \
569 + && sha256sum build-state > wheel.sha256); \
570 + fi
571 #################### EXTENSIONS BUILD IMAGE ####################
572
573 #################### WHEEL BUILD IMAGE ####################
574 FROM base AS build
575 ARG TARGETPLATFORM
576 +ARG INSTALL_RUBIN_PRERELEASE
577
578 ARG PIP_INDEX_URL UV_INDEX_URL
579 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -480,7 +595,11 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
595 ENV UV_LINK_MODE=copy
596
597 RUN --mount=type=cache,target=/opt/uv/cache \
483 - if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
598 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
599 + echo "Installing build requirements with Rubin prerelease PyTorch..." \
600 + && python3 use_existing_torch.py --prefix \
601 + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt; \
602 + elif [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
603 echo "Installing build requirements without torch..." \
604 && python3 use_existing_torch.py --prefix \
605 && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \
@@ -514,10 +633,15 @@ ENV VLLM_TARGET_DEVICE=${vllm_target_device}
633 # Skip adding +precompiled suffix to version (preserves git-derived version)
634 ENV VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX=1
635
517 -# Use existing torch for nightly builds
636 +# Use existing torch for nightly and Rubin prerelease builds
637 RUN --mount=type=cache,target=/opt/uv/cache \
519 - if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
638 + if [ "${PYTORCH_NIGHTLY}" = "1" ] || \
639 + [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
640 python3 use_existing_torch.py --prefix; \
641 + fi \
642 + && if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
643 + sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\])[[:space:]]*==/Id' \
644 + requirements/cuda.txt; \
645 fi
646
647 # Build the vLLM wheel
@@ -607,6 +731,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
731 # image with vLLM installed
732 FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base
733
734 +ARG INSTALL_RUBIN_PRERELEASE
735 ARG CUDA_VERSION
736 ARG PYTHON_VERSION
737 ARG NCCL_VERSION
@@ -614,6 +739,11 @@ ARG DEADSNAKES_MIRROR_URL
739 ARG DEADSNAKES_GPGKEY_URL
740 ARG GET_PIP_URL
741
742 +# Keep the output image's CUDA metadata aligned with the selected toolkit.
743 +# This also normalizes prerelease base images with four-component metadata for
744 +# OCI runtimes that evaluate NVIDIA's CUDA driver requirement.
745 +ENV CUDA_VERSION=${CUDA_VERSION}
746 +
747 ENV DEBIAN_FRONTEND=noninteractive
748 WORKDIR /vllm-workspace
749
@@ -676,6 +806,34 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
806 numactl && \
807 rm -rf /var/lib/apt/lists/*
808
809 +# Apply the CUDA 13.4-only libcurand compatibility link described in the build
810 +# stage. CUDA 13.5 uses its native libcurand.so.11. DeepEP expects NCCL directly
811 +# under its configured /usr prefix, while Ubuntu installs it in a multiarch
812 +# subdirectory.
813 +RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
814 + CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \
815 + && if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \
816 + CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \
817 + -print -quit) \
818 + && test -n "${CURAND_LIBRARY}" \
819 + && ln -sf "${CURAND_LIBRARY}" \
820 + /usr/local/cuda/lib64/libcurand.so.11; \
821 + elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \
822 + CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \
823 + -print -quit) \
824 + && test -n "${CURAND_LIBRARY}" \
825 + && ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \
826 + else \
827 + echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
828 + exit 1; \
829 + fi \
830 + && NCCL_LIBRARY=$(find /usr/lib -name libnccl.so.2 -print -quit) \
831 + && test -n "${NCCL_LIBRARY}" \
832 + && ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2 \
833 + && ln -sf libnccl.so.2 /usr/lib/libnccl.so \
834 + && ldconfig; \
835 + fi
836 +
837 # Install uv for faster pip installs
838 RUN python3 -m pip install uv
839
@@ -753,12 +911,33 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
911 COPY requirements/common.txt /tmp/common.txt
912 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
913 RUN --mount=type=cache,target=/opt/uv/cache \
914 + --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
915 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ] && \
916 + [ ! -f /tmp/vllm-requirements/rubin-prerelease.txt ]; then \
917 + echo "INSTALL_RUBIN_PRERELEASE=true requires requirements/rubin-prerelease.txt" >&2; \
918 + exit 1; \
919 + fi && \
920 if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
921 sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
922 sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' /tmp/requirements-cuda.txt; \
923 fi && \
760 - uv pip install --system -r /tmp/requirements-cuda.txt \
761 - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
924 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
925 + uv pip install --system \
926 + -r /tmp/vllm-requirements/rubin-prerelease.txt \
927 + && sed -Ei '/^[[:space:]]*(torch|torchvision|torchaudio)[[:space:]]*==/Id' \
928 + /tmp/requirements-cuda.txt \
929 + && uv pip install --system -r /tmp/requirements-cuda.txt \
930 + && uv pip install --system \
931 + -r /tmp/vllm-requirements/rubin-prerelease.txt \
932 + && printf '%s\n' \
933 + "/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvpl/lib" \
934 + > /etc/ld.so.conf.d/vllm-rubin-python.conf \
935 + && ldconfig; \
936 + else \
937 + uv pip install --system \
938 + -r /tmp/requirements-cuda.txt \
939 + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
940 + fi && \
941 rm /tmp/requirements-cuda.txt /tmp/common.txt
942
943 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
@@ -766,8 +945,10 @@ RUN --mount=type=cache,target=/opt/uv/cache \
945 # From versions.json: .flashinfer.version
946 ARG FLASHINFER_VERSION=0.6.17
947 RUN --mount=type=cache,target=/opt/uv/cache \
769 - uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
770 - --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
948 + if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \
949 + uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
950 + --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
951 + fi
952
953 # ============================================================
954 # OPENAI API SERVER DEPENDENCIES
@@ -781,12 +962,18 @@ ARG GDRCOPY_OS_VERSION=Ubuntu24_04
962 ARG TARGETPLATFORM
963 COPY tools/install_gdrcopy.sh /tmp/install_gdrcopy.sh
964 RUN set -eux; \
784 - case "${TARGETPLATFORM}" in \
785 - linux/arm64) UUARCH="aarch64" ;; \
786 - linux/amd64) UUARCH="x64" ;; \
787 - *) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
788 - esac; \
789 - /tmp/install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}" && \
965 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
966 + # The Rubin CUDA image already provides GDRCopy. \
967 + ldconfig -p | grep -q 'libgdrapi\.so'; \
968 + else \
969 + case "${TARGETPLATFORM}" in \
970 + linux/arm64) UUARCH="aarch64" ;; \
971 + linux/amd64) UUARCH="x64" ;; \
972 + *) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
973 + esac; \
974 + /tmp/install_gdrcopy.sh \
975 + "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"; \
976 + fi; \
977 rm /tmp/install_gdrcopy.sh
978
979 # Install vllm-openai dependencies (saves ~2.6s per build)
@@ -994,6 +1181,8 @@ RUN mv vllm src/vllm
1181 FROM vllm-base AS vllm-openai-base
1182 ARG TARGETPLATFORM
1183 ARG INSTALL_KV_CONNECTORS=false
1184 +ARG INSTALL_RUBIN_PRERELEASE
1185 +ARG TRITON_INSTALL_FROM_SOURCE_REPO
1186 ARG CUDA_VERSION
1187 ARG VLLM_BUILD_COMMIT
1188 ARG VLLM_BUILD_PIPELINE
@@ -1045,6 +1234,50 @@ RUN --mount=type=cache,target=/opt/uv/cache \
1234 # time; see docs/features/mooncake_connector_usage.md. Deliberately not an ENV
1235 # here so the image does not override mooncake's own default.
1236
1237 +# Install Rubin prerelease dependencies when explicitly requested. Keep this
1238 +# after the public vLLM and KV connector dependency pins installed above.
1239 +RUN --mount=type=cache,target=/opt/uv/cache \
1240 + --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
1241 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
1242 + echo "Installing Rubin prerelease requirements..." \
1243 + && uv pip install --system \
1244 + -r /tmp/vllm-requirements/rubin-prerelease.txt; \
1245 + fi
1246 +
1247 +# Install source-built Triton last so dependency resolution cannot restore the
1248 +# previously installed version.
1249 +COPY --from=extensions-build \
1250 + /tmp/triton_from_source_workspace/dist/wheel.sha256 \
1251 + /tmp/triton-source-wheel.sha256
1252 +COPY tools/install_triton_from_source.sh /tmp/install_triton_from_source.sh
1253 +RUN --mount=type=cache,target=/opt/uv/cache \
1254 + --mount=type=bind,from=extensions-build,source=/tmp/triton_from_source_workspace/dist,target=/tmp/triton-source-dist,ro \
1255 + if [ -n "${TRITON_INSTALL_FROM_SOURCE_REPO}" ]; then \
1256 + CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2); \
1257 + if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
1258 + case "${CUDA_MAJOR_MINOR}" in \
1259 + 13.4|13.5) \
1260 + TRITON_SYSTEM_PTXAS_PATH=/usr/local/cuda/bin/ptxas \
1261 + TRITON_EXPECTED_PTXAS_VERSION="${CUDA_MAJOR_MINOR}" \
1262 + bash /tmp/install_triton_from_source.sh \
1263 + /tmp/triton-source-dist \
1264 + /tmp/triton-source-wheel.sha256 \
1265 + ;; \
1266 + *) \
1267 + echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
1268 + exit 1 \
1269 + ;; \
1270 + esac; \
1271 + else \
1272 + bash /tmp/install_triton_from_source.sh \
1273 + /tmp/triton-source-dist \
1274 + /tmp/triton-source-wheel.sha256; \
1275 + fi; \
1276 + fi \
1277 + && rm -f \
1278 + /tmp/install_triton_from_source.sh \
1279 + /tmp/triton-source-wheel.sha256
1280 +
1281 ENV VLLM_USAGE_SOURCE production-docker-image
1282 ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
1283 VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
docker/versions.json
+3
@@ -10,6 +10,9 @@
10 "UBUNTU_VERSION": {
11 "default": "24.04"
12 },
13 + "INSTALL_RUBIN_PRERELEASE": {
14 + "default": "false"
15 + },
16 "NCCL_VERSION": {
17 "default": "2.30.7"
18 },
docs/assets/contributing/dockerfile-stages-dependency.png
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ
docs/contributing/dockerfile/dockerfile.md
+4
@@ -20,6 +20,10 @@ The edges of the build graph represent:
20 The `test-deps` stage branches from `vllm-runtime-base` so Git and the test
21 requirements remain cached independently of per-commit vLLM wheels and source.
22
23 +The `extensions-build` stage can also produce an optional source-built Triton
24 +wheel. `vllm-openai-base` installs that wheel after its other Python
25 +dependencies so dependency resolution cannot restore an older Triton version.
26 +
27 > <figure markdown="span">
28 > ![](../../assets/contributing/dockerfile-stages-dependency.png){ align="center" alt="query" width="100%" }
29 > </figure>
docs/getting_started/installation/gpu.cuda.inc.md
+83 -2
@@ -354,8 +354,9 @@ DOCKER_BUILDKIT=1 docker build . \
354
355 !!! note
356 By default vLLM will build for all GPU types for widest distribution. If you are just building for the
357 - current GPU type the machine is running on, you can add the argument `--build-arg torch_cuda_arch_list=""`
358 - for vLLM to find the current GPU type and build for that.
357 + current GPU type, you can add `--build-arg torch_cuda_arch_list=""` to delegate architecture selection
358 + to PyTorch. This requires the GPU to be visible to the container build; standard Docker BuildKit builds
359 + do not expose it and PyTorch instead falls back to its common architecture list.
360
361 If you are using Podman instead of Docker, you might need to disable SELinux labeling by
362 adding `--security-opt label=disable` when running `podman build` command to avoid certain [existing issues](https://github.com/containers/buildah/discussions/4184).
@@ -422,6 +423,86 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.
423
424 After setting up QEMU, you can use the `--platform "linux/arm64"` flag in your `docker build` command.
425
426 +#### [Preview] Building vLLM's Docker Image from Source for NVIDIA Rubin GPU Architecture
427 +
428 +At this time, the Rubin preview is available only to NVIDIA engineers and
429 +partners with access to NVIDIA's internal prerequisites. Set
430 +`INSTALL_RUBIN_PRERELEASE=true` to enable the Rubin build path. It requires a
431 +`requirements/rubin-prerelease.txt` file containing pointers to NVIDIA's
432 +internal prerelease packages, plus NVIDIA-internal CUDA 13.4 or CUDA 13.5
433 +development images for `BUILD_BASE_IMAGE` and `FINAL_BASE_IMAGE`.
434 +
435 +Triton must currently be installed from source for Rubin compatibility.
436 +Specify its repository with `TRITON_INSTALL_FROM_SOURCE_REPO`; an empty
437 +`TRITON_INSTALL_FROM_SOURCE_REVISION` selects the repository's latest `main`,
438 +while a commit, branch, or tag selects that revision. The tested revision
439 +lowers SM107 through LLVM's SM100 target and uses the final CUDA image's
440 +version-matched `ptxas` for SM107 assembly. This enables vLLM's default compiled
441 +mode on VR200 and R100.
442 +
443 +BuildKit does not automatically invalidate cached layers when a mutable Git
444 +ref changes. Use `--no-cache-filter extensions-build` to refresh an empty,
445 +branch, or tag revision.
446 +
447 +Use `ARCH=arm64` on VR200 or `ARCH=amd64` on R100:
448 +
449 +??? console "CUDA 13.4 command"
450 +
451 + ```bash
452 + ARCH=arm64
453 + docker buildx build --progress=plain --load \
454 + --file docker/Dockerfile \
455 + --target vllm-openai \
456 + --platform "linux/${ARCH}" \
457 + --tag "vllm/vllm-rubin-openai:prerelease-cu134-${ARCH}" \
458 + --build-arg max_jobs="$(nproc)" \
459 + --build-arg nvcc_threads=2 \
460 + --build-arg RUN_WHEEL_CHECK=false \
461 + --build-arg INSTALL_RUBIN_PRERELEASE=true \
462 + --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \
463 + --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
464 + --build-arg CUDA_VERSION=13.4.1 \
465 + --build-arg BUILD_BASE_IMAGE="<cuda-13.4-manylinux-devel-image-${ARCH}>" \
466 + --build-arg FINAL_BASE_IMAGE="<cuda-13.4-ubuntu24.04-devel-image-${ARCH}>" \
467 + .
468 + ```
469 +
470 +??? console "CUDA 13.5 command"
471 +
472 + ```bash
473 + ARCH=arm64
474 + docker buildx build --progress=plain --load \
475 + --file docker/Dockerfile \
476 + --target vllm-openai \
477 + --platform "linux/${ARCH}" \
478 + --tag "vllm/vllm-rubin-openai:prerelease-cu135-${ARCH}" \
479 + --build-arg max_jobs="$(nproc)" \
480 + --build-arg nvcc_threads=2 \
481 + --build-arg RUN_WHEEL_CHECK=false \
482 + --build-arg INSTALL_RUBIN_PRERELEASE=true \
483 + --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \
484 + --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
485 + --build-arg CUDA_VERSION=13.5.0 \
486 + --build-arg BUILD_BASE_IMAGE="<cuda-13.5-manylinux-devel-image-${ARCH}>" \
487 + --build-arg FINAL_BASE_IMAGE="<cuda-13.5-ubuntu24.04-devel-image-${ARCH}>" \
488 + .
489 + ```
490 +
491 +CUDA 13.5 compiled-mode qualification passed with NVIDIA driver versions
492 +615.62.03 and 620.05. Driver 610.47.04 failed the tested compiled CUDA 13.5
493 +kernel. Validate the exact driver version on the target system.
494 +
495 +!!! note
496 + Keep the default explicit `torch_cuda_arch_list`. GPU-less BuildKit builds
497 + cannot inspect the host GPU. R100 and VR200 report compute capability 10.7,
498 + for which the generic `10.0` target provides family-compatible kernels.
499 + The Ubuntu `devel` final image is also required: the corresponding `base`
500 + image lacks the CUDA runtime/JIT package closure used by vLLM and the
501 + prerelease PyTorch wheel.
502 +
503 + `RUN_WHEEL_CHECK=false` disables only the PyPI publication-size guard for
504 + this private staging image.
505 +
506 #### Use the custom-built vLLM Docker image**
507
508 To run vLLM with the custom-built Docker image:
tools/build_triton_from_source.sh new
+43
@@ -0,0 +1,43 @@
1 +#!/usr/bin/env bash
2 +set -euxo pipefail
3 +
4 +: "${TARGETARCH:?TARGETARCH must be set}"
5 +: "${TRITON_INSTALL_FROM_SOURCE_REPO:?TRITON_INSTALL_FROM_SOURCE_REPO must be set}"
6 +
7 +TRITON_REVISION="${TRITON_INSTALL_FROM_SOURCE_REVISION:-main}"
8 +TRITON_WORKSPACE="${TRITON_WORKSPACE:-/tmp/triton_from_source_workspace}"
9 +TRITON_SOURCE="${TRITON_WORKSPACE}/src"
10 +TRITON_DIST="${TRITON_WORKSPACE}/dist"
11 +TRITON_BUILD_JOBS="${MAX_JOBS:-$(nproc)}"
12 +
13 +case "${TARGETARCH}" in
14 + amd64) TRITON_LLVM_SYSTEM_SUFFIX=almalinux-x64 ;;
15 + arm64) TRITON_LLVM_SYSTEM_SUFFIX=almalinux-arm64 ;;
16 + *)
17 + echo "Unsupported Triton build architecture: ${TARGETARCH}" >&2
18 + exit 1
19 + ;;
20 +esac
21 +export TRITON_LLVM_SYSTEM_SUFFIX
22 +
23 +mkdir -p "${TRITON_DIST}"
24 +git init -q "${TRITON_SOURCE}"
25 +git -C "${TRITON_SOURCE}" remote add origin \
26 + "${TRITON_INSTALL_FROM_SOURCE_REPO}"
27 +git -C "${TRITON_SOURCE}" fetch --depth=1 origin "${TRITON_REVISION}"
28 +git -C "${TRITON_SOURCE}" checkout -q --detach FETCH_HEAD
29 +git -C "${TRITON_SOURCE}" rev-parse HEAD
30 +
31 +uv pip install --python /opt/venv/bin/python3 \
32 + -r "${TRITON_SOURCE}/python/requirements.txt"
33 +
34 +TRITON_BUILD_WITH_CCACHE=true \
35 +MAX_JOBS="${TRITON_BUILD_JOBS}" \
36 + uv build --python /opt/venv/bin/python3 \
37 + --wheel \
38 + --no-build-isolation \
39 + --out-dir "${TRITON_DIST}" \
40 + "${TRITON_SOURCE}"
41 +
42 +(cd "${TRITON_DIST}" && sha256sum triton-*.whl > wheel.sha256)
43 +rm -rf "${TRITON_SOURCE}"
tools/install_triton_from_source.sh new
+41
@@ -0,0 +1,41 @@
1 +#!/usr/bin/env bash
2 +set -euxo pipefail
3 +
4 +if [[ $# -ne 2 ]]; then
5 + echo "Usage: $0 <wheel-directory> <checksum-file>" >&2
6 + exit 1
7 +fi
8 +
9 +TRITON_DIST="$1"
10 +TRITON_CHECKSUM="$2"
11 +
12 +test -d "${TRITON_DIST}"
13 +test -f "${TRITON_CHECKSUM}"
14 +(cd "${TRITON_DIST}" && sha256sum -c "${TRITON_CHECKSUM}")
15 +
16 +shopt -s nullglob
17 +TRITON_WHEELS=("${TRITON_DIST}"/triton-*.whl)
18 +if [[ ${#TRITON_WHEELS[@]} -ne 1 ]]; then
19 + echo "Expected exactly one Triton wheel in ${TRITON_DIST}" >&2
20 + exit 1
21 +fi
22 +
23 +uv pip install --system --reinstall --no-deps "${TRITON_WHEELS[0]}"
24 +
25 +if [[ -n "${TRITON_SYSTEM_PTXAS_PATH:-}" \
26 + || -n "${TRITON_EXPECTED_PTXAS_VERSION:-}" ]]; then
27 + : "${TRITON_SYSTEM_PTXAS_PATH:?TRITON_SYSTEM_PTXAS_PATH must be set}"
28 + : "${TRITON_EXPECTED_PTXAS_VERSION:?TRITON_EXPECTED_PTXAS_VERSION must be set}"
29 + PTXAS_PATH="${TRITON_SYSTEM_PTXAS_PATH}"
30 + test -x "${PTXAS_PATH}"
31 + TRITON_NVIDIA_BIN=$(python3 -c \
32 + 'from pathlib import Path; import triton; print(Path(triton.__file__).parent / "backends" / "nvidia" / "bin")')
33 + mkdir -p "${TRITON_NVIDIA_BIN}"
34 + ln -sfn "${PTXAS_PATH}" \
35 + "${TRITON_NVIDIA_BIN}/ptxas-blackwell"
36 +
37 + env -u TRITON_PTXAS_BLACKWELL_PATH \
38 + TRITON_EXPECTED_PTXAS_PATH="${PTXAS_PATH}" \
39 + python3 -c \
40 + 'import os; from triton.backends.nvidia.compiler import get_ptxas; tool = get_ptxas(107); assert tool.version == os.environ["TRITON_EXPECTED_PTXAS_VERSION"]; assert os.path.samefile(tool.path, os.environ["TRITON_EXPECTED_PTXAS_PATH"])'
41 +fi