[Build] Download CUDA 12.9 nvcc/ptxas/cicc for better codegen in Docker build

Use CUDA 12.9.86 compiler toolchain (nvcc, ptxas, cicc) for better GPU code generation while keeping CUDA 12.8 headers/libs for runtime compatibility with PyTorch cu128 wheels. nvcc is replaced (not just ptxas/cicc) so that setup.py's get_nvcc_cuda_version() reports 12.9, which is required for FlashMLA registration. PyTorch's CMake nvcc-vs-header version check is patched to a warning in base (for extensions-build), csrc-build, and build stages (re-applied after torch pip reinstall overwrites cmake files). Controlled by ARG MIN_NVCC_TOOLCHAIN_VERSION (default 12.9.86): - Skipped if system nvcc is already >= the specified version - Set to empty string to disable entirely - Archive integrity verified via SHA256 checksums from NVIDIA's redistrib manifest Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
[Build] Downgrade default CUDA version to 12.8 for PyTorch cu128 compatibility
2026-02-18 19:15:17 -05:00 · 2026-02-18 17:59:30 -05:00
3 changed files with 78 additions and 22 deletions
@@ -31,7 +31,7 @@ steps:
        commands:
          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -70,7 +70,7 @@ steps:
        agents:
          queue: cpu_queue_postmerge
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -123,7 +123,7 @@ steps:
          queue: cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
@@ -137,7 +137,7 @@ steps:
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"

      - block: "Build release image for x86_64 CPU"
@@ -23,11 +23,6 @@
 # =============================================================================

 ARG CUDA_VERSION=12.8.1
-# BUILDER_CUDA_VERSION controls the CUDA toolkit used to compile csrc/ and
-# extensions (DeepGEMM, EP kernels).  It can differ from CUDA_VERSION, which
-# is the CUDA version shipped in the final runtime image and used to select
-# the matching PyTorch wheel.
-ARG BUILDER_CUDA_VERSION=12.9.1
 ARG PYTHON_VERSION=3.12

 # By parameterizing the base images, we allow third-party to use their own
@@ -41,7 +36,7 @@ ARG PYTHON_VERSION=3.12
 # compatibility with other Linux OSes. The main reason for this is that the
 # glibc version is baked into the distro, and binaries built with one glibc
 # version are not backwards compatible with OSes that use an earlier version.
-ARG BUILD_BASE_IMAGE=nvidia/cuda:${BUILDER_CUDA_VERSION}-devel-ubuntu20.04
+ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
 # Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels)
 ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04

@@ -97,13 +92,8 @@ ARG INSTALL_KV_CONNECTORS=false
 FROM ${BUILD_BASE_IMAGE} AS base

 ARG CUDA_VERSION
-ARG BUILDER_CUDA_VERSION
 ARG PYTHON_VERSION

-# Override the CUDA_VERSION env var inherited from the nvidia base image
-# (which equals BUILDER_CUDA_VERSION) so that $CUDA_VERSION in RUN commands
-# resolves to the runtime CUDA version used for PyTorch wheel selection.
-ENV CUDA_VERSION=${CUDA_VERSION}
 ENV DEBIAN_FRONTEND=noninteractive

 # Install system dependencies including build tools
@@ -143,7 +133,7 @@ ENV UV_LINK_MODE=copy
 RUN gcc --version

 # Ensure CUDA compatibility library is loaded
-RUN echo "/usr/local/cuda-$(echo "$BUILDER_CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig
+RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig

 # ============================================================
 # SLOW-CHANGING DEPENDENCIES BELOW
@@ -193,6 +183,50 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # From versions.json: .torch.cuda_arch_list
 ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
+
+# Download newer CUDA compiler toolchain for better GPU code generation
+# while keeping CUDA 12.8 headers/libs for runtime compatibility with PyTorch cu128.
+# We replace nvcc, ptxas, and cicc so that setup.py's get_nvcc_cuda_version() reports
+# the newer version (required for FlashMLA registration which gates on >= 12.9).
+# PyTorch's CMake checks nvcc version vs CUDA header version; that check is patched
+# to a warning in the csrc-build and build stages (after torch is pip-installed).
+# Skipped if system nvcc is already >= this version (e.g. when building with CUDA 13+).
+# Set MIN_NVCC_TOOLCHAIN_VERSION="" to disable.
+ARG MIN_NVCC_TOOLCHAIN_VERSION=12.9.86
+# SHA256 checksums from https://developer.download.nvidia.com/compute/cuda/redist/redistrib_<cuda_version>.json
+ARG NVCC_TOOLCHAIN_SHA256_X86_64=7a1a5b652e5ef85c82b721d10672fc9a2dbaab44e9bd3c65a69517bf53998c35
+ARG NVCC_TOOLCHAIN_SHA256_AARCH64=2432ef8a7c12d0a9ce3332a8af42b123c07f256390b3390802b1b2c6254c6c74
+RUN if [ -n "${MIN_NVCC_TOOLCHAIN_VERSION}" ]; then \
+        CURRENT_VERSION=$(nvcc --version | sed -n 's/.*V\([0-9.]*\).*/\1/p') && \
+        SMALLEST=$(printf '%s\n' "${MIN_NVCC_TOOLCHAIN_VERSION}" "${CURRENT_VERSION}" | sort -V | head -n1) && \
+        if [ "${SMALLEST}" = "${MIN_NVCC_TOOLCHAIN_VERSION}" ]; then \
+            echo "System nvcc ${CURRENT_VERSION} >= ${MIN_NVCC_TOOLCHAIN_VERSION}, skipping toolchain download"; \
+        else \
+            ARCH=$(uname -m) && \
+            case "${ARCH}" in \
+                x86_64)  EXPECTED_SHA256="${NVCC_TOOLCHAIN_SHA256_X86_64}" ;; \
+                aarch64) EXPECTED_SHA256="${NVCC_TOOLCHAIN_SHA256_AARCH64}" ;; \
+                *) echo "Unsupported architecture: ${ARCH}" >&2; exit 1 ;; \
+            esac && \
+            curl -sL -o cuda_nvcc.tar.xz \
+                "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-${ARCH}/cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive.tar.xz" && \
+            echo "${EXPECTED_SHA256}  cuda_nvcc.tar.xz" | sha256sum -c && \
+            tar xJf cuda_nvcc.tar.xz && \
+            cp -f cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive/bin/nvcc /usr/local/cuda/bin/nvcc && \
+            cp -f cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive/bin/ptxas /usr/local/cuda/bin/ptxas && \
+            cp -f cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive/nvvm/bin/cicc /usr/local/cuda/nvvm/bin/cicc && \
+            rm -rf cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive cuda_nvcc.tar.xz && \
+            echo "Upgraded nvcc/ptxas/cicc from ${CURRENT_VERSION} to ${MIN_NVCC_TOOLCHAIN_VERSION}"; \
+        fi; \
+    fi
+
+# Patch PyTorch cmake to allow nvcc/header version mismatch (for newer nvcc toolchain).
+# Stages that reinstall torch (csrc-build, build) must re-apply this patch.
+RUN python3 -c "\
+import pathlib; \
+files = list(pathlib.Path('/opt/venv').rglob('Caffe2/public/cuda.cmake')); \
+assert files, 'ERROR: Caffe2/public/cuda.cmake not found under /opt/venv'; \
+[print(f'Patching {f}') or f.write_text(f.read_text().replace('FATAL_ERROR \"FindCUDA says', 'WARNING \"FindCUDA says')) for f in files]"
 #################### BUILD BASE IMAGE ####################

 #################### CSRC BUILD IMAGE ####################
@@ -232,6 +266,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
    fi

+# Patch PyTorch cmake to allow nvcc/header version mismatch (for newer nvcc toolchain).
+# This must run after pip install since torch reinstall overwrites the cmake files.
+RUN python3 -c "\
+import pathlib; \
+files = list(pathlib.Path('/opt/venv').rglob('Caffe2/public/cuda.cmake')); \
+assert files, 'ERROR: Caffe2/public/cuda.cmake not found under /opt/venv'; \
+[print(f'Patching {f}') or f.write_text(f.read_text().replace('FATAL_ERROR \"FindCUDA says', 'WARNING \"FindCUDA says')) for f in files]"
+
 WORKDIR /workspace

 COPY pyproject.toml setup.py CMakeLists.txt ./
@@ -319,7 +361,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 # Build DeepGEMM, pplx-kernels, DeepEP - runs in PARALLEL with csrc-build
 # This stage is independent and doesn't affect csrc cache
 FROM base AS extensions-build
-ARG BUILDER_CUDA_VERSION
+ARG CUDA_VERSION

 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 ENV UV_HTTP_TIMEOUT=500
@@ -335,7 +377,7 @@ COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
 RUN --mount=type=cache,target=/root/.cache/uv \
    mkdir -p /tmp/deepgemm/dist && \
    VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh \
-        --cuda-version "${BUILDER_CUDA_VERSION}" \
+        --cuda-version "${CUDA_VERSION}" \
        ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} \
        --wheel-dir /tmp/deepgemm/dist || \
    echo "DeepGEMM build skipped (CUDA version requirement not met)"
@@ -398,6 +440,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
    fi

+# Patch PyTorch cmake to allow nvcc/header version mismatch (for newer nvcc toolchain).
+# This must run after pip install since torch reinstall overwrites the cmake files.
+RUN python3 -c "\
+import pathlib; \
+files = list(pathlib.Path('/opt/venv').rglob('Caffe2/public/cuda.cmake')); \
+assert files, 'ERROR: Caffe2/public/cuda.cmake not found under /opt/venv'; \
+[print(f'Patching {f}') or f.write_text(f.read_text().replace('FATAL_ERROR \"FindCUDA says', 'WARNING \"FindCUDA says')) for f in files]"
+
 WORKDIR /workspace

 # Copy pre-built csrc wheel directly
@@ -4,14 +4,11 @@
    "CUDA_VERSION": {
      "default": "12.8.1"
    },
-    "BUILDER_CUDA_VERSION": {
-      "default": "12.9.1"
-    },
    "PYTHON_VERSION": {
      "default": "3.12"
    },
    "BUILD_BASE_IMAGE": {
-      "default": "nvidia/cuda:12.9.1-devel-ubuntu20.04"
+      "default": "nvidia/cuda:12.8.1-devel-ubuntu20.04"
    },
    "FINAL_BASE_IMAGE": {
      "default": "nvidia/cuda:12.8.1-base-ubuntu22.04"
@@ -34,6 +31,15 @@
    "TORCH_CUDA_ARCH_LIST": {
      "default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
    },
+    "MIN_NVCC_TOOLCHAIN_VERSION": {
+      "default": "12.9.86"
+    },
+    "NVCC_TOOLCHAIN_SHA256_X86_64": {
+      "default": "7a1a5b652e5ef85c82b721d10672fc9a2dbaab44e9bd3c65a69517bf53998c35"
+    },
+    "NVCC_TOOLCHAIN_SHA256_AARCH64": {
+      "default": "2432ef8a7c12d0a9ce3332a8af42b123c07f256390b3390802b1b2c6254c6c74"
+    },
    "MAX_JOBS": {
      "default": "2"
    },