Schnelle Embeddings auf GPUs
Eine schnelle und präzise Suche ist für ganz Perplexity von vitaler Bedeutung, von Search und Computer bis hin zu unserer API-Plattform. Hinter den Kulissen wird die Hauptarbeit von Embedding- und Ranking-Modellen geleistet, die unseren Systemen helfen, die relevantesten Ergebnisse für eine gegebene Anfrage zu identifizieren. Wir erreichen state-of-the
Eine schnelle und präzise Suche ist für ganz Perplexity von vitaler Bedeutung, von Search und Computer bis hin zu unserer API-Plattform. Hinter den Kulissen wird die Hauptarbeit von Embedding- und Ranking-Modellen geleistet, die unseren Systemen helfen, die relevantesten Ergebnisse für eine gegebene Anfrage zu identifizieren. Wir erreichen branchenführende Qualität und Latenz, indem wir unsere eigenen Modelle wie pplx-embed trainieren und bereitstellen.
Dieser Artikel bietet einen Blick hinter die Kulissen der Serving-Infrastruktur von Perplexity für diese spezielle Klasse von Modellen. Wir diskutieren unsere Techniken zur effizienten Bewältigung der Inferenzanforderungen der KI-nativen Suche, die ein schnelles Prototyping und die Evaluierung von Modellen ermöglichen und gleichzeitig unseren Suchindex im Exabytemaßstab antreiben. Diese Techniken erweitern gemeinsam die Pareto-Grenze für Suchqualität und -effizienz und versetzen uns in die Lage, Agenten und Benutzern die bestmöglichen Ergebnisse bei geringsten Kosten und Latenzen zu liefern.
Embeddings für die Suche
In einem typischen Such-Setup werden indizierte Dokumente mithilfe eines Embedding-Modells auf einen hochdimensionalen Vektorraum abgebildet und in einer Vektordatenbank gespeichert. Durch das Einbetten einer Abfrage mit demselben Modell können ähnliche Dokumente gefunden werden, indem die Vektoren ermittelt werden, die dem der Abfrage am nächsten liegen. Dies führt zu zwei unterschiedlichen Traffic-Mustern, die eine Inferenz-Engine bedienen muss:
- Batch Embedding: Beim Erstellen, Erweitern oder Re-Indexieren der Datenbank müssen große Mengen an Dokumenten in den Vektorraum eingebettet werden, um den Durchsatz zu maximieren und die Kosten zu minimieren.
Im Anschluss an die Vektorsuche müssen große Batches von Dokumenten bewertet werden, wobei ein Gleichgewicht zwischen Durchsatz und Latenz gefunden werden muss.
- Online Embedding: Bei der Abfrage der Datenbank muss eine kurze Abfrage für Lookups eingebettet werden, um die Latenz zu minimieren.
Wir haben unsere Inferenzinfrastruktur so aufgebaut, dass über Anwendungsfälle hinweg so viele gemeinsame Komponenten wie möglich genutzt werden. Da wir typischerweise kleine Transformer-Modelle verwenden, um Embeddings zu erzeugen, teilen wir den Großteil der Implementierung mit unserem LLM-Inferenzcode: Batch-Embeddings ähneln rechengebundenem Prefill, während Online-Embeddings, die oft auf wenigen Token laufen, rechentechnisch speichergebundenem Decode ähneln. Wir wiederverwenden daher unsere optimierten Prefill- und Decode-Kernels, um Embedding-Modelle zu bedienen. Dadurch können wir einen enormen Batch-Inferenz-Durchsatz bei minimalem zusätzlichen Engineering-Aufwand erzielen und gleichzeitig eine niedrige Latenz für Online-Embedding-Workloads beibehalten.
Tulips, Roses und etwas Ivy
Wir stellen Inferenz über standardisierte APIs zur Verfügung, sowohl intern als auch extern über unsere API-Plattform. Hinter den Kulissen sind mehrere Dienste an der Verarbeitung einer Embedding-Anfrage beteiligt:
- Ivy ist ein Rust-HTTP-Gateway, das von Perplexity-Diensten aufgerufen wird.
Es übernimmt die CPU-seitige Arbeit für Anfragen wie JSON-Parsing, Tokenisierung, Input-Templating und Batch-Splitting und übersetzt Anfragen in ein benutzerdefiniertes gRPC-Protokoll für nachgeschaltete Server. Diese Trennung ermöglicht es uns, bestimmte Parameter rund um Tokenisierung und Eingabeformatierung zu konfigurieren, ohne die schwereren Inferenzinstanzen anfassen zu müssen.
- Tulip ist die Inferenz-Server-Schnittstelle.
Es handelt sich um einen gRPC-Server, der mit Rust, tokio und `tonic` implementiert ist. Tulip empfängt gRPC-Inferenzanfragen und übernimmt Planung und Batching. Anschließend sendet es die Batches an die ROSE-Engine und gibt fertige Antworten an Clients zurück.
- **ROSE** (Runtime-Optimized Serving Engine) implementiert Modellinferenz.
Es ist primär in Python definiert und stellt Kernels, Layer und Definitionen für eine Vielzahl von Modellen bereit. ROSE implementiert die Forward-Pässe durch Modelle und bietet zudem eine CUDA-Graphenverwaltung, die auf Embeddings spezialisiert ist. Es ist über eine step()-Funktion an Tulip angebunden, die einen Batch entgegennimmt und eine Referenz auf die Berechnung zurückgibt, die sie auf dem Beschleuniger durchführt.
Aufmerksamkeit über den Kernel hinaus
Sowohl Transformer-basierte Modelle als auch die zugrundeliegenden Hopper/Blackwell-Architekturen sind ausgereifte Technologien, sodass die Einbettungs-Inferenz auf der GPU-Seite über verschiedene Inferenz-Engines hinweg zu einer weitgehend optimalen Implementierung konvergiert ist. Dennoch haben wir zusätzliche Verbesserungsmöglichkeiten bei Laufzeiten und Harnesses entdeckt, die die Modelle Ende-zu-Ende für einen Client verfügbar machen. Insbesondere stellten wir fest, dass wir die Latenzen verbessern können, indem wir CUDA-Graphen sorgfältig verwalten und eine LazyTensor-Abstraktion entwickeln, um ein GPU-seitiges Ergebnis in der nativen Rust-Engine asynchron zu verfolgen. Wir haben diese Funktionen in Tulip implementiert, damit es effektiv mit den Modellimplementierungen von ROSE interagieren kann.
Tulip
Wir haben Tulip so konzipiert, dass es eine möglichst leichtgewichtige Schnittstelle über unserem Modell-Serving ist. Es verarbeitet eingehende Anfragen in Tokio-Async-Tasks, verwaltet einen Pool von Anfragen, die es verfolgt, und plant Batches für die Weiterleitung an den Beschleuniger ein. Der Planungsmechanismus in Tulip ist sehr einfach: Anfragen akkumulieren sich, während Tulip Arbeit verteilt oder auf Ergebnisse wartet. Aus den akkumulierten Anfragen werden Sequenzen nach dem Prinzip „Wer zuerst kommt, mahlt zuerst“ ausgewählt, um durch das Modell ausgeführt zu werden.
Der einfache Planungsmechanismus wird durch eine Beobachtung der Modellleistung motiviert. Bei kleinen Embedding-Modellen stellten wir bei den von uns bedienten Sequenzlängen fest, dass die linearen Kosten von Dense-Schichten gegenüber den quadratischen Kosten von Attention dominieren. Daher ist die Latenz größtenteils proportional zur Anzahl der Token, nicht zur Anzahl der Sequenzen. Folglich verbessert das Hinzufügen weiterer Sequenzen zu einem Batch, sobald dieser groß genug ist, um die GPU zu saturieren (was bei einem Modell unter einer Milliarde Parametern bei etwa 512 Token der Fall ist), die Effizienz nicht mehr.
Um effektiv mit dem Modell zu interagieren, stützt sich Tulip auf CUDA-Graphen und Lazy-Ergebnisverfolgung, um GPU- und CPU-Arbeit zu überlappen und die verfügbaren Ressourcen vollständig auszunutzen.
CUDA-Graphenverwaltung
Die Ausführung des Forward-Passes eines Modells umfasst sowohl CPU- als auch GPU-seitige Arbeit. Die CPU ist für das Planen von Batches und das Starten von Kernels mit den entsprechenden Parametern verantwortlich, während die GPU die relevanten Matrixmultiplikations-, Attention-, Norm- oder Aktivierungskernels ausführt. Bei Workloads mit hohem Durchsatz wie Training und Reindexing sind CPU-seitige Overheads vernachlässigbar, da Batchgrößen und GPU-seitige Latenz beide groß sind. Bei kleineren Batchgrößen kann die CPU-seitige Arbeit jedoch die GPU-seitige überwiegen.
Um Overheads zu reduzieren, kann anstelle des Startens unabhängiger Kernels ein CUDA-Graph erstellt werden, um die Metadaten zu erfassen, die erforderlich sind, um alle Kernels eines Forward-Passes mit einem einzigen Aufruf des CUDA-Treibers zu starten. Dadurch entfällt die Notwendigkeit, teuren Python- und PyTorch-Code für die Konfigurationen, für die CUDA-Graphen erfasst werden können, erneut auszuführen.
Bei jedem Modell verfolgen wir einen Wendepunkt und bestimmen die minimale Anzahl von Token, bei der die GPU-Ausführung teurer ist als der CPU-seitige Kernel-Start. Da Embedding-Modelle klein sind, stellen wir fest, dass dieser Wendepunkt bei Batches von Tausenden von Token und Dutzenden von Sequenzen liegt. Einige Attention-Implementierungen verlassen sich auf dynamische hostseitige Eingaben, um Kernel-Starts zu konfigurieren, was Full-Model-Prefill/Dense-CUDA-Graphen verhindert. Wir haben Änderungen an relevanten Kernels upgestreamt, um sie in unserer Inferenz-Engine zu aktivieren.
Um Overheads anzugehen, erstellen wir Full-Model-CUDA-Graphen für alle Embedding-Modelle und überlappen CPU-Arbeit mit GPU-Arbeit. Da CUDA-Graphen CPU-seitige Overheads minimieren, haben wir nach dem Start eines Graphen freie Zeit, um die Ausführung des nächsten Batches anzustoßen und in die Warteschlange einzureihen, sobald dieser verfügbar ist. Die Ergebnisse des ausstehenden Batches werden mit einem LazyTensor verfolgt, was es einem Async-Task in Rust ermöglicht zu blockieren, bis der vorherige Batch seine Ausführung beendet hat. CUDA-Graphen unterstützen das Serving mit niedriger Latenz, indem sie sicherstellen, dass wir nicht durch die Kosten von Kernel-Starts ausgebremst werden, und erleichtern eine verbesserte Planung im Fall eines hohen Durchsatzes, da sie die CPU entlasten, um früher an der nächsten Batch zu arbeiten.
CUDA-Graphen müssen für jede einzelne Konfiguration erfasst werden, was für Embeddings einen Graphen pro Kombination aus Sequenzanzahl und Tokenanzahl bedeutet. Da dieses Raster umfangreich ist, füllen wir Token-Anzahlen (Padding) auf Buckets auf, die ein Vielfaches von 64 oder 256 sind. Dies führt dennoch zu Tausenden von Graphen, deren Erfassung für ein typisches Modell mehrere Minuten dauern kann. Die Kosten für die Erfassung stammen aus zwei Quellen: einem Eager Forward-Pass, der ausgeführt werden muss, um Kernels zu kompilieren und Puffer für verschiedene Kernels einzurichten, die sie benötigen, gefolgt von dem Erfassungslauf, der den Python-Code erneut ausführt.
Wir mildern die Startkosten ab, indem wir CUDA-Graphen während des Servings der Engine lazy erfassen. Wir verfolgen jede Konfiguration und stellen sicher, dass sie einen Eager-Warmup-Lauf durchläuft, bevor beim zweiten Aufruf die Erfassung und das Replay des Graphen ausgelöst werden. Alle nachfolgenden Ausführungen derselben Graphenkonfiguration durchlaufen dann das CUDA-Graphen-Replay. Die Lazy-Graphenerfassung hat Auswirkungen auf die p99-Latenzen während des Startvorgangs; sie ist jedoch wertvoll, um die mehrinütige Eager-Arbeit über mehrere Stunden zu verteilen. Kürzere Startzeiten ermöglichen es uns, Embedding-Deployments besser zu skalieren und zu verwalten.
Lazy Tensors
Durch CUDA ist die GPU-Arbeit asynchron. Da das asynchrone Starten eines Kernels diesen in einer Warteschlange (Stream) einreiht, muss Host-Code explizit synchronisieren, um die resultierenden Vektoren auszulesen. Um einen höheren Grad an Parallelität zu ermöglichen und zukünftige Batches anstoßen zu können, während auf den Abschluss des vorherigen auf dem Device gewartet wird, stützen wir uns zur Verfolgung von Werten auf eine LazyTensor-Abstraktion.
Der LazyTensor verfolgt einen Host-Puffer inseitig gesperrtem Speicher (page-locked memory) und eine cudaMemcpyAsync-Operation über ein Event, das Daten vom Device kopiert. Er wird nach dem Start des Forward-Passes auf demselben Stream angestoßen. Da die Kopiervorgangsoperation auf den Abschluss aller vorherigen Kernels auf dem Stream warten muss, verfolgt das zugehörige Event sowohl den Abschluss des Forward-Passes als auch die Verfügbarkeit des Ergebnisses auf der CPU.
Wir nutzen LazyTensors in unserer ROSE-Encoder-Engine, um GPU- und CPU-Arbeit zu überlappen. Anstatt dass jeder step()-Aufruf den CUDA-Graphen ausführt und auf dessen Abschluss wartet, gibt step() einen LazyTensor zurück, um sein Ergebnis asynchron zu verfolgen. In Verbindung mit CUDA-Graphen hilft uns dies, niedrige Latenzen und einen besseren Durchsatz zu erreichen.
ROSE
Wir haben unsere ROSE-Engine, die wir ursprünglich für das LLM-Serving entwickelt hatten, so angepasst, dass sie auch die Ausführung von Embedding-Modellen übernimmt. Um den Aufwand für die Unterstützung von Embedding-Modellen zu minimieren, reuse ROSE Code zwischen LLMs und Embeddings aggressiv. Zum Beispiel durchlaufen pplx-embed Serving und Qwen3.5 LLM-Decoding alle dieselben Kernels. Diese gemeinsame Nutzung ermöglicht es uns, problemlos ein Embedding-Modell bereitzustellen, das ursprünglich für Prototyping, Evaluierung und Produktionsinferenz aus einem LLM fine-getuned wurde.
Für Dense-Layer sind Embedding- und LLM-Inferenz identisch, da Token-Vektoren unabhängig voneinander verarbeitet werden. In Attention-Schichten werden Unterschiede durch die Unterstützung unregelmäßiger (ragged) Eingaben sowie der für LLMs erforderlichen gepaged Prefill- und Decode-Setups gehandhabt. Beim Serving eines Embedding-Modells instanziieren wir keinen KV-Cache und leiten an Varianten von Attention-Kernels weiter, die das Ragged-Format unterstützen, um Padding zu vermeiden. Die unterstützenden Konvertierungs- und Kalibrierungsroutinen werden ebenfalls mit den LLMs geteilt.
Ivy
Ivy, unsere Inferenz-HTTP-Proxyschicht, spielt ebenfalls eine wichtige Nutzlastrolle für die Leistung. Da die Anfrage-Payrolls in der Produktion variieren, kann das Routing einzelner Anfragen an einzelne Replikate zu Lastenausgleichsproblemen führen. Ivy unterteilt große Batch-Anfragen in Chunks und verteilt sie per Lastenausgleich zwischen den Replikaten, wodurch die Auslastung verbessert und die Latenz glattgezogen wird. Unsere jüngste Arbeit an der eigenen Unigram-Tokenisierung, die vollständig in Ivy ausgerollt ist, verbessert die Latenzen im Vergleich zu Standard-Tokenizern drastisch.
...aber die Kernels zählen nach wie vor
ROSE unterstützt eine Vielzahl von Attention-Backends. Verschiedene Kernels eignen sich möglicherweise für spezifische Problemgrößen. Im Laufe der Zeit haben wir FlashInfer 2-, FlashInfer 3- und FlashAttention 4-Kernels integriert, um Ragged Attention zu implementieren.
Im Allgemeinen stellen wir fest, dass FlashAttention 4 schneller ist. FlashInfer 3 übertrifft es jedoch bei Qwen-basierten Modellen bei sehr langen Sequenzlängen. Da Leistung und Tuning von Anzahl und Dimension der Attention-Heads abhängen können, behalten wir die Unterstützung für mehrere Konfigurationen bei und treffen beim Serving eine Einzelfallentscheidung.
Benchmarks
Wir benchmarken gegen vLLM v0.22.0 und führen Inferenz mit BF16-Präzision auf tatsächlichen Modellgewichten und Eingaben aus, die aus Evaluierungsdatensätzen stammen. Allen Timing-Läufen gingen Warmup-Läufe voraus, die sicherstellten, dass die Abweichung der Kosinusähnlichkeit innerhalb von 0,1 % liegt.
Low-Latency Embeddings (p50 / p90 / p99 / max ms)
Wir berichten über Laufzeiten für prä-tokenisierte Anfrage-Batchgröße 1, voll sequentielle Anfragen, Sequenzlängen von 128, 512 und 4096 Token.
Low-Latency Scoring (p50 / p90 / p99 / max ms)
Prä-tokenisierte Anfrage-Batchgrößen 5, 25 und 50, Sequenzlänge von 512 Token.
High-Throughput Embeddings (emb/s)
Anfrage-Batchgröße 100, vier gleichzeitige Prozesse, die Anfragen übermitteln, Sequenzlängen von 512, 1024 und 4096 Token.
High-Concurrency Embeddings (p50 / p90 / p99 / max ms)
Sequenzlänge 512, Batchgröße 1, aber wir senden 1, 2, 4, 8 und 16 gleichzeitige Anfragen. Dieser Benchmark umfasst auch die Tokenisierungskosten über Ivy sowie den Netzwerk-Overhead zwischen Ivy und Tulip.
Fazit und zukünftige Arbeiten
Die aus Ivy, Tulip und ROSE bestehende Serving-Infrastruktur ermöglicht es uns, Embeddings für Perplexity mit geringerer Latenz und besserem Durchsatz bereitzustellen, was im Vergleich zu Standardlösungen zu einer präziseren Suche zu reduzierten Kosten führt.
Indem wir uns auf spezifische Modelle konzentrieren und die Verantwortung für den gesamten Stack übernehmen, erhalten wir die nötige Freiheit, um eine effektive Balance zwischen Leistung und Flexibilität zu finden, wobei wir hochgradig wiederverwendbare und performante Rust-Primitive mit generischerem Python-Modellierungscode kombinieren. Viele Open-Source-Inferenz-Engines wie vLLM, SGLang und TokenSpeed integrieren Sprachen wie Rust und C++ in ihren Stack. Wir haben in den letzten zwei Jahren in Rust investiert und große Erfolge sowohl bei der Leistung als auch bei der Wartbarkeit erzielt. Indem wir den Großteil der Embedding-Implementierung mit unserem LLM-Serving-Stack teilen, erzielen wir auch Gewinne beim Durchsatz, ohne dass nennenswerter Engineering-Aufwand für die Wartung von Embedding-Modellen betrieben werden muss.
Während sich Modelle weiterentwickeln, werden wir jede Schicht unseres Stacks kontinuierlich verbessern, um sowohl CPU- als auch GPU-gebundene Latenzen zu reduzieren. Unsere benutzerdefinierten gRPC-basierten Protokolle innerhalb von Ivy und Tulip ermöglichen es uns, die Kommunikation anzupassen, um Netzwerklatenzen zu reduzieren, während ROSE eine Grundlage zur Verbesserung des Durchsatzes bietet. Da die Unterstützung für Free-Threaded Python im gesamten Ökosystem wächst, werden wir zudem in der Lage sein, die Python-Rust-Interoperabilität weiter zu verbessern, um Overheads zu reduzieren.