$ show huggingface/text-embeddings-inference
Text Embeddings Inference (TEI) is a production-grade inference server for open-source text embedding models, built in Rust with multiple backend implementations (Candle, Python, ONNX). The system exposes both HTTP (REST) and gRPC APIs supporting dense and sparse embeddings, reranking, sequence classification, and OpenAI-compatible endpoints. It features dynamic batching with cumulative sequence length management, flash attention optimization across CUDA compute capabilities, and comprehensive observability with Prometheus metrics and OpenTelemetry tracing. The architecture cleanly separates the router layer (HTTP/gRPC handling) from the core inference pipeline (tokenization, batching, backend dispatch) via the `Backend` trait, allowing seamless integration of different model backends and hardware targets including CPU, CUDA, ROCm, Intel, and ARM64. The codebase demonstrates strong engineering practices: modular neural network layer implementations (Linear, LayerNorm, RMSNorm, Rotary) reused across model architectures, snapshot-based regression testing for all supported models, and feature-gated builds for precise dependency control. The router provides production features such as CORS, API key authentication, request size limits, health checks, and graceful shutdown. The batching system (`Queue` struct) efficiently accumulates requests and dispatches them to the backend in configurable token-bucket sizes. While the engineering quality is high, the repository exhibits some maintenance risks: significant model duplication across `flash_` and non-flash variants for each architecture, patched upstream dependencies (candle, cudarc) that may diverge from official releases, and commented-out integration tests. The Python backend, wrapping a subprocess, introduces platform-dependent fragility, and the large number of hardware-specific Dockerfiles increases build complexity. Overall, TEI is a well-architected, actively maintained inference server suitable for production deployments, with reusable patterns in its backend abstraction, batch processing, and layer implementations.
The LLM identified 4 files worth highlighting for reuse:
Provides a unified flash attention interface that dynamically selects between flash attention v1 and v2 implementations based on GPU compute capability (Turing, Ampere, Ada, Hopper, Blackwell). The module abstracts away the complexity of different CUDA architectures and kernel versions, falling back to non-flash implementations when unsupported. This pattern is reusable in any Rust-based ML inference engine that needs to target multiple GPU generations with optimal attention kernels.
Implements a fully-featured linear (dense) layer with support for cuBLASLt acceleration, custom fused kernels, and standard matrix multiplication fallback. The module handles weight loading from different tensor formats, bias handling, and optional matrix rearrangement for optimized execution. This is a self-contained neural network primitive that can be extracted for use in any Rust-based transformer implementation requiring performant linear transformations.
Implements a background batching queue that accumulates inference requests and dispatches them as optimally-sized batches. The `Queue` struct uses a configurable maximum token budget (`max_batch_tokens`) and background task that continuously groups pending requests into `Batch` objects. This pattern is directly reusable in any inference serving system that needs efficient, token-aware batching without requiring a dedicated batching framework.
A lightweight, self-contained Rust crate that injects OpenTelemetry trace context into gRPC request metadata. It provides a `inject` function that serializes the current span context into gRPC metadata key-value pairs, enabling distributed tracing across service boundaries. This module has minimal dependencies (only `opentelemetry` and `tonic` types) and can be extracted into any Rust gRPC project requiring telemetry propagation.