Skip to content

Commit 289a4fb

Browse files
Abdennacer-BadaouiAbdennacer-Badaoui
authored
Merge branch 'main' into amd-minimal-support
2 parents 100afcf + c14ee47 commit 289a4fb

9 files changed

Lines changed: 108 additions & 38 deletions

File tree

‎.github/workflows/build.yaml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ jobs:
126126
context: .
127127
file: ${{ matrix.dockerfile }}
128128
push: ${{ github.event_name != 'pull_request' }}
129-
platforms: "linux/amd64"
129+
platforms: ${{ matrix.platforms || 'linux/amd64' }}
130130
build-args: |
131131
SCCACHE_GHA_ENABLED=${{ matrix.sccache }}
132132
CUDA_COMPUTE_CAP=${{ matrix.cudaComputeCap }}
@@ -168,7 +168,7 @@ jobs:
168168
target: grpc
169169
file: ${{ matrix.dockerfile }}
170170
push: ${{ github.event_name != 'pull_request' }}
171-
platforms: "linux/amd64"
171+
platforms: ${{ matrix.platforms || 'linux/amd64' }}
172172
build-args: |
173173
SCCACHE_GHA_ENABLED=${{ matrix.sccache }}
174174
CUDA_COMPUTE_CAP=${{ matrix.cudaComputeCap }}

‎.github/workflows/matrix.json‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,16 @@
6363
"grpc": true,
6464
"dockerfile": "Dockerfile-cuda"
6565
},
66+
{
67+
"name": "blackwell-121",
68+
"imageNamePrefix": "121-",
69+
"runOn": "always",
70+
"sccache": true,
71+
"cudaComputeCap": 121,
72+
"platforms": "linux/amd64,linux/arm64",
73+
"grpc": true,
74+
"dockerfile": "Dockerfile-cuda"
75+
},
6676
{
6777
"name": "all",
6878
"imageNamePrefix": "cuda-",
@@ -79,6 +89,15 @@
7989
"grpc": true,
8090
"dockerfile": "Dockerfile"
8191
},
92+
{
93+
"name": "cpu-arm64",
94+
"imageNamePrefix": "cpu-arm64-",
95+
"runOn": "always",
96+
"sccache": true,
97+
"platforms": "linux/arm64",
98+
"grpc": true,
99+
"dockerfile": "Dockerfile-arm64"
100+
},
82101
{
83102
"name": "cpu-ipex",
84103
"imageNamePrefix": "cpu-ipex-",

‎Cargo.lock‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ homepage = "https://github.com/huggingface/text-embeddings-inference"
3131
anyhow = "1.0.75"
3232
clap = { version = "4.1", features = ["derive", "env"] }
3333
hf-hub = { version = "0.4", features = ["tokio"], default-features = false }
34-
metrics = "0.23"
34+
metrics = "0.23.1"
3535
nohash-hasher = "0.2"
3636
num_cpus = "1.16.0"
3737
tokenizers = { version = "0.21.0", default-features = false, features = ["onig", "esaxx_fast"] }

‎Dockerfile-cuda‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,14 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
1212
pkg-config \
1313
&& rm -rf /var/lib/apt/lists/*
1414

15-
# Donwload and configure sccache
16-
RUN curl -fsSL https://github.com/mozilla/sccache/releases/download/v$SCCACHE/sccache-v$SCCACHE-x86_64-unknown-linux-musl.tar.gz | tar -xzv --strip-components=1 -C /usr/local/bin sccache-v$SCCACHE-x86_64-unknown-linux-musl/sccache && \
15+
# Download and configure sccache (multi-arch)
16+
ARG TARGETARCH
17+
RUN case "${TARGETARCH}" in \
18+
"amd64") SCCACHE_ARCH=x86_64-unknown-linux-musl ;; \
19+
"arm64") SCCACHE_ARCH=aarch64-unknown-linux-musl ;; \
20+
*) echo "Unsupported arch: ${TARGETARCH}"; exit 1 ;; \
21+
esac && \
22+
curl -fsSL https://github.com/mozilla/sccache/releases/download/v$SCCACHE/sccache-v$SCCACHE-${SCCACHE_ARCH}.tar.gz | tar -xzv --strip-components=1 -C /usr/local/bin sccache-v$SCCACHE-${SCCACHE_ARCH}/sccache && \
1723
chmod +x /usr/local/bin/sccache
1824

1925
COPY rust-toolchain.toml rust-toolchain.toml
@@ -65,6 +71,9 @@ RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
6571
elif [ ${CUDA_COMPUTE_CAP} -eq 120 ]; \
6672
then \
6773
nvprune --generate-code code=sm_120 /usr/local/cuda/lib64/libcublas_static.a -o /usr/local/cuda/lib64/libcublas_static.a; \
74+
elif [ ${CUDA_COMPUTE_CAP} -eq 121 ]; \
75+
then \
76+
nvprune --generate-code code=sm_120 --generate-code code=sm_121 /usr/local/cuda/lib64/libcublas_static.a -o /usr/local/cuda/lib64/libcublas_static.a; \
6877
else \
6978
echo "cuda compute cap ${CUDA_COMPUTE_CAP} is not supported"; exit 1; \
7079
fi;
@@ -103,7 +112,13 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
103112
unzip \
104113
&& rm -rf /var/lib/apt/lists/*
105114

106-
RUN PROTOC_ZIP=protoc-21.12-linux-x86_64.zip && \
115+
ARG TARGETARCH
116+
RUN case "${TARGETARCH}" in \
117+
"amd64") PROTOC_ARCH=x86_64 ;; \
118+
"arm64") PROTOC_ARCH=aarch_64 ;; \
119+
*) echo "Unsupported arch: ${TARGETARCH}"; exit 1 ;; \
120+
esac && \
121+
PROTOC_ZIP=protoc-21.12-linux-${PROTOC_ARCH}.zip && \
107122
curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v21.12/$PROTOC_ZIP && \
108123
unzip -o $PROTOC_ZIP -d /usr/local bin/protoc && \
109124
unzip -o $PROTOC_ZIP -d /usr/local 'include/*' && \

‎README.md‎

Lines changed: 39 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ length of 512 tokens:
4343
- [Apple Silicon (Homebrew)](#apple-silicon-homebrew)
4444
- [Docker Build](#docker-build)
4545
- [Apple M1/M2 Arm](#apple-m1m2-arm64-architectures)
46+
- [ARM64 / aarch64](#arm64--aarch64)
4647
- [AMD Instinct GPUs (ROCm)](#amd-instinct-gpus-rocm-experimental)
4748
- [Examples](#examples)
4849

@@ -337,17 +338,19 @@ Options:
337338

338339
Text Embeddings Inference ships with multiple Docker images that you can use to target a specific backend:
339340

340-
| Architecture | Image |
341-
|----------------------------------------|-------------------------------------------------------------------------|
342-
| CPU | ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 |
343-
| Volta | NOT SUPPORTED |
344-
| Turing (T4, RTX 2000 series, ...) | ghcr.io/huggingface/text-embeddings-inference:turing-1.9 (experimental) |
345-
| Ampere 8.0 (A100, A30) | ghcr.io/huggingface/text-embeddings-inference:1.9 |
346-
| Ampere 8.6 (A10, A40, ...) | ghcr.io/huggingface/text-embeddings-inference:86-1.9 |
347-
| Ada Lovelace (RTX 4000 series, ...) | ghcr.io/huggingface/text-embeddings-inference:89-1.9 |
348-
| Hopper (H100) | ghcr.io/huggingface/text-embeddings-inference:hopper-1.9 |
349-
| Blackwell 10.0 (B200, GB200, ...) | ghcr.io/huggingface/text-embeddings-inference:100-1.9 (experimental) |
350-
| Blackwell 12.0 (GeForce RTX 50X0, ...) | ghcr.io/huggingface/text-embeddings-inference:120-1.9 (experimental) |
341+
| Architecture | Platform | Image |
342+
|----------------------------------------|----------|-------------------------------------------------------------------------|
343+
| CPU | x86_64 | ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 |
344+
| CPU | aarch64 | ghcr.io/huggingface/text-embeddings-inference:cpu-arm64-1.9 |
345+
| Volta | x86_64 | NOT SUPPORTED |
346+
| Turing (T4, RTX 2000 series, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:turing-1.9 (experimental) |
347+
| Ampere 8.0 (A100, A30) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:1.9 |
348+
| Ampere 8.6 (A10, A40, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:86-1.9 |
349+
| Ada Lovelace (RTX 4000 series, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:89-1.9 |
350+
| Hopper (H100) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:hopper-1.9 |
351+
| Blackwell 10.0 (B200, GB200, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:100-1.9 (experimental) |
352+
| Blackwell 12.0 (GeForce RTX 50X0, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:120-1.9 (experimental) |
353+
| Blackwell 12.1 (DGX Spark GB10, ...) | multi | ghcr.io/huggingface/text-embeddings-inference:121-1.9 (experimental) |
351354

352355
**Warning**: Flash Attention is turned off by default for the Turing image as it suffers from precision issues.
353356
You can turn Flash Attention v1 ON by using the `USE_FLASH_ATTENTION=True` environment variable.
@@ -610,21 +613,40 @@ runtime_compute_cap=100
610613
# Example for Blackwell (GeForce RTX 50X0, RTX PRO 6000, ...)
611614
runtime_compute_cap=120
612615

616+
# Example for Blackwell GB10 (DGX Spark)
617+
runtime_compute_cap=121
618+
613619
docker build . -f Dockerfile-cuda --build-arg CUDA_COMPUTE_CAP=$runtime_compute_cap
614620
```
615621

616-
### Apple M1/M2 arm64 architectures
622+
### ARM64 / aarch64
617623

618-
#### DISCLAIMER
624+
#### CPU-only (Apple Silicon, Ampere, Graviton)
619625

620-
As explained here [MPS-Ready, ARM64 Docker Image](https://github.com/pytorch/pytorch/issues/81224), Metal / MPS is not
621-
supported via Docker. As such inference will be CPU bound and most likely pretty slow when using this docker image on an
622-
M1/M2 ARM CPU.
626+
For ARM64 hosts without NVIDIA GPUs, use the CPU Dockerfile. Inference runs on CPU cores
627+
only (no Metal/MPS support via Docker).
623628

624-
```
629+
```shell
625630
docker build . -f Dockerfile-arm64 --platform=linux/arm64
626631
```
627632

633+
#### CUDA on ARM64 (DGX Spark, Jetson)
634+
635+
For ARM64 hosts with NVIDIA GPUs, build `Dockerfile-cuda` with the appropriate compute
636+
capability and `--platform linux/arm64`:
637+
638+
```shell
639+
# DGX Spark (GB10, sm_121)
640+
docker build . -f Dockerfile-cuda \
641+
--build-arg CUDA_COMPUTE_CAP=121 \
642+
--platform linux/arm64
643+
644+
# Future ARM64 + Blackwell devices (sm_120)
645+
docker build . -f Dockerfile-cuda \
646+
--build-arg CUDA_COMPUTE_CAP=120 \
647+
--platform linux/arm64
648+
```
649+
628650
## AMD Instinct GPUs (ROCm) — experimental
629651

630652
TEI has experimental support for AMD Instinct GPUs (MI200, MI300 series) via ROCm. You can use the `rocm/pytorch:latest` Docker image or a bare-metal ROCm installation. TEI will auto-detect the GPU at startup.

‎backends/candle/src/compute_cap.rs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@ fn compute_cap_matching(runtime_compute_cap: usize, compile_compute_cap: usize)
3131
(89, 89) => true,
3232
(90, 90) => true,
3333
(100, 100) => true,
34-
(120, 120) => true,
34+
(120..=121, 120) => true,
35+
(121, 121) => true,
3536
(_, _) => false,
3637
}
3738
}
@@ -57,6 +58,8 @@ mod tests {
5758
assert!(compute_cap_matching(89, 89));
5859
assert!(compute_cap_matching(90, 90));
5960
assert!(compute_cap_matching(120, 120));
61+
assert!(compute_cap_matching(121, 121));
62+
assert!(compute_cap_matching(121, 120));
6063

6164
assert!(compute_cap_matching(86, 80));
6265
assert!(compute_cap_matching(89, 80));
@@ -96,5 +99,12 @@ mod tests {
9699
assert!(!compute_cap_matching(120, 89));
97100
assert!(!compute_cap_matching(120, 90));
98101
assert!(!compute_cap_matching(120, 100));
102+
103+
assert!(!compute_cap_matching(121, 75));
104+
assert!(!compute_cap_matching(121, 80));
105+
assert!(!compute_cap_matching(121, 86));
106+
assert!(!compute_cap_matching(121, 89));
107+
assert!(!compute_cap_matching(121, 90));
108+
assert!(!compute_cap_matching(121, 100));
99109
}
100110
}

‎backends/candle/src/flash_attn.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ pub(crate) fn flash_attn_varlen(
6565
|| runtime_compute_cap == 90
6666
|| runtime_compute_cap == 100
6767
|| runtime_compute_cap == 120
68+
|| runtime_compute_cap == 121
6869
{
6970
#[cfg(feature = "flash-attn")]
7071
{

‎docs/source/en/supported_models.md‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,8 @@ Below are some examples of the currently supported models:
6868
## Supported hardware
6969

7070
Text Embeddings Inference supports can be used on CPU, Turing (T4, RTX 2000 series, ...), Ampere 80 (A100, A30),
71-
Ampere 86 (A10, A40, ...), Ada Lovelace (RTX 4000 series, ...), Hopper (H100), and Blackwell (B200, ...) architectures.
71+
Ampere 86 (A10, A40, ...), Ada Lovelace (RTX 4000 series, ...), Hopper (H100), and Blackwell (B200, RTX 5090, DGX Spark, ...) architectures.
72+
ARM64 (aarch64) is supported for both CPU-only and CUDA (Blackwell 12.1) workloads.
7273

7374
The library does **not** support CUDA compute capabilities < 7.5, which means V100, Titan V, GTX 1000 series, etc. are not supported.
7475

@@ -78,17 +79,19 @@ NVIDIA drivers with CUDA version 12.2 or higher.
7879

7980
Find the appropriate Docker image for your hardware in the following table:
8081

81-
| Architecture | Image |
82-
|----------------------------------------|-------------------------------------------------------------------------|
83-
| CPU | ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 |
84-
| Volta | NOT SUPPORTED |
85-
| Turing (T4, RTX 2000 series, ...) | ghcr.io/huggingface/text-embeddings-inference:turing-1.9 (experimental) |
86-
| Ampere 8.0 (A100, A30) | ghcr.io/huggingface/text-embeddings-inference:1.9 |
87-
| Ampere 8.6 (A10, A40, ...) | ghcr.io/huggingface/text-embeddings-inference:86-1.9 |
88-
| Ada Lovelace (RTX 4000 series, ...) | ghcr.io/huggingface/text-embeddings-inference:89-1.9 |
89-
| Hopper (H100) | ghcr.io/huggingface/text-embeddings-inference:hopper-1.9 |
90-
| Blackwell 10.0 (B200, GB200, ...) | ghcr.io/huggingface/text-embeddings-inference:100-1.9 (experimental) |
91-
| Blackwell 12.0 (GeForce RTX 50X0, ...) | ghcr.io/huggingface/text-embeddings-inference:120-1.9 (experimental) |
82+
| Architecture | Platform | Image |
83+
|----------------------------------------|----------|-------------------------------------------------------------------------|
84+
| CPU | x86_64 | ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 |
85+
| CPU | aarch64 | ghcr.io/huggingface/text-embeddings-inference:cpu-arm64-1.9 |
86+
| Volta | x86_64 | NOT SUPPORTED |
87+
| Turing (T4, RTX 2000 series, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:turing-1.9 (experimental) |
88+
| Ampere 8.0 (A100, A30) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:1.9 |
89+
| Ampere 8.6 (A10, A40, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:86-1.9 |
90+
| Ada Lovelace (RTX 4000 series, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:89-1.9 |
91+
| Hopper (H100) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:hopper-1.9 |
92+
| Blackwell 10.0 (B200, GB200, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:100-1.9 (experimental) |
93+
| Blackwell 12.0 (GeForce RTX 50X0, ...) | x86_64 | ghcr.io/huggingface/text-embeddings-inference:120-1.9 (experimental) |
94+
| Blackwell 12.1 (DGX Spark GB10, ...) | multi | ghcr.io/huggingface/text-embeddings-inference:121-1.9 (experimental) |
9295

9396
**Warning**: Flash Attention is turned off by default for the Turing image as it suffers from precision issues.
9497
You can turn Flash Attention v1 ON by using the `USE_FLASH_ATTENTION=True` environment variable.

0 commit comments

Comments
 (0)