add Jetson Orin support

[HiddenPeak]

Motivation and Context

NVIDIA Jetson Orin devices have a compute capability of 8.7, which is not currently supported in the compute_cap_matching function. This PR ensures that these devices can be used with the library by adding the necessary support.

What does this PR do?

This PR adds support for NVIDIA Jetson Orin devices by including the compute capability 8.7 in the compute_cap_matching function and updating the tests to ensure the new capability is correctly supported.

Fixes #466

Checklist

• I have read the contributor guidelines.
• I have added tests to verify my changes.
• I have tagged the appropriate reviewers.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

@OlivierDehaene OR @Narsil

[r0kk]

@HiddenPeak I am wondering if you could share reproducible steps how were you able to run text-embeddings-inference on Jetson AGX Orin. It would be greatly appreciated🙏.

Unfortunately I don't have deep enough knowledge to review you PR.

[HiddenPeak]

@r0kk
The Jetson Orin series uses the CUDA architecture SM8.7, which is part of the Ampere architecture. Theoretically, it should be compatible with TEI. However, in practical applications, there are many incompatibilities, making direct support generally impossible. During the process of attempting to use it, I encountered the following issues:

1. The compute_cap_matching() function does not support the SM87 architecture, so I modified the source code and recompiled it.
2. It is necessary to ensure that the GPU driver, CUDA runtime, and CUDA compiler are correctly installed and available in the environment variable paths.（on Jetpack 6.1 with cuda 12.6）
3. The compilation process is extremely long, and memory usage exceeds 90% (60GB).

I attempted to compile and deploy TEI on a Jetson AGX Orin 64G and found that it could not recognize SM87. Therefore, I modified the compute_cap_matching() function in backends/candle/src/compute_cap.rs to add support for the SM87 environment and architecture. Such modifications may not be effective in many cases, but fortunately, after making these changes, I was able to achieve support on the Jetson AGX Orin 64G. Not only did it not produce any errors, but it also showed excellent performance.

curl 127.0.0.1:8080/rerank \
    -X POST \
    -d '{"query": "What is Deep Learning?", "texts": ["Deep Learning is not...", "Deep learning is..."]}' \
    -H 'Content-Type: application/json'

logs

2025-01-04T20:38:18.706787Z  INFO text_embeddings_backend_candle: backends/candle/src/lib.rs:292: Starting FlashBert model on Cuda(CudaDevice(DeviceId(1)))
2025-01-04T20:38:31.539445Z  INFO text_embeddings_router: router/src/lib.rs:248: Warming up model
2025-01-04T20:38:32.189069Z  INFO text_embeddings_router::http::server: router/src/http/server.rs:1812: Starting HTTP server: 0.0.0.0:8080
2025-01-04T20:38:32.189098Z  INFO text_embeddings_router::http::server: router/src/http/server.rs:1813: Ready
2025-01-04T20:44:11.047170Z  INFO rerank{total_time="177.15121ms" tokenization_time="727.783µs" queue_time="79.024583ms" inference_time="87.618256ms"}: text_embeddings_router::http::server: router/src/http/server.rs:459: Success

More screenshots are as follows:

Therefore, I created a branch and added test code. After testing it in my application, I submitted a merge request.
Additionally, I also tried other embedding and rerank models, which ran well.

[r0kk]

@HiddenPeak
I can confirm that this is working on Jetson AGX 64GB. Thank you very much 🙏.

[HiddenPeak]

@HiddenPeak I can confirm that this is working on Jetson AGX 64GB. Thank you very much 🙏.

It's very cool~