feat: faster CPU image on AMD (#35)

Author: OlivierDehaene

Cargo.lock

@@ -18,10 +18,10 @@ homepage = "https://github.com/huggingface/text-embeddings-inference"
 
 [patch.crates-io]
 cudarc = { git = "https://github.com/OlivierDehaene/cudarc", rev = "4c8e6d36a4a4c31e2e4649ae5246226452a01fc1" }
-candle = { git = "https://github.com/huggingface/candle", rev = "122da875806f274a7aa9048f76d7a676b473e56f", package = "candle-core" }
-candle-nn = { git = "https://github.com/huggingface/candle", rev = "122da875806f274a7aa9048f76d7a676b473e56f", package = "candle-nn" }
-candle-transformers = { git = "https://github.com/huggingface/candle", rev = "122da875806f274a7aa9048f76d7a676b473e56f", package = "candle-transformers" }
-candle-flash-attn = { git = "https://github.com/huggingface/candle", rev = "122da875806f274a7aa9048f76d7a676b473e56f", package = "candle-flash-attn" }
+candle = { git = "https://github.com/OlivierDehaene/candle", rev = "9f2b4081b83a0e47ec1b12caa71d3cac7cc2161e", package = "candle-core" }
+candle-nn = { git = "https://github.com/OlivierDehaene/candle", rev = "9f2b4081b83a0e47ec1b12caa71d3cac7cc2161e", package = "candle-nn" }
+candle-transformers = { git = "https://github.com/OlivierDehaene/candle", rev = "9f2b4081b83a0e47ec1b12caa71d3cac7cc2161e", package = "candle-transformers" }
+candle-flash-attn = { git = "https://github.com/OlivierDehaene/candle", rev = "9f2b4081b83a0e47ec1b12caa71d3cac7cc2161e", package = "candle-flash-attn" }
 
 [profile.release]
 debug = 1

Cargo.toml

@@ -35,31 +35,53 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
 
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     intel-oneapi-mkl-devel \
+    build-essential \
     && rm -rf /var/lib/apt/lists/*
 
+RUN echo "int mkl_serv_intel_cpu_true() {return 1;}" > fakeintel.c && \
+    gcc -shared -fPIC -o libfakeintel.so fakeintel.c
+
 COPY --from=planner /usr/src/recipe.json recipe.json
 
-RUN cargo chef cook --release --features candle --features mkl --no-default-features --recipe-path recipe.json && sccache -s
+RUN cargo chef cook --release --features candle --features mkl-dynamic --no-default-features --recipe-path recipe.json && sccache -s
 
 COPY backends backends
 COPY core core
 COPY router router
 COPY Cargo.toml ./
 COPY Cargo.lock ./
 
-RUN cargo build --release --bin text-embeddings-router -F candle -F mkl --no-default-features && sccache -s
+RUN cargo build --release --bin text-embeddings-router -F candle -F mkl-dynamic --no-default-features && sccache -s
 
 FROM debian:bookworm-slim
 
 ENV HUGGINGFACE_HUB_CACHE=/data \
-    PORT=80
+    PORT=80 \
+    MKL_ENABLE_INSTRUCTIONS=AVX512_E4 \
+    RAYON_NUM_THREADS=8 \
+    LD_PRELOAD=/usr/local/libfakeintel.so \
+    LD_LIBRARY_PATH=/usr/local/lib
 
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     libomp-dev \
     ca-certificates \
     libssl-dev \
     && rm -rf /var/lib/apt/lists/*
 
+# Copy a lot of the Intel shared objects because of the mkl_serv_intel_cpu_true patch...
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_intel_lp64.so.2 /usr/local/lib/libmkl_intel_lp64.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_intel_thread.so.2 /usr/local/lib/libmkl_intel_thread.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_core.so.2 /usr/local/lib/libmkl_core.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_vml_def.so.2 /usr/local/lib/libmkl_vml_def.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_def.so.2 /usr/local/lib/libmkl_def.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_vml_avx.so.2 /usr/local/lib/libmkl_vml_avx.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_vml_avx2.so.2 /usr/local/lib/libmkl_vml_avx2.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_vml_avx512.so.2 /usr/local/lib/libmkl_vml_avx512.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_avx.so.2 /usr/local/lib/libmkl_avx.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_avx2.so.2 /usr/local/lib/libmkl_avx2.so.2
+COPY --from=builder /opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_avx512.so.2 /usr/local/lib/libmkl_avx512.so.2
+COPY --from=builder /usr/src/libfakeintel.so /usr/local/libfakeintel.so
+
 COPY --from=builder /usr/src/target/release/text-embeddings-router /usr/local/bin/text-embeddings-router
 
 ENTRYPOINT ["text-embeddings-router"]

Dockerfile

@@ -19,6 +19,7 @@ clap = ["dep:clap", "text-embeddings-backend-core/clap"]
 python = ["dep:text-embeddings-backend-python"]
 candle = ["dep:text-embeddings-backend-candle"]
 mkl = ["text-embeddings-backend-candle?/mkl"]
+mkl-dynamic = ["text-embeddings-backend-candle?/mkl-dynamic"]
 accelerate = ["text-embeddings-backend-candle?/accelerate"]
 static-linking = ["text-embeddings-backend-candle?/static-linking"]
 flash-attn = ["text-embeddings-backend-candle?/flash-attn"]

backends/Cargo.toml

@@ -7,7 +7,7 @@ homepage.workspace = true
 
 [dependencies]
 accelerate-src = { version = "0.3.2", optional = true }
-intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"], optional = true }
+intel-mkl-src = { version = "0.8.1", optional = true  }
 candle = { version = "0.3.0", package = "candle-core", default-features = false }
 candle-nn = { version = "0.3.0" }
 candle-transformers = { version = "0.3.0" }
@@ -29,7 +29,8 @@ anyhow = { version = "1", features = ["backtrace"] }
 
 [features]
 accelerate = ["dep:accelerate-src", "candle/accelerate", "candle-nn/accelerate"]
-mkl = ["dep:intel-mkl-src", "candle/mkl", "candle-nn/mkl"]
+mkl = ["dep:intel-mkl-src", "intel-mkl-src/mkl-static-lp64-iomp", "candle/mkl", "candle-nn/mkl"]
+mkl-dynamic = ["dep:intel-mkl-src", "intel-mkl-src/mkl-dynamic-lp64-iomp", "candle/mkl-dynamic", "candle-nn/mkl-dynamic"]
 cuda = ["candle/cuda", "candle-nn/cuda", "dep:candle-cublaslt", "dep:candle-layer-norm"]
 flash-attn-v1 = ["dep:candle-flash-attn-v1", "cuda"]
 flash-attn = ["dep:candle-flash-attn", "cuda"]

backends/candle/Cargo.toml

@@ -15,7 +15,6 @@ use text_embeddings_backend_core::{BackendError, Batch, Embedding, EmbeddingBack
 
 pub struct CandleBackend {
     model: Box<dyn EmbeddingModel + Send>,
-    device: Device,
 }
 
 impl CandleBackend {
@@ -126,7 +125,7 @@ impl CandleBackend {
             }
         };
 
-        Ok(Self { model, device })
+        Ok(Self { model })
     }
 }
 
@@ -142,10 +141,7 @@ impl EmbeddingBackend for CandleBackend {
     }
 
     fn max_batch_size(&self) -> Option<usize> {
-        match self.device {
-            Device::Cpu => Some(1),
-            Device::Cuda(_) => None,
-        }
+        None
     }
 }
 

backends/candle/src/lib.rs

@@ -1,4 +1,4 @@
-#[cfg(feature = "mkl")]
+#[cfg(any(feature = "mkl", feature = "mkl-dynamic"))]
 extern crate intel_mkl_src;
 
 #[cfg(feature = "accelerate")]