From dc0b8d76b58eb8e16f16b0ecee13b06f0694c4c7 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Wed, 15 May 2024 13:58:59 +0000
Subject: [PATCH] Change the dockerfile. It builds locally, something might be
 up in AWS env.

---
 Dockerfile | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index ac5cefa8..630372b7 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -96,6 +96,7 @@ WORKDIR /usr/src
 COPY server/Makefile-flash-att Makefile
 
 # Build specific version of flash attention
+ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
 RUN make build-flash-attention
 
 # Build Flash Attention v2 CUDA kernels
@@ -107,6 +108,7 @@ COPY server/Makefile-flash-att-v2 Makefile
 
 # Build specific version of flash attention v2
 RUN make build-flash-attention-v2-cuda
+RUN TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0+PTX" make build-flash-attention
 
 # Build Transformers exllama kernels
 FROM kernel-builder as exllama-kernels-builder
@@ -181,7 +183,7 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
         ca-certificates \
         make \
         curl \
-        build-essential \
+        git \
         && rm -rf /var/lib/apt/lists/*
 
 # Copy conda with PyTorch installed