diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 3753015b5ebb..fe238cb1518f 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -18,6 +18,13 @@ ARG TRITON_IMAGE=nvcr.io/nvidia/tritonserver ARG BASE_TAG=26.09-py3 ARG TRITON_BASE_TAG=26.09-py3 ARG DEVEL_IMAGE=devel +# Tag of ${BASE_IMAGE} whose CUDA forward-compatibility driver (/usr/local/cuda-/compat) +# replaces the base image's own; it must ship the same CUDA toolkit version as BASE_TAG. The +# 26.09-py3 compat driver (615.71.09) runs the trtllm-gen context FMHA kernels ~10% slower on +# GB200 than the 26.08-py3 one (615.65.02). Set it equal to BASE_TAG to keep the base image's own. +ARG CUDA_COMPAT_TAG=26.08-py3 + +FROM ${BASE_IMAGE}:${CUDA_COMPAT_TAG} AS cuda_compat FROM ${BASE_IMAGE}:${BASE_TAG} AS base @@ -72,6 +79,14 @@ RUN --mount=type=bind,source=docker/common,target=/opt/docker/common \ bash /opt/docker/common/install.sh --base --cmake --ccache --cuda_toolkit \ --cuda_libs --polygraphy --mpi4py --pytorch +ARG BASE_TAG +ARG CUDA_COMPAT_TAG +RUN --mount=type=bind,source=docker/common,target=/opt/docker/common \ + --mount=type=bind,from=cuda_compat,source=/usr/local,target=/tmp/cuda_compat_usr_local \ + if [ "${CUDA_COMPAT_TAG}" != "${BASE_TAG}" ]; then \ + bash /opt/docker/common/install_cuda_compat.sh /tmp/cuda_compat_usr_local; \ + fi + # Install constraints after install.sh so cleanup() doesn't delete the file mid-RUN COPY constraints.txt /tmp/constraints.txt RUN --mount=type=cache,target=/root/.cache/pip \ diff --git a/docker/Makefile b/docker/Makefile index dcc11fc291d7..dba8258d87b9 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -3,6 +3,7 @@ BASE_IMAGE ?= $(shell grep '^ARG BASE_IMAGE=' Dockerfile.multi | grep -o BASE_TAG ?= $(shell grep '^ARG BASE_TAG=' Dockerfile.multi | grep -o '=.*' | tr -d '="') TRITON_IMAGE ?= $(shell grep '^ARG TRITON_IMAGE=' Dockerfile.multi | grep -o '=.*' | tr -d '="') TRITON_BASE_TAG ?= $(shell grep '^ARG TRITON_BASE_TAG=' Dockerfile.multi | grep -o '=.*' | tr -d '="') +CUDA_COMPAT_TAG ?= $(shell grep '^ARG CUDA_COMPAT_TAG=' Dockerfile.multi | grep -o '=.*' | tr -d '="') # Name of the new image IMAGE_NAME ?= tensorrt_llm IMAGE_TAG ?= latest @@ -99,6 +100,7 @@ base_pull: $(if $(BASE_TAG), --build-arg BASE_TAG=$(BASE_TAG)) \ $(if $(TRITON_IMAGE), --build-arg TRITON_IMAGE=$(TRITON_IMAGE)) \ $(if $(TRITON_BASE_TAG), --build-arg TRITON_BASE_TAG=$(TRITON_BASE_TAG)) \ + $(if $(CUDA_COMPAT_TAG), --build-arg CUDA_COMPAT_TAG=$(CUDA_COMPAT_TAG)) \ $(if $(BUILD_WHEEL_ARGS), --build-arg BUILD_WHEEL_ARGS="$(BUILD_WHEEL_ARGS)") \ $(if $(BUILD_WHEEL_SCRIPT), --build-arg BUILD_WHEEL_SCRIPT="$(BUILD_WHEEL_SCRIPT)") \ $(if $(TORCH_INSTALL_TYPE), --build-arg TORCH_INSTALL_TYPE="$(TORCH_INSTALL_TYPE)") \ @@ -207,20 +209,24 @@ jenkins-rockylinux8_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags. jenkins-rockylinux8_%: STAGE = tritondevel jenkins-rockylinux8_%: BASE_IMAGE = nvcr.io/nvidia/cuda jenkins-rockylinux8_%: BASE_TAG = 13.4.1-devel-rockylinux8 +jenkins-rockylinux8_%: CUDA_COMPAT_TAG = $(BASE_TAG) rockylinux8_%: STAGE = tritondevel rockylinux8_%: BASE_IMAGE = nvcr.io/nvidia/cuda rockylinux8_%: BASE_TAG = 13.4.1-devel-rockylinux8 +rockylinux8_%: CUDA_COMPAT_TAG = $(BASE_TAG) # For x86_64 ubuntu22_%: STAGE = tritondevel ubuntu22_%: BASE_IMAGE = nvcr.io/nvidia/cuda ubuntu22_%: BASE_TAG = 13.4.1-devel-ubuntu22.04 +ubuntu22_%: CUDA_COMPAT_TAG = $(BASE_TAG) # For x86_64 and aarch64 ubuntu24_%: STAGE = tritondevel ubuntu24_%: BASE_IMAGE = nvcr.io/nvidia/cuda ubuntu24_%: BASE_TAG = 13.4.1-devel-ubuntu24.04 +ubuntu24_%: CUDA_COMPAT_TAG = $(BASE_TAG) trtllm_%: STAGE = release trtllm_%: PUSH_TO_STAGING := 0 diff --git a/docker/common/install_cuda_compat.sh b/docker/common/install_cuda_compat.sh new file mode 100644 index 000000000000..c9e72cd73aac --- /dev/null +++ b/docker/common/install_cuda_compat.sh @@ -0,0 +1,60 @@ +#!/bin/bash + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Replace this image's CUDA forward-compatibility driver with the one from another image. +# +# Usage: install_cuda_compat.sh +# +# is the /usr/local tree of the image that provides the compat driver. Only +# the driver libraries under /usr/local/cuda-/compat (lib.real, and lib.real-nvgpu +# where present) are replaced; the CUDA toolkit is untouched, so the source image must ship the +# same CUDA major.minor as this one. The base image's shell init (/etc/shinit_v2) checks and +# selects the compat driver from that directory at container start, so nothing else changes. + +set -euo pipefail + +SRC_USR_LOCAL="${1:?usage: install_cuda_compat.sh }" + +CUDA_DIR="$(readlink -f /usr/local/cuda)" +CUDA_DIR_NAME="$(basename "${CUDA_DIR}")" +DST_COMPAT="${CUDA_DIR}/compat" +SRC_COMPAT="${SRC_USR_LOCAL}/${CUDA_DIR_NAME}/compat" + +if ! ls "${SRC_COMPAT}"/lib.real/libcuda.so.* >/dev/null 2>&1; then + echo "No ${CUDA_DIR_NAME}/compat/lib.real/libcuda.so.* in the compat source image; it must" \ + "ship the same CUDA version as this image (${CUDA_DIR_NAME})." >&2 + exit 1 +fi +if ! ls "${DST_COMPAT}"/lib.real/libcuda.so.* >/dev/null 2>&1; then + echo "This image has no ${DST_COMPAT}/lib.real/libcuda.so.*; nothing to replace." >&2 + exit 1 +fi + +echo "Replacing CUDA compat driver $(basename "$(ls "${DST_COMPAT}"/lib.real/libcuda.so.*.* | head -1)")" \ + "with $(basename "$(ls "${SRC_COMPAT}"/lib.real/libcuda.so.*.* | head -1)"):" +for entry in lib.real lib.real-nvgpu; do + rm -rf "${DST_COMPAT:?}/${entry}" + if [ -d "${SRC_COMPAT}/${entry}" ]; then + cp -a "${SRC_COMPAT}/${entry}" "${DST_COMPAT}/${entry}" + fi +done +# The shell init skips its compat check on hosts with a `...checked` marker; drop +# markers recorded against the replaced driver so the check runs against this one. +find "${DST_COMPAT}" -maxdepth 1 -name '.*.checked' -delete + +ldconfig +test -e "${DST_COMPAT}/lib.real/libcuda.so.1" +ls -l "${DST_COMPAT}"/lib.real/libcuda.so.* diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index 428f962e04f1..bbf9b59e08ab 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -13,8 +13,8 @@ # images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead. IMAGE_NAME=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm -LLM_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.09-py3-x86_64-ubuntu24.04-skip-tritondevel-202610130652-20175 -LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.09-py3-sbsa-ubuntu24.04-skip-tritondevel-202610130652-20175 -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py310-skip-tritondevel-202610130652-20175 -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py312-skip-tritondevel-202610130652-20175 -LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-skip-tritondevel-202610130652-20175 +LLM_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.09-py3-x86_64-ubuntu24.04-skip-tritondevel-202610150933-20252 +LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.09-py3-sbsa-ubuntu24.04-skip-tritondevel-202610150933-20252 +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py310-skip-tritondevel-202610150933-20252 +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py312-skip-tritondevel-202610150933-20252 +LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-skip-tritondevel-202610150933-20252