Rychlá vložení na GPU
Rychlé a přesné vyhledávání je klíčové pro celou Perplexity, od Search a Computer až po naši API Platform. V zákulisí odvádějí těžkou práci modely vkládání a řazení, které našim systémům pomáhají identifikovat nejrelevantnější výsledky pro daný dotaz. Špičkové kvality
Rychlé a přesné vyhledávání je klíčové pro celou Perplexity, od Search a Computer až po naši API Platform. V zákulisí odvádějí těžkou práci modely vkládání a řazení, které našim systémům pomáhají identifikovat nejrelevantnější výsledky pro daný dotaz. Špičkové kvality a latence dosahujeme trénováním a obsluhou vlastních modelů, jako je pplx-embed.
Tento článek přináší pohled pod pokličku obslužné infrastruktury Perplexity pro tuto speciální třídu modelů. Diskutujeme o našich technikách pro efektivní řešení potřeb odvozování vyhledávání nativního pro AI, což umožňuje rychlé vytváření prototypů a vyhodnocování modelů a zároveň pohání náš vyhledávací index v exabytovém měřítku. Tyto techniky dohromady rozšiřují Paretovu hranici kvality a efektivnosti vyhledávání a umožňují nám obsluhovat agenty a uživatele s co nejlepšími výsledky při nejnižších nákladech a latenci.
Vložení pro vyhledávání
V typickém nastavení vyhledávání jsou indexované dokumenty mapovány do vysokorozměrného vektorového prostoru pomocí modelu vkládání a uloženy ve vektorové databázi. Vložením dotazu pomocí stejného modelu lze vyhledat podobné dokumenty nalezením vektorů nejbližších k vektoru dotazu. To vede ke dvěma různým vzorcům provozu, které má odvozovací engine obsluhovat:
- Dávkové vkládání: při vytváření, rozšiřování nebo reindexaci databáze je nutné hromadně vkládat dokumenty do vektorového prostoru, což maximalizuje propustnost za účelem minimalizace nákladů.
Po vyhledávání vektorů je nutné ohodnotit velké dávky dokumentů, čímž se dosáhne rovnováhy mezi propustností a latencí.
- Online vkládání: při dotazování do databáze je nutné pro vyhledávání vložit krátký dotaz, čímž se minimalizuje latence.
Naši odvozovací infrastrukturu jsme vybudovali tak, aby napříč případy použití využívala co nejvíce společných komponent. Vzhledem k tomu, že k vytváření vložení obvykle používáme malé modely Transformer, sdílíme většinu implementace s naším kódem pro odvozování LLM: dávková vložení jsou podobná předběžnému vyplnění vázanému na výpočet, zatímco online vkládání, která často běží nad několika tokeny, jsou z hlediska výpočtů podobná dekódování vázanému na paměť. Pro obsluhu modelů vkládání tak opakovaně používáme naše optimalizované kernely pro předběžné vyplnění a dekódování. Díky tomu můžeme dosáhnout masivní propustnosti dávkového odvozování s minimálním dalším inženýrským úsilím při zachování nízké latence pro online úlohy vkládání.
Tulipány, růže a nějaký břečťan
Odvozování vystavujeme prostřednictvím standardizovaných rozhraní API, a to jak interně, tak externě prostřednictvím naší API Platform. V zákulisí se na zpracování požadavku na vložení podílí několik služeb:
- Ivy je HTTP brána v jazyce Rust, kterou služby Perplexity volají.
Zpracovává práci na straně CPU pro požadavky, jako je parsování JSON, tokenizace, šablonování vstupů a rozdělování dávek, a překládá požadavky na vlastní protokol gRPC pro navazující servery. Toto oddělení nám umožňuje konfigurovat určité parametry týkající se tokenizace a formátování vstupů, aniž bychom museli sahat na těžší instance odvozování.
- Tulip je rozhraní odvozovacího serveru.
Jedná se o server gRPC implementovaný v jazyce Rust pomocí knihoven tokio a `tonic`. Tulip přijímá požadavky na odvozování gRPC, přičemž vyřizuje plánování a dávkování. Následně odesílá dávky do enginu ROSE a vrací dokončené odpovědi klientům.
- **ROSE** (Runtime-Optimized Serving Engine) implementuje odvozování modelů.
Je definován primárně v Python, kde poskytuje kernely, vrstvy a definice pro širokou škálu modelů. ROSE implementuje průchody vpřed skrz modely a rovněž poskytuje správu grafů CUDA specializovanou pro vložení. S Tulip je propojen přes funkci step(), která přijímá dávku a vrací odkaz na výpočet, jenž provádí na akcelerátoru.
Věnování pozornosti nad rámec kernelu
Modely založené na Transformeru i podkladové architektury Hopper/Blackwell jsou zralé technologie, takže vkládání odvození na straně GPU dospělo k převážně optimální implementaci napříč různými odvozovacími enginy. Přesto jsme objevili další příležitosti ke zlepšení v runtimech a rozhraních, která vystavují modely klientovi end-to-end. Zejména jsme zjistili, že můžeme zlepšit latence pečlivou správou grafů CUDA a sestavením abstrakce LazyTensor pro asynchronní sledování výsledku na straně GPU v nativním enginu Rust. Tyto funkce jsme implementovali v Tulip, aby mohla efektivně propojit s implementacemi modelů ROSE.
Tulip
Navrhli jsme Tulip tak, aby byl co nejlehčím rozhraním nad naší obsluhou modelů. Zpracovává příchozí požadavky v asynchronních úlohách Tokio, udržuje fond požadavků, které sleduje, a plánuje z nich dávky k odeslání na akcelerátor. Mechanismus plánování v Tulip je velmi jednoduchý: požadavky se kumulují, zatímco Tulip odesílá práci nebo čeká na výsledky. Z nahromaděných požadavků jsou vybírány sekvence na základě principu „kdo dřív přijde, ten dřív mele“, které jsou spuštěny skrz model.
Jednoduchý mechanismus plánování je motivován pozorováním výkonu modelu. U malých modelů vkládání jsme si při délkách sekvencí, které obsluhujeme, všimli, že lineární náklad hustých vrstev převažuje nad kvadratickým nákladem pozornosti. Latence je tedy většinou úměrná počtu tokenů, nikoliv počtu sekvencí. Jakmile je tedy dávka dostatečně velká na to, aby nasycovala GPU, což je přibližně 512 tokenů u modelu s méně než jednou miliardou parametrů, napěchování dalších sekvencí do ní nezlepší efektivitu.
Aby se Tulip mohl efektivně propojit s modelem, spoléhá se na grafy CUDA a líné sledování výsledků k překrytí práce GPU a CPU a úplnému využití dostupných zdrojů.
Správa grafů CUDA
Spuštění průchodu modelu vpřed zahrnuje práci jak na straně CPU, tak na straně GPU. CPU odpovídá za plánování dávek a spouštění kernelů s příslušnými parametry, zatímco GPU provádí příslušné kernely násobení matic, pozornosti, normy nebo aktivace. U úloh s vysokou propustností, jako je trénování a reindexace, jsou režie na straně CPU zanedbatelné, protože velikost dávek i latence na straně GPU jsou velké. U menších velikostí dávek však práce na straně CPU může převážit nad prací na straně GPU.
Aby se zmírnily režijní náklady, lze namísto spouštění nezávislých kernelů sestavit graf CUDA, který zachytí metadata potřebná ke spuštění všech kernelů průchodu vpřed jediným voláním ovladače CUDA. Tím se eliminuje nutnost opětovného spouštění drahého kódu v Python a PyTorch pro konfigurace, pro které lze grafy CUDA zachytit.
U každého modelu sledujeme bod zvratu, který určuje minimální počet tokenů, při němž je spuštění na GPU dražší než spuštění kernelu na straně CPU. Vzhledem k tomu, že modely vkládání jsou malé, pozorujeme, že tento bod zvratu přichází u dávek čítajících tisíce tokenů a desítky sekvencí. Některé implementace pozornosti se spoléhají na dynamické vstupy na straně hostitele pro konfiguraci spouštění kernelů, což znemožňuje plné předběžné vyplnění modelu / husté grafy CUDA. Změny souvisejících kernelů jsme přenesli do upstreamu, abychom je mohli povolit v našem odvozovacím enginu.
K řešení režijních nákladů vytváříme grafy CUDA pro celé modely pro všechny modely vkládání a překrýváme práci CPU s prací GPU. Vzhledem k tomu, že grafy CUDA minimalizují režie na straně CPU, máme po spuštění grafu volný čas na zahájení a zařazení provádění další dávky do fronty, jakmile je k dispozici. Výsledky čekající dávky jsou sledovány pomocí LazyTensor, což umožňuje asynchronní úloze v jazyce Rust blokovat, dokud předchozí dávka nedokončí provádění. Grafy CUDA pomáhají obsluze s nízkou latencí tím, že zajišťují, že nás nezdržují náklady na spouštění kernelů, a usnadňují lepší plánování v případě vysoké propustnosti, protože uvolňují CPU, aby mohlo dříve pracovat na další dávce.
Grafy CUDA je nutné zachytit pro každou odlišnou konfiguraci, což pro vložení znamená graf pro každou kombinaci počtu sekvencí a počtu tokenů. Vzhledem k tomu, že tato mřížka je rozsáhlá, vyplňujeme počty tokenů do kbelíků, které jsou násobky 64 nebo 256. To stále vede k tisícům grafů, jejichž zachycení může u typického modelu trvat několik minut. Náklady na zachycení pocházejí ze dvou zdrojů: eager průchodu vpřed, který je nutné provést za účelem kompilace kernelů a nastavení vyrovnávacích pamětí pro různé kernely, které je potřebují, následovaného zachycovacím během, který znovu spouští kód v Python.
Náklady na spuštění zmírňujeme líným zachycováním grafů CUDA v průběhu obsluhy enginu. Sledujeme každou konfiguraci a zajišťujeme, aby před spuštěním zachycení a přehrání grafu prošla eager zahřívacím spuštěním při druhém zásahu. Všechna následná spuštění stejné konfigurace grafu pak procházejí přehráním grafu CUDA. Líné zachycování grafů má vliv na latence p99 během spouštění; je však cenné pro rozložení několika minut eager práce do několika hodin. Rychlejší doby spuštění nám umožňují lépe škálovat a spravovat nasazení vkládání.
Líné tenzory
Prostřednictvím CUDA je práce GPU asynchronní. Vzhledem k tomu, že spuštění kernelu jej asynchronně zařadí do fronty na streamu, musí kód hostitele explicitně synchronizovat, aby mohl přečíst výsledné vektory. Abychom usnadnili vyšší míru paralelismu a mohli spouštět budoucí dávky během čekání na dokončení předchozí dávky na zařízení, spoléháme se na abstrakci LazyTensor pro sledování hodnot.
LazyTensor sleduje vyrovnávací paměť hostitele v paměti uzamčené na stránkách a operaci cudaMemcpyAsync prostřednictvím události kopírující data ze zařízení. Spustí se po spuštění průchodu vpřed na stejném streamu. Vzhledem k tomu, že operace kopírování musí počkat na spuštění všech předchozích kernelů na streamu, sleduje přidružená událost jak dokončení průchodu vpřed, tak dostupnost výsledku na CPU.
V našem kódovacím enginu ROSE využíváme LazyTensory k překrytí práce GPU a CPU. Namísto toho, aby každé volání step() spouštělo graf CUDA a čekalo na jeho dokončení, vrací step() LazyTensor, který asynchronně sleduje svůj výsledek. Ve spojení s grafy CUDA nám to pomáhá dosahovat nízkých latencí a lepší propustnosti.
ROSE
Upravili jsme náš engine ROSE, který jsme původně vytvořili pro obsluhu LLM, tak, aby zvládal také provádění modelů vkládání. Abychom minimalizovali úsilí potřebné k podpoře modelů vkládání, ROSE agresivně sdílí kód mezi LLM a vloženími. Například obsluha pplx-embed a dekódování LLM Qwen3.5 procházejí stejnými kernely. Toto sdílení nám umožňuje snadno obsluhovat model vkládání, který byl původně doladěn z LLM pro účely prototypování, vyhodnocování a produkčního odvozování.
U hustých vrstev jsou vkládání a odvozování LLM totožné, protože vektory tokenů jsou zpracovávány nezávisle. Ve vrstvách pozornosti jsou rozdíly řešeny přidáním podpory pro roztříštěné vstupy spolu s nastavením stránkovaného předběžného vyplnění a dekódování požadovaným pro LLM. Při obsluze modelu vkládání nestencujeme mezipaměť KV a neodesíláme ji do variant kernelů pozornosti, které podporují roztříštěný formát, abychom se vyhnuli odsazení. Odpovídající konverzní a kalibrační rutiny jsou rovněž sdíleny s LLM.
Ivy
Na výkonu hraje důležitou roli také Ivy, naše vrstva HTTP proxy pro odvozování. Vzhledem k tomu, že datová zatížení požadavků se v produkčním prostředí liší, může směrování jednotlivých požadavků na jednotlivé repliky způsobit nerovnoměrné zatížení. Ivy rozděluje požadavky s velkými dávkami na menší části a vyvažuje jejich zatížení mezi replikami, čímž zlepšuje využití a vyhlazuje latenci. Naše nedávná práce na vlastní unigramové tokenizaci, plně zavedené v Ivy, radikálně zlepšuje latence oproti hotovým tokenizérům.
...ale kernely stále hrají roli
ROSE podporuje řadu backendů pozornosti. Různé kernely mohou vyhovovat určitým velikostech problémů. Postupem času jsme integrovali kernely FlashInfer 2, FlashInfer 3 a FlashAttention 4 pro implementaci roztříštěné pozornosti.
Obecně pozorujeme, že FlashAttention 4 je rychlejší. FlashInfer 3 ho však překonává u modelů založených na Qwen při velmi dlouhých délkách sekvencí. Jelikož se výkon a ladění mohou lišit podle počtu a dimenze hlaviček pozornosti, zachováváme podporu pro více konfigurovatelných možností a při obsluze se rozhodujeme případ od případu.
Benchmarky
Provádíme srovnání s vLLM v0.22.0, přičemž spouštíme odvozování v přesnosti BF16 na skutečných vahách modelů a vstupech odvozených z evaluačních datových sad. Všechna měření času byla předcházena zahřívacími spuštěními, která ověřila, že odchylka v kosínové podobnosti je v rozmezí 0,1 %.
Vložení s nízkou latencí (p50 / p90 / p99 / max ms)
Uvádíme runtimy pro předem tokenizovanou velikost dávky požadavků 1, plně sekvenční požadavky, délky sekvencí 128, 512 a 4096 tokenů.
Bodování s nízkou latencí (p50 / p90 / p99 / max ms)
Předem tokenizované velikosti dávek požadavků 5, 25 a 50, délka sekvence 512 tokenů.
Vložení s vysokou propustností (emb/s)
Velikost dávky požadavků 100, čtyři souběžné procesy odesílající požadavky, délky sekvencí 512, 1024 a 4096 tokenů.
Vložení s vysokou souběžností (p50 / p90 / p99 / max ms)
Délka sekvence 512, velikost dávky 1, ale odesíláme 1, 2, 4, 8 a 16 souběžných požadavků. Tento benchmark rovněž zahrnuje náklady na tokenizaci prostřednictvím Ivy a síťovou režii mezi Ivy a Tulip.
Závěr a budoucí práce
Obslužná infrastruktura složená z Ivy, Tulip a ROSE nám umožňuje obsluhovat vložení pro Perplexity s nižší latencí a lepší propustností, což má za následek přesnější vyhledávání za snížené náklady ve srovnání s hotovými řešeními.
Zaměřením se na konkrétní modely a převzetím odpovědnosti za celý zásobník získáváme svobodu nezbytnou k dosažení efektivní rovnováhy mezi výkonem a flexibilitou, kdy kombinujeme vysoce znovu použitelné a výkonné primitivní prvky jazyka Rust spolu s obecnějším kódem modelování v Python. Mnoho open-source odvozovacích enginů, jako jsou vLLM, SGLang a TokenSpeed, integruje do svého zásobníku jazyky jako Rust a C++. Do jazyka Rust jsme investovali během posledních dvou let a sklidili jsme velké úspěchy jak ve výkonu, tak v udržovatelnosti. Sdílením většiny implementace vkládání s naším zásobníkem pro obsluhu LLM získáváme rovněž nárůst propustnosti, aniž bychom museli vynakládat významné inženýrské úsilí na údržbu modelů vkládání.
S vývojem modelů budeme nadále vylepšovat každou vrstvu našeho zásobníku, abychom snížili latence vázané jak na CPU, tak na GPU. Naše vlastní protokoly založené na gRPC v rámci Ivy a Tulip nám umožňují vyladit komunikaci za účelem snížení síťových latencí, zatímco ROSE poskytuje základ pro zlepšení výpočetní propustnosti. Jakmile navíc v celém ekosystému poroste podpora pro Python bez zámku vláken (free-threaded), budeme moci dále vylepšit vzájemnou interoperabilitu mezi Python a Rust za účelem snížení režie.