Retrieval quality in an AI search product depends on two factors: the strength of the embedding model and the cost of running it at scale. Teams often face a trade‑off between high‑quality embeddings and affordable latency, especially when the same model must serve both offline index building and real‑time query embedding. Perplexity’s engineering write‑up shows how they solved this by reusing a single inference engine for both workloads instead of building separate pipelines.
The system splits work across three services. Ivy, a Rust HTTP gateway, handles JSON parsing, tokenization, input templating and batch splitting, then forwards requests via a custom gRPC protocol. It also balances load by chunking large batches and distributing them across replicas. Tulip, a Rust gRPC server built on tokio and tonic, schedules incoming requests, packs them using a simple first‑come‑first‑served policy, and hands the batch to the inference engine. ROSE, the Runtime‑Optimized Serving Engine written in Python, owns the model kernels, manages CUDA graphs and exposes a step() function that returns a handle to the GPU computation.
Two key optimizations improve throughput. First, whole‑model CUDA graphs capture every kernel launch into a single driver call, eliminating launch overhead for small batches. Because graphs must be built per configuration, Perplexity uses lazy capture: a configuration runs eagerly once, then triggers graph capture on its second use, spreading the upfront cost over hours of traffic. Second, a LazyTensor abstraction overlaps CPU preparation with GPU execution by returning a non‑blocking handle that tracks an async host‑to‑device copy and a CUDA event, letting Tulip prepare the next batch while the current one runs on the GPU.
The scheduler stays deliberately simple: for sub‑billion‑parameter embedding models at the served sequence lengths, dense‑layer cost dominates attention cost, so latency scales with token count, not batch size. Once a batch reaches roughly 512 tokens the GPU is saturated, and adding more sequences yields no extra efficiency. This insight lets the team pack requests without complex heuristics.
The result is a unified serving stack that delivers high‑quality embeddings at low cost for both batch re‑indexing and online queries, removing the need for separate tuning pipelines. #AI #Product #ML #Embeddings #GPU #Inference