[CI/Build][Hardware][NVIDIA] Add public CUDA 13.4 Rubin build path (#54640)

Signed-off-by: Meenakshi Venkataraman <meenakshiv@nvidia.com>
Signed-off-by: meena-at-work <80416898+meena-at-work@users.noreply.github.com>
Signed-off-by: Shang Wang <shangw@nvidia.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Shang Wang <shangw@nvidia.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

docker/Dockerfile

@@ -797,6 +797,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
         cuda-nvcc-${CUDA_VERSION_DASH} \
         cuda-cudart-${CUDA_VERSION_DASH} \
         cuda-nvrtc-${CUDA_VERSION_DASH} \
+        # FlashInfer's runtime JIT-compiled kernels (e.g. trtllm_fmha_kernel_launcher.cu)
+        # #include <nvrtc.h>, which only the -dev package ships.
+        cuda-nvrtc-dev-${CUDA_VERSION_DASH} \
         cuda-cuobjdump-${CUDA_VERSION_DASH} \
         libcurand-dev-${CUDA_VERSION_DASH} \
         libcublas-dev-${CUDA_VERSION_DASH} \
@@ -806,10 +809,21 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
         numactl && \
     rm -rf /var/lib/apt/lists/*
 
+# Ensures nvcc/ptxas (installed above, needed at runtime for FlashInfer/
+# DeepGEMM/EP kernel JIT compilation) are reachable on PATH. Harmless no-op
+# if already set by FINAL_BASE_IMAGE.
+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.
+# 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 \
@@ -827,10 +841,13 @@ RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
             echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \
             exit 1; \
         fi \
-        && NCCL_LIBRARY=$(find /usr/lib -name libnccl.so.2 -print -quit) \
-        && test -n "${NCCL_LIBRARY}" \
-        && ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2 \
-        && ln -sf libnccl.so.2 /usr/lib/libnccl.so \
+        && 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
 
@@ -961,10 +978,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" ]; then \
-        # The Rubin CUDA image already provides GDRCopy. \
-        ldconfig -p | grep -q 'libgdrapi\.so'; \
+    if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ] && ldconfig -p | grep -q 'libgdrapi\.so'; then \
+        true; \
     else \
         case "${TARGETPLATFORM}" in \
           linux/arm64) UUARCH="aarch64" ;; \

docs/getting_started/installation/gpu.cuda.inc.md

@@ -425,12 +425,7 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.
 
 #### [Preview] Building vLLM's Docker Image from Source for NVIDIA Rubin GPU Architecture
 
-At this time, the Rubin preview is available only to NVIDIA engineers and
-partners with access to NVIDIA's internal prerequisites. Set
-`INSTALL_RUBIN_PRERELEASE=true` to enable the Rubin build path. It requires a
-`requirements/rubin-prerelease.txt` file containing pointers to NVIDIA's
-internal prerelease packages, plus NVIDIA-internal CUDA 13.4 or CUDA 13.5
-development images for `BUILD_BASE_IMAGE` and `FINAL_BASE_IMAGE`.
+Set `INSTALL_RUBIN_PRERELEASE=true` to enable the Rubin build path.
 
 Triton must currently be installed from source for Rubin compatibility.
 Specify its repository with `TRITON_INSTALL_FROM_SOURCE_REPO`; an empty
@@ -444,54 +439,53 @@ BuildKit does not automatically invalidate cached layers when a mutable Git
 ref changes. Use `--no-cache-filter extensions-build` to refresh an empty,
 branch, or tag revision.
 
-Use `ARCH=arm64` on VR200 or `ARCH=amd64` on R100:
+For `FINAL_BASE_IMAGE`, use the public, multi-arch
+`nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04` image.
+For `BUILD_BASE_IMAGE`, use:
 
-??? console "CUDA 13.4 command"
+- `pytorch/manylinux2_28-builder:cuda13.4` for x86_64 CPUs.
+- `pytorch/manylinuxaarch64-builder:cuda13.4` for ARM64/AArch64 CPUs.
+
+??? console "ARM64/AArch64 build command"
 
     ```bash
-    ARCH=arm64
     docker buildx build --progress=plain --load \
       --file docker/Dockerfile \
       --target vllm-openai \
-      --platform "linux/${ARCH}" \
-      --tag "vllm/vllm-rubin-openai:prerelease-cu134-${ARCH}" \
+      --platform "linux/arm64" \
+      --tag "vllm/vllm-rubin-openai:prerelease-cu134-public-arm64" \
       --build-arg max_jobs="$(nproc)" \
       --build-arg nvcc_threads=2 \
       --build-arg RUN_WHEEL_CHECK=false \
       --build-arg INSTALL_RUBIN_PRERELEASE=true \
       --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \
       --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
-      --build-arg CUDA_VERSION=13.4.1 \
-      --build-arg BUILD_BASE_IMAGE="<cuda-13.4-manylinux-devel-image-${ARCH}>" \
-      --build-arg FINAL_BASE_IMAGE="<cuda-13.4-ubuntu24.04-devel-image-${ARCH}>" \
+      --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" \
       .
     ```
 
-??? console "CUDA 13.5 command"
+??? console "x86_64 build command"
 
     ```bash
-    ARCH=arm64
     docker buildx build --progress=plain --load \
       --file docker/Dockerfile \
       --target vllm-openai \
-      --platform "linux/${ARCH}" \
-      --tag "vllm/vllm-rubin-openai:prerelease-cu135-${ARCH}" \
+      --platform "linux/amd64" \
+      --tag "vllm/vllm-rubin-openai:prerelease-cu134-public-amd64" \
       --build-arg max_jobs="$(nproc)" \
       --build-arg nvcc_threads=2 \
       --build-arg RUN_WHEEL_CHECK=false \
       --build-arg INSTALL_RUBIN_PRERELEASE=true \
       --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \
       --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
-      --build-arg CUDA_VERSION=13.5.0 \
-      --build-arg BUILD_BASE_IMAGE="<cuda-13.5-manylinux-devel-image-${ARCH}>" \
-      --build-arg FINAL_BASE_IMAGE="<cuda-13.5-ubuntu24.04-devel-image-${ARCH}>" \
+      --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" \
       .
     ```
 
-CUDA 13.5 compiled-mode qualification passed with NVIDIA driver versions
-615.62.03 and 620.05. Driver 610.47.04 failed the tested compiled CUDA 13.5
-kernel. Validate the exact driver version on the target system.
-
 !!! note
     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,

requirements/rubin-prerelease.txt

@@ -0,0 +1,16 @@
+# Public CUDA 13.4 prerelease PyTorch nightly pins for Rubin bring-up.
+# Installed when INSTALL_RUBIN_PRERELEASE=true (see docker/Dockerfile).
+# TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels
+# (vllm-project/vllm#49735).
+--pre
+--extra-index-url https://download.pytorch.org/whl/nightly/cu134
+
+torch==2.15.0.dev20260903+cu134
+torchvision==0.30.0.dev20260903+cu134
+torchaudio==2.11.0.dev20260903+cu134
+
+# docker/Dockerfile strips these two 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.6.18