From 8085e8509aa885c5acf2bdd2d057359e18ed547e Mon Sep 17 00:00:00 2001
From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
Date: Tue, 13 Oct 2026 19:59:56 -0700
Subject: [PATCH 1/3] [https://nvbugs/6919405][fix] Keep the 26.08 CUDA
 forward-compat driver on the 26.09 image

The NGC PyTorch 26.09 base ships CUDA forward-compatibility driver 615.71.09
(26.08 shipped 615.65.02). Under it the unchanged trtllm-gen context FMHA cubin
runs 7-13% slower per call on GB200, which costs
aggr-ctx_only-gb200_deepseek-r1-fp4_128k8k_con1_ctx1_pp8_gen1_tep8_eplb0_mtp3_ccb-NIXL
~6% Total Token Throughput since #19679.

Add a CUDA_COMPAT_TAG build arg (default 26.08-py3) and replace the base
image's /usr/local/cuda-13.4/compat driver libraries with that tag's during the
devel stage. The CUDA toolkit is untouched, and the base image's shell init still
validates and selects the compat driver at container start. Targets built on
nvcr.io/nvidia/cuda keep their own compat driver.

Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
---
 docker/Dockerfile.multi              | 15 +++++++
 docker/Makefile                      |  6 +++
 docker/common/install_cuda_compat.sh | 60 ++++++++++++++++++++++++++++
 3 files changed, 81 insertions(+)
 create mode 100644 docker/common/install_cuda_compat.sh

diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi
index c6712f3c8b4e..4e893717ac9d 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-<ver>/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 <source /usr/local>
+#
+# <source /usr/local> is the /usr/local tree of the image that provides the compat driver. Only
+# the driver libraries under /usr/local/cuda-<major.minor>/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 <source /usr/local>}"
+
+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 `.<driver>.<host>.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.*

From aaaa8a4b85359a3b5bded36d7e6013662ed520c3 Mon Sep 17 00:00:00 2001
From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
Date: Wed, 14 Oct 2026 17:35:48 -0700
Subject: [PATCH 2/3] [https://nvbugs/6919405][infra] Run CI on the images
 built with the 26.08 compat driver

Point jenkins/current_image_tags.properties at the CI images that L0
build 64146 built from this branch (commit 33d6b6f), so pre-merge tests
run on the pinned compat driver. These are staging tags; they will be
promoted to tensorrt-llm tags with scripts/rename_docker_images.py
before merge.

Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
---
 jenkins/current_image_tags.properties | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties
index 428f962e04f1..05a3ef82dbda 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-staging/tritondevel:x86_64-tritondevel-torch_skip-33d6b6f-github-pr-20252-64146
+LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:sbsa-tritondevel-torch_skip-33d6b6f-github-pr-20252-64146
+LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py310-33d6b6f-github-pr-20252-64146
+LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py312-33d6b6f-github-pr-20252-64146
+LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:sbsa-ubuntu24-torch_skip-py312-33d6b6f-github-pr-20252-64146

From 2db87501413b4adcbf2835927d1e8c3e8e2a26a0 Mon Sep 17 00:00:00 2001
From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
Date: Thu, 15 Oct 2026 02:50:45 -0700
Subject: [PATCH 3/3] [https://nvbugs/6919405][infra] Use the retagged CI
 images built with the 26.08 compat driver

Replace the tensorrt-llm-staging tags in jenkins/current_image_tags.properties
with tensorrt-llm tags. These images are the ones L0 build 64146 built from
this branch's Dockerfile change (commit 33d6b6f), retagged with
scripts/rename_docker_images.py 64146 20252 --commit 33d6b6f run against
main's copy of the tags file.

Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
---
 jenkins/current_image_tags.properties | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties
index 05a3ef82dbda..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-staging/tritondevel:x86_64-tritondevel-torch_skip-33d6b6f-github-pr-20252-64146
-LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:sbsa-tritondevel-torch_skip-33d6b6f-github-pr-20252-64146
-LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py310-33d6b6f-github-pr-20252-64146
-LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py312-33d6b6f-github-pr-20252-64146
-LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm-staging/tritondevel:sbsa-ubuntu24-torch_skip-py312-33d6b6f-github-pr-20252-64146
+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