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

- Add (87, 87) and (110, 110) match arms in compute_cap.rs for dedicated
  Jetson Orin (SM87) and Jetson Thor (SM110) binary support
- Add Dockerfile-jetson: multi-target Dockerfile building both SM87 and
  SM110 binaries with architecture-aware sccache for aarch64/x86_64
- Add jetson-entrypoint.sh: runtime GPU detection and routing to the
  correct binary based on compute capability
- Add comprehensive test coverage for SM87 and SM110 cross-compatibility

Author: thomas-hiddenpeak

Dockerfile-jetson

@@ -0,0 +1,107 @@
+FROM nvidia/cuda:12.9.1-devel-ubuntu24.04 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 (detect architecture for Jetson aarch64 support)
+RUN ARCH=$(uname -m) && \
+    if [ "$ARCH" = "aarch64" ]; then SCCACHE_TARGET="aarch64-unknown-linux-musl"; \
+    else SCCACHE_TARGET="x86_64-unknown-linux-musl"; fi && \
+    curl -fsSL https://github.com/mozilla/sccache/releases/download/v$SCCACHE/sccache-v$SCCACHE-$SCCACHE_TARGET.tar.gz | tar -xzv --strip-components=1 -C /usr/local/bin sccache-v$SCCACHE-$SCCACHE_TARGET/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;
+
+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=110 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
+
+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=110 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-110
+
+FROM nvidia/cuda:12.9.1-runtime-ubuntu24.04 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 \
+    cuda-compat-12-9 \
+    && 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 --from=builder /usr/src/target/release/text-embeddings-router-110 /usr/local/bin/text-embeddings-router-110
+
+COPY --chmod=775 jetson-entrypoint.sh entrypoint.sh
+
+ENTRYPOINT ["./entrypoint.sh"]
+CMD ["--json-output"]

backends/candle/src/compute_cap.rs

@@ -28,9 +28,11 @@ 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,
         (_, _) => false,
     }
@@ -54,47 +56,79 @@ 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(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));
     }
 }

jetson-entrypoint.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+
+if ! command -v nvidia-smi &>/dev/null; then
+    echo "Error: 'nvidia-smi' command not found."
+    exit 1
+fi
+
+# NOTE: Given that we need to support CUDA versions earlier than CUDA 12.9.1, we
+# need to include the `cuda-compat-12-9` in `LD_LIBRARY_PATH` when the host CUDA
+# version is lower than that; whilst we shouldn't include that when CUDA is 13.0+
+# as otherwise it will fail due to it.
+if [ -d /usr/local/cuda/compat ]; then
+    DRIVER_CUDA=$(nvidia-smi 2>/dev/null | awk '/CUDA Version/ {print $3; exit}')
+
+    IFS='.' read -r MAJ MIN PATCH <<EOF
+${DRIVER_CUDA:-0.0.0}
+EOF
+    : "${MIN:=0}"
+    : "${PATCH:=0}"
+
+    DRIVER_INT=$((10#${MAJ} * 10000 + 10#${MIN} * 100 + 10#${PATCH}))
+    TARGET_INT=$((12 * 10000 + 9 * 100 + 1))
+
+    if [ "$DRIVER_INT" -lt "$TARGET_INT" ]; then
+        export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH}"
+    fi
+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 "$@"
+elif [ ${compute_cap} -eq 110 ]; then
+    exec text-embeddings-router-110 "$@"
+else
+    echo "cuda compute cap ${compute_cap} is not supported by the Jetson image (supported: 87, 110)"
+    exit 1
+fi