Adding intel openmp tunings do Docker file

This patch adds more tunins for CPU backend on intel openmp. These tunings improves CPU backend performance greatly, especially on throughput related tests. Signed-off-by: Yuan Zhou <[email protected]>
vllm-project · Jul 1, 2024 · 30ee3ce · 30ee3ce
1 parent 80ca1e6
commit 30ee3ce
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/Dockerfile.cpu b/Dockerfile.cpu
@@ -6,7 +6,16 @@ RUN apt-get update  -y \
     && apt-get install -y git wget vim numactl gcc-12 g++-12 python3 python3-pip libtcmalloc-minimal4 \
     && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12
 
-RUN echo 'export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4:$LD_PRELOAD' >> ~/.bashrc
+RUN pip install accelerate mkl
+
+ENV LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4:/usr/local/lib/libiomp5.so:$LD_PRELOAD"
+
+ENV KMP_BLOCKTIME=1
+ENV KMP_TPAUSE=0
+ENV KMP_SETTINGS=1
+ENV KMP_FORKJOIN_BARRIER_PATTERN=dist,dist
+ENV KMP_PLAIN_BARRIER_PATTERN=dist,dist
+ENV KMP_REDUCTION_BARRIER_PATTERN=dist,dist
 
 RUN pip install https://intel-extension-for-pytorch.s3.amazonaws.com/ipex_dev/cpu/intel_extension_for_pytorch-2.3.100%2Bgit0eb3473-cp310-cp310-linux_x86_64.whl