feat: add SM87 and SM110 compute capability support for Jetson devices

Dockerfile-jetson

@@ -0,0 +1,92 @@
+FROM nvcr.io/nvidia/l4t-jetpack:r36.4.0 AS base-builder
+
+ENV SCCACHE=0.10.0
+ENV RUSTC_WRAPPER=/usr/local/bin/sccache
+ENV PATH="/root/.cargo/bin:${PATH}"
+# aligned with `cargo-chef` version in `lukemathwalker/cargo-chef:latest-rust-1.92-bookworm`
+ENV CARGO_CHEF=0.1.73
+
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    curl \
+    libssl-dev \
+    pkg-config \
+    && rm -rf /var/lib/apt/lists/*
+
+# Download and configure sccache for aarch64
+RUN curl -fsSL https://github.com/mozilla/sccache/releases/download/v$SCCACHE/sccache-v$SCCACHE-aarch64-unknown-linux-musl.tar.gz | tar -xzv --strip-components=1 -C /usr/local/bin sccache-v$SCCACHE-aarch64-unknown-linux-musl/sccache && \
+    chmod +x /usr/local/bin/sccache
+
+COPY rust-toolchain.toml rust-toolchain.toml
+RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
+RUN cargo install cargo-chef --version $CARGO_CHEF --locked
+
+FROM base-builder AS planner
+
+WORKDIR /usr/src
+
+COPY backends backends
+COPY core core
+COPY router router
+COPY Cargo.toml ./
+COPY Cargo.lock ./
+
+RUN cargo chef prepare  --recipe-path recipe.json
+
+FROM base-builder AS builder
+
+ARG GIT_SHA
+ARG DOCKER_LABEL
+
+# sccache specific variables
+ARG SCCACHE_GHA_ENABLED
+
+# Limit parallelism
+ARG RAYON_NUM_THREADS=4
+ARG CARGO_BUILD_JOBS
+ARG CARGO_BUILD_INCREMENTAL
+
+WORKDIR /usr/src
+
+COPY --from=planner /usr/src/recipe.json recipe.json
+
+RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
+    --mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
+    cargo chef cook --release --recipe-path recipe.json && sccache -s;
+
+RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
+    --mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
+    CUDA_COMPUTE_CAP=87 cargo chef cook --release --features candle-cuda --recipe-path recipe.json && sccache -s;
+
+COPY backends backends
+COPY core core
+COPY router router
+COPY Cargo.toml ./
+COPY Cargo.lock ./
+
+RUN --mount=type=secret,id=actions_results_url,env=ACTIONS_RESULTS_URL \
+    --mount=type=secret,id=actions_runtime_token,env=ACTIONS_RUNTIME_TOKEN \
+    CUDA_COMPUTE_CAP=87 cargo build --release --bin text-embeddings-router -F candle-cuda && sccache -s;
+
+RUN mv /usr/src/target/release/text-embeddings-router /usr/src/target/release/text-embeddings-router-87
+
+FROM nvcr.io/nvidia/l4t-cuda:12.6.11-runtime AS base
+
+ARG DEFAULT_USE_FLASH_ATTENTION=True
+
+ENV HUGGINGFACE_HUB_CACHE=/data \
+    PORT=80 \
+    USE_FLASH_ATTENTION=$DEFAULT_USE_FLASH_ATTENTION \
+    LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
+
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    ca-certificates \
+    libssl-dev \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY --from=builder /usr/src/target/release/text-embeddings-router-87 /usr/local/bin/text-embeddings-router-87
+
+COPY --chmod=775 jetson-entrypoint.sh entrypoint.sh
+
+ENTRYPOINT ["./entrypoint.sh"]
+CMD ["--json-output"]

backends/candle/src/compute_cap.rs

@@ -28,9 +28,12 @@ fn compute_cap_matching(runtime_compute_cap: usize, compile_compute_cap: usize)
         (75, 75) => true,
         (80..=89, 80) => true,
         (86..=89, 80..=86) => true,
+        (87, 87) => true,
         (89, 89) => true,
         (90, 90) => true,
         (100, 100) => true,
+        (110, 110) => true,
+        (120, 120) => true,
         (120..=121, 120) => true,
         (121, 121) => true,
         (_, _) => false,
@@ -55,50 +58,82 @@ mod tests {
         assert!(compute_cap_matching(75, 75));
         assert!(compute_cap_matching(80, 80));
         assert!(compute_cap_matching(86, 86));
+        assert!(compute_cap_matching(87, 87));
         assert!(compute_cap_matching(89, 89));
         assert!(compute_cap_matching(90, 90));
+        assert!(compute_cap_matching(110, 110));
         assert!(compute_cap_matching(120, 120));
         assert!(compute_cap_matching(121, 121));
         assert!(compute_cap_matching(121, 120));
 
         assert!(compute_cap_matching(86, 80));
+        assert!(compute_cap_matching(87, 80));
         assert!(compute_cap_matching(89, 80));
+        assert!(compute_cap_matching(87, 86));
         assert!(compute_cap_matching(89, 86));
 
         assert!(!compute_cap_matching(75, 80));
         assert!(!compute_cap_matching(75, 86));
+        assert!(!compute_cap_matching(75, 87));
         assert!(!compute_cap_matching(75, 89));
         assert!(!compute_cap_matching(75, 90));
+        assert!(!compute_cap_matching(75, 110));
 
         assert!(!compute_cap_matching(80, 75));
         assert!(!compute_cap_matching(80, 86));
+        assert!(!compute_cap_matching(80, 87));
         assert!(!compute_cap_matching(80, 89));
         assert!(!compute_cap_matching(80, 90));
+        assert!(!compute_cap_matching(80, 110));
 
         assert!(!compute_cap_matching(86, 75));
+        assert!(!compute_cap_matching(86, 87));
         assert!(!compute_cap_matching(86, 89));
         assert!(!compute_cap_matching(86, 90));
+        assert!(!compute_cap_matching(86, 110));
+
+        assert!(!compute_cap_matching(87, 75));
+        assert!(!compute_cap_matching(87, 89));
+        assert!(!compute_cap_matching(87, 90));
+        assert!(!compute_cap_matching(87, 110));
 
         assert!(!compute_cap_matching(89, 75));
+        assert!(!compute_cap_matching(89, 87));
         assert!(!compute_cap_matching(89, 90));
+        assert!(!compute_cap_matching(89, 110));
 
         assert!(!compute_cap_matching(90, 75));
         assert!(!compute_cap_matching(90, 80));
         assert!(!compute_cap_matching(90, 86));
+        assert!(!compute_cap_matching(90, 87));
         assert!(!compute_cap_matching(90, 89));
+        assert!(!compute_cap_matching(90, 110));
 
         assert!(!compute_cap_matching(100, 75));
         assert!(!compute_cap_matching(100, 80));
         assert!(!compute_cap_matching(100, 86));
+        assert!(!compute_cap_matching(100, 87));
         assert!(!compute_cap_matching(100, 89));
         assert!(!compute_cap_matching(100, 90));
+        assert!(!compute_cap_matching(100, 110));
 
         assert!(!compute_cap_matching(120, 75));
         assert!(!compute_cap_matching(120, 80));
         assert!(!compute_cap_matching(120, 86));
+        assert!(!compute_cap_matching(120, 87));
         assert!(!compute_cap_matching(120, 89));
         assert!(!compute_cap_matching(120, 90));
         assert!(!compute_cap_matching(120, 100));
+        assert!(!compute_cap_matching(120, 110));
+
+        assert!(!compute_cap_matching(110, 75));
+        assert!(!compute_cap_matching(110, 80));
+        assert!(!compute_cap_matching(110, 86));
+        assert!(!compute_cap_matching(110, 87));
+        assert!(!compute_cap_matching(110, 89));
+        assert!(!compute_cap_matching(110, 90));
+        assert!(!compute_cap_matching(110, 100));
+        assert!(!compute_cap_matching(110, 120));
 
         assert!(!compute_cap_matching(121, 75));
         assert!(!compute_cap_matching(121, 80));

jetson-entrypoint.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+if ! command -v nvidia-smi &>/dev/null; then
+    echo "Error: 'nvidia-smi' command not found."
+    exit 1
+fi
+
+# On Jetson L4T, CUDA libraries are provided by the host via nvidia-container-runtime.
+# Add compat path if it exists.
+if [ -d /usr/local/cuda/compat ]; then
+    export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH}"
+fi
+
+compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv | sed -n '2p' | sed 's/\.//g')
+
+if [ ${compute_cap} -eq 87 ]; then
+    exec text-embeddings-router-87 "$@"
+else
+    echo "cuda compute cap ${compute_cap} is not supported by the Jetson image (supported: 87)"
+    exit 1
+fi