Rýchle vkladania na GPU
Rýchle a presné vyhľadávanie je životne dôležité pre celú Perplexity, od vyhľadávania a počítača až po našu platformu API. V pozadí ťažkú prácu robia modely vkladávania a hodnotenia, ktoré pomáhajú našim systémom identifikovať najrelevantnejšie výsledky pre daný dopyt. Dosahujeme špičkovú kvalitu
Rýchle a presné vyhľadávanie je životne dôležité pre celú spoločnosť Perplexity, od vyhľadávania a počítača až po našu platformu API. V pozadí ťažkú prácu robia modely vkladávania a hodnotenia, ktoré pomáhajú našim systémom identifikovať najrelevantnejšie výsledky pre daný dopyt. Špičkovú kvalitu a latenciu dosahujeme trénovaním a poskytovaním vlastných modelov, ako je napríklad pplx-embed.
Tento článok prináša pohľad pod kapotu serverovej infraštruktúry spoločnosti Perplexity pre túto špeciálnu triedu modelov. Diskutujeme o našich technikách na efektivne riešenie potrieb inferencie vyhľadávania natívneho pre AI, ktoré umožňujú rýchle prototypovanie a vyhodnocovanie modelov a zároveň poháňajú náš vyhľadávací index v exabajtovom rozsahu. Tieto techniky spoločne rozširujú Paretovu hranicu kvality a efektívnosti vyhľadávania, čo nám umožňuje poskytovať agentom a používateľom tie najlepšie možné výsledky pri najnižších nákladoch a latencii.
Vkladania na účely vyhľadávania
V typickom nastavení vyhľadávania sa indexované dokumenty mapujú do vo vysokorozmerného vektorového priestoru pomocou modelu vkladania a ukladajú sa do vektorovej databázy. Vložením dopytu pomocou rovnakého modelu možno nájsť podobné dokumenty nájdením vektorov najbližších k vektoru dopytu. To vedie k dvom rôznym vzorcom prevádzky pre inferenčný engine, ktoré je potrebné obslúžiť:
- Dávkové vkladanie: pri vytváraní, rozširovaní alebo opätovnom indexovaní databázy sa hromadné dokumenty musia vkladať do vektorového priestoru, čím sa maximalizuje priepustnosť, aby sa minimalizovali náklady.
Po vyhľadávaní vektorov je potrebné ohodnotiť veľké dávky dokumentov, čím sa dosiahne rovnováha medzi priepustnosťou a latenciou.
- Online vkladanie: pri dopytovaní v databáze sa musí vložiť krátky dopyt na účely vyhľadávania, čím sa minimalizuje latencia.
Našu inferenčnú infraštruktúru sme vybudovali tak, aby v rôznych prípadoch použitia využívala čo najviac spoločných komponentov. Keďže na vytváranie vkladaní zvyčajne používame malé modely Transformer, väčšinu implementácie zdieľame s naším kódom inferencie LLM: dávkové vkladania sú podobné predbežnému vyplneniu viazanému na výpočet, zatiaľ čo online vkladania, ktoré často bežia na niekoľkých tokenoch, sú výpočtovo podobné dekódovaniu viazanému na pamäť. Na poskytovanie modelov vkladaní preto opätovne používame naše optimalizované predbežné vyplnenie a dekódovacie kernely. V dôsledku toho môžeme dosiahnuť obrovskú priepustnosť dávkovej inferencie s minimálnym dodatočným inžinierskym úsilím pri zachovaní nízkej latencie pre online pracovné zaťaženia vkladaní.
Tulipy, ruže a nejaký brečtan
Inferenciu sprístupňujeme prostredníctvom štandardizovaných API, a to interných aj externých, prostredníctvom našej platformy API. Pod kapotou sa na spracovaní požiadavky na vkladanie podieľa viacero služieb:
- Ivy je HTTP brána v Ruste, ktorú služby Perplexity volajú.
Zvláda prácu na strane CPU pre požiadavky, ako je analýza JSON, tokenizácia, tvorba šablón vstupov a rozdeľovanie dávok, pričom preklada požiadavky na vlastný protokol gRPC pre následné servery. Toto oddelenie nám umožňuje konfigurovať určité parametre týkajúce sa tokenizácie a formátovania vstupov bez toho, aby sme sa museli dotýkať ťažších inferenčných inštancií.
- Tulip je rozhranie inferenčného servera.
Je to server gRPC implementovaný v Ruste, tokio a `tonic`. Tulip prijíma požiadavky na inferenciu gRPC, pričom riadi plánovanie a dávkovanie. Následne posiela dávky do enginu ROSE a vracia dokončené odpovede klientom.
- **ROSE** (Runtime-Optimized Serving Engine) implementuje inferenciu modelov.
Je primárne definovaný v jazyku Python, pričom poskytuje kernely, vrstvy a definície pre širokú škálu modelov. ROSE implementuje dopredné prechody cez modely a zároveň poskytuje správu grafov CUDA špecializovanú na vkladania. S Tulipom je prepojený prostredníctvom funkcie step(), ktorá prevezme dávku a vráti referenciatu na výpočet, ktorý vykonáva na akcelerátore.
Venovanie pozornosti za hranice kernelu
Modely založené na transformeroch aj základné architektúry Hopper/Blackwell sú zrelé technológie, takže inferencia vkladania na strane GPU sa zlúčila do do značnej miery optimálnej implementácie naprieč rôznymi inferenčnými enginmi. Napriek tomu sme objavili ďalšie možnosti na zlepšenie v behových prostrediach a nástrojoch, ktoré sprístupňujú modely klientovi od začiatku do konca. Zistili sme najmä to, že dokážeme zlepšiť latencie starostlivou správou grafov CUDA a vytvorebím abstrakcie LazyTensor na asynchrónne sledovanie výsledku na strane GPU v natívnom engine Rust. Tieto funkcie sme implementovali v Tulipe, aby mohol efektívne spolupracovať s implementáciami modelov v ROSE.
Tulip
Tulip sme navrhli tak, aby bol čo najľahším rozhraním pre naše poskytovanie modelov. Spracováva prichádzajúce požiadavky v asynchrónnych úlohách Tokio, pričom udržiava fond požiadaviek, ktoré sleduje, a plánuje z nich dávky na odoslanie do akcelerátora. Mechanizmus plánovania v Tulipe je veľmi jednoduchý: požiadavky sa hromadia, zatiaľ čo Tulip odosiela prácu alebo čaká na výsledky. Z nahromadených požiadaviek sa sekvencie vyberajú na základe zásady „kto dřív přijde, ten dřív mele“ na spustenie cez model.
Jednoduchý mechanizmus plánovania je motivovaný pozorovaním výkonu modelu. Pre malé modely vkladaní pri dĺžkach sekvencií, ktoré obsluhujeme, sme si všimli, že lineárne náklady hustých vrstiev sú dominantné nad kvadratickými nákladmi pozornosti. Latencia je teda väčšinou úmerná počtu tokenov, nie počtu sekvencií. V dôsledku toho, akonáhle je dávka dostatočne veľká na nasýtenie GPU, čo je okolo 512 tokenov na modeli s menej ako miliardou parametrov, zabalenie viacerých sekvencií do nej nezlepšuje efektivitu.
Aby Tulip efektívne spolupracoval s modelom, spolieha sa na grafy CUDA a sledovanie lenivých výsledkov na prekrývanie práce GPU a CPU a plné využitie dostupných zdrojov.
Správa grafov CUDA
Spustenie dopredného prechodu modelu zahŕňa prácu na strane CPU aj GPU. CPU je zodpovedný za plánovanie dávok a spúšťanie kernelov s príslušnými parametrami, zatiaľ čo GPU vykonáva príslušné kernely násobenia matíc, pozornosti, normovania alebo aktivácie. Pri pracovných zaťaženiach s vysokou priepustnosťou, ako je trénovanie a opätovné indexovanie, sú režijné náklady na strane CPU zanedbateľné, pretože veľkosti dávok aj latencia na strane GPU sú veľké. Pri menších veľkostiach dávok však práca na strane CPU môže prevážiť prácu na strane GPU.
Na zmiernenie režijných nákladov možno namiesto spúšťania nezávislých kernelov vytvoriť graf CUDA na zachytenie metadát potrebných na spustenie všetkých kernelov dopredného prechodu jediným volaním ovládača CUDA. Tým sa eliminuje potreba opätovného spúšťania drahého kódu Pythonu a PyTorchu pre konfigurácie, pre ktoré možno zachytiť grafy CUDA.
Pri každom modeli sledujeme bod zvratu, ktorý určuje minimálny počet tokenov, pri ktorom je spúšťanie na GPU drahšie ako spúšťanie kernela na strane CPU. Keďže modely vkladaní sú malé, pozorujeme, že tento bod zvratu prichádza pri dávkach tisícok tokenov a desiatok sekvencií. Niektoré implementácie pozornosti sa spoliehajú na dynamické vstupy na strane hostiteľa na konfiguráciu spustenia kernelov, čomu bránia úplné grafy CUDA pre predbežné vyplnenie/husté modely. Do príslušných kernelov sme upstreamovali zmeny, aby sme ich mohli povoliť v našom inferenčnom engine.
Na riešenie režijných nákladov zostavujeme grafy CUDA pre celé modely pre všetky modely vkladaní a prekrývame prácu CPU s prácou GPU. Keďže grafy CUDA minimalizujú réžiu na strane CPU, akonáhle je graf spustený, máme voľný čas na spustenie a zaradenie vykonávania ďalšej dávky do frontu, kedykoľvek je k dispozícii. Výsledky čakajúcej dávky sa sledujú pomocou LazyTensor, čo umožňuje asynchrónnej úlohe v Ruste blokovať, kým predchádzajúca dávka nedokončí vykonávanie. Grafy CUDA pomáhajú pri poskytovaní s nízkou latenciou tým, že zabezpečujú, že nás nezdržiavajú náklady na spúšťanie kernelov, a uľahčujú lepšie plánovanie v prípade vysokej priepustnosti, pretože uvoľňujú CPU, aby mohol skôr pracovať na ďalšej dávke.
Grafy CUDA sa musia zachytiť pre každú odlišnú konfiguráciu, čo v prípade vkladaní znamená jeden graf na kombináciu počtu sekvencií a počtu tokenov. Keďže táto mriežka je rozsiahla, zarovnávame počty tokenov na bloky, ktoré sú násobkami 64 alebo 256. To stále vedie k tisíckam grafov, ktorých zachytenie pre typický model môže trvať niekoľko minút. Náklady na zachytenie pochádzajú z dvoch zdrojov: netrpezlivý dopredný prechod, ktorý sa musí vykonať na skompilovanie kernelov a nastavenie vyrovnávacích pamätí pre rôzne kernely, ktoré ich potrebujú, po ktorom nasleduje zachytávací beh, ktorý opätovne spúšťa kód v Pythone.
Náklady na spustenie zmierňujeme lenivým zachytávaním grafov CUDA počas obsluhy enginu. Sledujeme každú konfiguráciu a zaisťujeme, že prejde zahrievacím behom pred spustením zachytávania a prehrávania grafu pri druhom výskyte. Všetky následné spustenia rovnakej konfigurácie grafu potom prebiehajú prostredníctvom prehrávania grafu CUDA. Lenivé zachytávanie grafov má vplyv na latencie p99 počas spúšťania; je však cenné pri rozložení niekoľkých minút netrpezlivej práce do viacerých hodín. Rýchlejšie časy spustenia nám umožňujú lepšie škálovať a spravovať nasadenia vkladaní.
Lenivé tenzory
Prostredníctvom CUDA je práca GPU asynchrónna. Keďže asynchrónne spustenie kernelu ho zaradí do frontu prúdu, kód hostiteľa sa musí explicitne synchronizovať, aby prečítal výsledné vektory. Na uľahčenie vyššieho stupňa paralelizmu a schopnosti spúšťať budúce dávky počas čakajúcich na dokončenie predchádzajúcej na zariadení sa spoliehame na abstrakciu LazyTensor na sledovanie hodnôt.
LazyTensor sleduje vyrovnávaciu pamäť hostiteľa v uzamknutej pamäti stránky a operáciu cudaMemcpyAsync prostredníctvom udalosti kopírujúcej dáta zo zariadenia. Spúšťa sa po spustení dopredného prechodu na rovnakom prúde. Keďže operácia kopírovania musí počkať na vykonanie všetkých predchádzajúcich kernelov v prúde, pridružená udalosť sleduje dokončenie dopredného prechodu aj dostupnosť výsledku na CPU.
V našom kódovacom engine ROSE využívame LazyTensory na prekrývanie práce GPU a CPU. Namiesto toho, aby každé volanie step() spustilo graf CUDA a čakalo na jeho dokončenie, step() vráti LazyTensor na asynchrónne sledovanie svojho výsledku. V spojení s grafmi CUDA nám to pomáha dosiahnuť nízke latencie a lepšiu priepustnosť.
ROSE
Upravili sme náš engine ROSE, ktorý sme pôvodne vytvorili na poskytovanie LLM, tak, aby zvládal aj spúšťanie modelov vkladaní. Aby sa minimalizovalo úsilie potrebné na podporu modelov vkladaní, ROSE agresívne opätovne používa kód medzi LLM a vkladaniami. Napríklad poskytovanie pplx-embed a dekódovanie LLM Qwen3.5 prechádzajú cez rovnaké kernely. Toto zdieľanie nám umožňuje ľahko poskytovať model vkladaní, ktorý bol pôvodne jemne vyladený z LLM na prototypovanie, vyhodnocovanie a produkčnú inferenciu.
Pre husté vrstvy sú vkladanie a inferencia LLM totožné, pretože vektory tokenov sa spracovávajú nezávisle. Vo vrstvách pozornosti sa rozdiely riešia pridaním podpory pre neúplné vstupy spolu s nastaveniami stránkovaného predbežného vyplnenia a dekódovania vyžadovanými modelmi LLM. Pri poskytovaní modelu vkladaní neinštancujeme vyrovnávaciu pamäť KV a neodosielame do variantov kernelov pozornosti, ktoré podporujú neúplný formát, aby sme sa vyhli vypchávaniu. Podporné konverzné a kalibračné rutiny sa taktiež zdieľajú s modelmi LLM.
Ivy
Ivy, naša vrstva HTTP proxy pre inferenciu, tiež zohráva dôležitú úlohu vo výkone. Keďže sa užívateľské zaťaženia požiadaviek vo výrobe líšia, smerovanie jednotlivých požiadaviek na jednotlivé repliky môže spôsobiť nerovnováhu záťaže. Ivy rozdeľuje požiadavky s veľkými dávkami na časti a vyvažuje ich zaťaženie medzi replikami, čím zlepšuje využitie a vyhladzuje latenciu. Naša nedávna práca na vlastnej unigramovej tokenizácii, ktorá je úplne nasadená v Ivy, drasticky zlepšuje latencie v porovnaní s bežnými tokenizátormi.
...ale kernely sú stále dôležité
ROSE podporuje rôzne backendy pozornosti. Rôzne kernely môžu byť vhodné pre špecifické veľkosti problémov. Časom sme integrovali kernely FlashInfer 2, FlashInfer 3 a FlashAttention 4 na implementáciu neúplnej pozornosti.
Všeobecne pozorujeme, že FlashAttention 4 je rýchlejší. FlashInfer 3 však prekonáva FlashAttention pri modeloch založených na Qwen pri veľmi dĺžkach sekvencií. Keďže výkon a ladenie sa môžu líšiť v závislosti od počtu a rozmeru hláv pozornosti, zachovávame podporu pre viaceré konfigurácie a pri poskytovaní služieb sa rozhodujeme od prípadu k prípadu.
Benchmarky
Porovnávame sa s vLLM v0.22.0, pričom spúšťame inferenciu v presnosti BF16 na skutočných váhach modelov a vstupoch odvodených z evaluačných dátových setov. Všetkým načasovanimám predchádzali zahrievacie behy, ktoré overili, že odchýlka v kosínusovej podobnosti je v rozmedzí 0,1 %.
Vkladania s nízkou latenciou (p50 / p90 / p99 / max ms)
Uvádzame časy behov pre veľkosť dávky vopred tokenizovaných požiadaviek 1, úplne sekvenčné požiadavky, dĺžky sekvencií 128, 512 a 4096 tokenov.
Bodovanie s nízkou latenciou (p50 / p90 / p99 / max ms)
Veľkosti dávok vopred tokenizovaných požiadaviek 5, 25 a 50, dĺžka sekvencie 512 tokenov.
Vkladania s vysokou priepustnosťou (emb/s)
Veľkosť dávky požiadaviek 100, štyri súbežné procesy odosielajúce požiadavky, dĺžky sekvencií 512, 1024 a 4096 tokenov.
Vkladania s vysokou súbežnosťou (p50 / p90 / p99 / max ms)
Dĺžka sekvencie 512, veľkosť dávky 1, ale posielame 1, 2, 4, 8 a 16 súbežných požiadaviek. Tento benchmark zahŕňa aj náklady na tokenizáciu prostredníctvom Ivy spolu so sieťovou réžiou medzi Ivy a Tulipom.
Záver a budúca práca
Serverová infraštruktúra zložená z Ivy, Tulipu a ROSE nám umožňuje poskytovať vkladania pre Perplexity s nižšou latenciou a lepšou priepustnosťou, čoho výsledkom je presnejšie vyhľadávanie za znížené náklady v porovnaní s hotovými riešeniami.
Zameraním sa na špecifické modely a prevzatím kontroly nad celým zásobníkom získavame slobodu potrebnú na dosiahnutie efektívnej rovnováhy medzi výkonom a flexibilitou, pričom kombinujeme vysoko opakovane použiteľné a výkonné primitíva jazyka Rust spolu s generickejším kódom modelovania v Pythone. Mnoho open-source inferenčných enginov, ako napríklad vLLM, SGLang a TokenSpeed, integruje do svojho zásobníka jazyky ako Rust a C++. Do jazyka Rust sme investovali za posledné dva roky a zožali sme veľké úspechy z hľadiska výkonu aj udržiavateľnosti. Zdieľaním väčšiny implementácie vkladaní s naším zásobníkom poskytovania LLM získavame aj nárast priepustnosti bez toho, aby sme museli vynaložiť značné inžinierske úsilie na údržbu modelov vkladaní.
S vývojom modelov budeme aj naďalej vylepšovať každú vrstvu nášho zásobníka, aby sme znížili latencie viazané na CPU aj GPU. Naše vlastné protokoly založené na gRPC v rámci Ivy a Tulipu nám umožňujú ladiť komunikáciu s cieľom znížiť sieťové latencie, zatiaľ čo ROSE poskytuje základ na zlepšenie výpočtovej priepustnosti. Okrem toho, s tým, ako rastie podpora pre Python bez vláknových zámkov v celom ekosystéme, budeme môcť ďalej zlepšovať interoperabilitu medzi Pythonom a Rustom s cieľom znížiť réžiu.