From 71e700a6ea192269d229ea464292a32a981a2541 Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <morgan@huggingface.co>
Date: Mon, 2 Dec 2024 00:06:24 +0100
Subject: [PATCH] feat(backend): use latest trtllm main version to have g++ >=
 13 compatibility

---
 Dockerfile_trtllm                  | 16 ++++++++++------
 backends/trtllm/cmake/trtllm.cmake |  5 +++--
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/Dockerfile_trtllm b/Dockerfile_trtllm
index 3ccb0310..7e1ffd61 100644
--- a/Dockerfile_trtllm
+++ b/Dockerfile_trtllm
@@ -1,5 +1,5 @@
 ARG CUDA_ARCH_LIST="75-real;80-real;86-real;89-real;90-real"
-ARG OMPI_VERSION="4.1.6"
+ARG OMPI_VERSION="4.1.7rc1"
 
 # Build dependencies resolver stage
 FROM lukemathwalker/cargo-chef:latest AS chef
@@ -10,7 +10,7 @@ COPY . .
 RUN cargo chef prepare --recipe-path recipe.json
 
 # CUDA dependent dependencies resolver stage
-FROM nvidia/cuda:12.6.1-cudnn-devel-ubuntu22.04 AS cuda-builder
+FROM nvidia/cuda:12.6.1-cudnn-devel-ubuntu24.04 AS cuda-builder
 
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
@@ -23,13 +23,16 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     git \
     git-lfs \
     libssl-dev \
+    libucx-dev \
     ninja-build \
     pkg-config \
+    pipx \
     python3 \
     python3-dev \
     python3-setuptools \
     tar \
-    wget
+    wget && \
+    pipx ensurepath
 
 ENV TGI_INSTALL_PREFIX=/usr/local/tgi
 ENV TENSORRT_INSTALL_PREFIX=/usr/local/tensorrt
@@ -83,10 +86,11 @@ RUN mkdir $TGI_INSTALL_PREFIX && mkdir "$TGI_INSTALL_PREFIX/include" && mkdir "$
     cd backends/trtllm && \
     CMAKE_INSTALL_PREFIX=$TGI_INSTALL_PREFIX cargo build --release
 
-FROM nvidia/cuda:12.6.1-cudnn-runtime-ubuntu22.04 AS runtime
-RUN apt update && apt install -y python3-minimal python3-dev python3-pip && \
+FROM nvidia/cuda:12.6.1-cudnn-runtime-ubuntu24.04 AS runtime
+RUN apt update && apt install -y libucx0 pipx python3-minimal python3-dev python3-pip python3-venv && \
     rm -rf /var/lib/{apt,dpkg,cache,log}/ && \
-    python3 -m pip install transformers tokenizers
+    pipx ensurepath && \
+    pipx install --include-deps transformers tokenizers
 
 WORKDIR /usr/local/tgi/bin
 
diff --git a/backends/trtllm/cmake/trtllm.cmake b/backends/trtllm/cmake/trtllm.cmake
index 5f1b6c19..91b7753a 100644
--- a/backends/trtllm/cmake/trtllm.cmake
+++ b/backends/trtllm/cmake/trtllm.cmake
@@ -11,6 +11,7 @@ set(CMAKE_CUDA_ARCHITECTURES ${TGI_TRTLLM_BACKEND_TARGET_CUDA_ARCH_LIST})
 
 message(STATUS "Building for CUDA Architectures: ${CMAKE_CUDA_ARCHITECTURES}")
 
+set(ENABLE_UCX OFF)
 if (${CMAKE_BUILD_TYPE} STREQUAL "Debug")
     set(FAST_BUILD ON)
     set(NVTX_DISABLE OFF)
@@ -23,8 +24,8 @@ endif ()
 fetchcontent_declare(
         trtllm
         GIT_REPOSITORY https://github.com/NVIDIA/TensorRT-LLM.git
-        GIT_TAG 201135e58aa525af7e523d091d4c9584229524bc
-        GIT_SHALLOW FALSE
+        GIT_TAG c994b697313b8d36527ba346a0c5e8e374616c95
+        GIT_SHALLOW ON
         DOWNLOAD_EXTRACT_TIMESTAMP
 )
 fetchcontent_makeavailable(trtllm)