From bf9ecfcae451531fea4daa8009b81ed2ebffd108 Mon Sep 17 00:00:00 2001 From: Thomas Hiddenpeak <118246850+thomas-hiddenpeak@users.noreply.github.com> Date: Fri, 13 Mar 2026 00:51:01 +0800 Subject: [PATCH] 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: builds SM87 binary using L4T JetPack r36.4.0 (CUDA 12.6) as build base and l4t-cuda:12.6.11-runtime for deployment - Add jetson-entrypoint.sh: runtime GPU detection for Jetson Orin (SM87) - Add comprehensive test coverage for SM87 and SM110 cross-compatibility --- Dockerfile-jetson | 92 ++++++++++++++++++++++++++++++ backends/candle/src/compute_cap.rs | 34 +++++++++++ jetson-entrypoint.sh | 21 +++++++ 3 files changed, 147 insertions(+) create mode 100644 Dockerfile-jetson create mode 100644 jetson-entrypoint.sh diff --git a/Dockerfile-jetson b/Dockerfile-jetson new file mode 100644 index 000000000..ee1a10942 --- /dev/null +++ b/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"] diff --git a/backends/candle/src/compute_cap.rs b/backends/candle/src/compute_cap.rs index 782806279..c211cbc0a 100644 --- a/backends/candle/src/compute_cap.rs +++ b/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)); } } diff --git a/jetson-entrypoint.sh b/jetson-entrypoint.sh new file mode 100644 index 000000000..e53fc9d91 --- /dev/null +++ b/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