From c52f08351fa6dcc00dd09ced48166304c35fec44 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= <angt@huggingface.co>
Date: Wed, 5 Feb 2025 10:57:50 +0000
Subject: [PATCH] Set TGI_LLAMA_PKG_CUDA from CUDA_VERSION
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Adrien Gallouët <angt@huggingface.co>
---
 Dockerfile_llamacpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Dockerfile_llamacpp b/Dockerfile_llamacpp
index ed8783d6..b020778f 100644
--- a/Dockerfile_llamacpp
+++ b/Dockerfile_llamacpp
@@ -2,6 +2,7 @@ FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu24.04 AS deps
 
 ARG llama_version=b4628
 ARG llama_cuda_arch=75-real;80-real;86-real;89-real;90-real
+ENV TGI_LLAMA_PKG_CUDA=cuda-${CUDA_VERSION%.*}
 
 WORKDIR /opt/src
 
@@ -50,7 +51,6 @@ RUN cargo chef cook \
     --profile release-opt \
     --package text-generation-router-llamacpp
 COPY . .
-ENV TGI_LLAMA_PKG_CUDA=cuda-12.6
 RUN cargo build \
     --profile release-opt \
     --package text-generation-router-llamacpp --frozen