Gyors beágyazások GPU-kon
A gyors és pontos keresés létfontosságú a Perplexity egésze számára, a Search-től és a Computer-től kezdve az API Platformunkig. A kulisszák mögött a nehéz munkát beágyazási és rangsorolási modellek végzik, amelyek segítik a rendszereinket az adott lekérdezéshez leginkább releváns eredmények azonosításában. State-of-the
A gyors és pontos keresés létfontosságú a Perplexity egésze számára, a Search-től és a Computer-től kezdve az API Platformunkig. A kulisszák mögött a nehéz munkát beágyazási és rangsorolási modellek végzik, amelyek segítik a rendszereinket az adott lekérdezéshez leginkább releváns eredmények azonosításában. State-of-the-art minőséget és latenciát érünk el saját modellek, például a pplx-embed betanításával és kiszolgálásával.
Ez a cikk a Perplexity ezen speciális modellosztályhoz tartozó kiszolgálási infrastruktúrájának motorháztető alatti nézetét mutatja be. Megvitatjuk azokat a technikáinkat, amelyekkel hatékonyan kezeljük az AI-alapú keresés inferenciaigényeit, lehetővé téve a modellek gyors prototípuskészítését és értékelését, miközben az exabájt skálájú keresési indexünket működtetjük. Ezek a technikák együttesen kiterjesztik a keresési minőség és hatékonyság Pareto-határát, lehetővé téve számunkra, hogy a legjobb lehetséges eredményeket szolgáljuk ki az ügynököknek és a felhasználóknak a legalacsonyabb költséggel és latenciával.
Beágyazások a kereséshez
Egy tipikus keresési beállításban az indexelt dokumentumokat egy nagy dimenziójú vektortérbe képezik le egy beágyazási modell segítségével, és egy vektoradatbázisban tárolják. Egy lekérdezés ugyanezzel a modellnel történő beágyazásával a hasonló dokumentumok megtalálhatók a lekérdezés vektorához legközelebbi vektorok megkeresésével. Ez két különböző forgalmi mintát eredményez egy inferenciamotor kiszolgálásához:
- Kötegelt beágyazás: az adatbázis felépítése, bővítése vagy újraindexelése során a tömeges dokumentumokat be kell ágyazni a vektortérbe, maximalizálva az áteresztőképességet a költségek minimalizálása érdekében.
A vektorkeresést követően dokumentumok nagy kötegeit kell pontozni, egyensúlyt teremtve az áteresztőképesség és a latencia között.
- Online beágyazás: az adatbázis lekérdezésekor egy rövid lekérdezést be kell ágyazni a keresésekhez, minimalizálva a latenciát.
Inferencia-infrastruktúránkat úgy építettük ki, hogy a felhasználási esetek között a lehető legtöbb közös komponenst használja ki. Mivel a beágyazások előállításához jellemzően kis Transformer modelleket használunk, a megvalósítás nagy részét megosztjuk az LLM-inferencia kódunkkal: a kötegelt beágyazások hasonlítanak a számításigényes prefillhez, míg az online beágyazások, amelyek gyakran néhány tokenen futnak, számítási szempontból a memóriakorlátos decode-hoz hasonlítanak. Így a beágyazási modellek kiszolgálásához újrafelhasználjuk az optimalizált prefill és decode kerneleket. Ennek eredményeként hatalmas kötegelt inferencia-áteresztőképességet érhetünk el minimális további mérnöki munkával, miközben fenntartjuk az alacsony latenciát az online beágyazási munkaterhelések számára.
Tulipák, rózsák és némi Ivy
Az inferenciát szabványosított API-kon keresztül tesszük elérhetővé, mind belsőleg, mind külsőleg az API Platformunkon keresztül. A motorháztető alatt több szolgáltatás vesz részt egy beágyazási kérés feldolgozásában:
- Az Ivy egy olyan Rust HTTP-átjáró (gateway), amelyet a Perplexity szolgáltatásai hívnak.
Kezeli a kérések CPU-oldali munkáját, mint például a JSON-elemzést, a tokenizációt, a bemeneti sablonkészítést és a kötegfelosztást, és a kéréseket egyedi gRPC protokollra fordítja le az alsó rétegbeli (downstream) szerverek számára. Ez a szétválasztás lehetővé teszi számunkra, hogy bizonyos paramétereket konfiguráljunk a tokenizáció és a bemeneti formázás körül anélkül, hogy hozzá kellene nyúlnunk a nehezebb inferencia-példányokhoz.
- A Tulip az inferencia-szerver interfész.
Ez egy Rusttal, tokio-val és `tonic`-kal implementált gRPC-szerver. A Tulip gRPC inferenciakéréseket fogad, kezelve az ütemezést és a kötegelést. Ezután elküldi a kötegeket a ROSE motornak, és a kész válaszokat visszaküldi a klienseknek.
Elsősorban Pythonban van definiálva, kerneleket, rétegeket és definíciókat biztosít a modellek széles skálájához. A ROSE valósítja meg a modelleken áthaladó előremenő (forward) fázisokat, emellett a beágyazásokra szakosodott CUDA-gráf-kezelést is biztosít. A Tuliphez egy step() függvény kötözi be, amely vesz egy köteget, és visszaadja az akcelerátoron végzett számítás referenciáját.
Figyelem a kernelen túl
Mind a Transformer-alapú modellek, mind az alapul szolgáló Hopper/Blackwell architektúrák érett technológiák, így a GPU-oldali beágyazási inferencia nagyrészt optimális megvalósításba konvergált a különböző inferencia-motorok között. Ennek ellenére további fejlesztési lehetőségeket fedeztünk fel a futtatókörnyezetekben és a keretrendszerekben, amelyek end-to-end teszik ki a modelleket egy kliens számára. Különösen azt találtuk, hogy javíthatjuk a latenciákat a CUDA-gráfok gondos kezelésével és egy LazyTensor absztrakció felépítésével, amellyel aszinkron módon követhető egy GPU-oldali eredmény a natív Rust motorban. Ezeket a funkciókat a Tulipben valósítottuk meg, így az hatékonyan tudott együttműködni a ROSE modellmegvalósításaival.
Tulip
A Tulipet úgy terveztük, hogy a modellkiszolgálás felett a lehető legkönnyebb interfész legyen. A bejövő kéréseket Tokio aszinkron feladatokban kezeli, fenntartva a kérések egy poolját, amelyet nyomon követ, és ütemezi a kötegeket az akcelerátorra történő elküldéshez. A Tulipben az ütemezési mechanizmus nagyon egyszerű: a kérések felhalmozódnak, amíg a Tulip munkát küld ki, vagy az eredményekre vár. A felhalmozott kérések közül a szekvenciákat érkezési sorrendben (first-come, first-served) választják ki, hogy átfuttassák a modellen.
Az egyszerű ütemezési mechanizmust a modell teljesítményével kapcsolatos megfigyelés motiválja. Kis beágyazási modellek esetében, azon szekvenciahosszak mellett, amelyeket kiszolgálunk, észrevettük, hogy a sűrű rétegek lineáris költsége dominál az attention kvadratikus költségével szemben. Így a latencia nagyrészt arányos a tokenek számával, nem pedig a szekvenciák számával. Következésképpen, ha egy köteg elég nagy ahhoz, hogy telítse a GPU-t (ami egy egymilliárd paraméter alatti modellen körülbelül 512 token), akkor több szekvencia beletömörítése nem javítja a hatékonyságot.
A modellel való hatékony együttműködés érdekében a Tulip a CUDA-gráfokra és a lusta eredménykövetésre támaszkodik a GPU- és CPU-munkák átfedésére, valamint a rendelkezésre álló erőforrások teljes kihasználására.
CUDA-gráf kezelés
A modell előremenő (forward) fázisának futtatása CPU- és GPU-oldali munkát egyaránt magában foglal. A CPU felelős a kötegek ütemezéséért és a kernelek megfelelő paraméterekkel történő indításáért, míg a GPU futtatja a releváns mátrixszorzási, attention-, norm- vagy aktiválási kerneleket. A nagy áteresztőképességű munkaterhelések, például a betanítás és az újraindexelés esetén a CPU-oldali overheadok elhanyagolhatóak, mivel a kötegméretek és a GPU-oldali latencia egyaránt nagyok. Kisebb kötegméreteknél azonban a CPU-oldali munka felülmúlhatja a GPU-oldali munkát.
Az overheadok mérséklésére a független kernelek indítása helyett építhető egy CUDA-gráf, amely rögzíti az előremenő fázis összes kernelének indításához szükséges metaadatokat a CUDA-illesztőprogram (driver) egyetlen hívásával. Ez kiküszöböli a drága Python- és PyTorch-kódok újrafuttatásának szükségességét azon konfigurációk esetében, amelyekhez rögzíthetők a CUDA-gráfok.
Minden egyes modell esetében nyomon követünk egy inflexiós pontot, amely meghatározza a tokenek minimális számát, amelynél a GPU-futtatás drágább, mint a CPU-oldali kernelindítás. Mivel a beágyazási modellek kicsik, azt tapasztaljuk, hogy ez az inflexiós pont ezer tokenes és tucatnyi szekvenciás kötegeknél következik be. Bizonyos attention-megvalósítások dinamikus gazdagép-oldali (host-side) bemenetekre támaszkodnak a kernelindítások konfigurálásához, ami megakadályozza a teljes modell prefill/sűrű CUDA-gráfjait. Változtatásokat upstreameltünk a releváns kerneleken, hogy lehetővé tegyük azokat az inferencia-motorunkban.
Az overheadok kezelésére teljes modellre vonatkozó CUDA-gráfokat építünk az összes beágyazási modellhez, és átfedésbe hozzuk a CPU-munkát a GPU-munkával. Mivel a CUDA-gráfok minimalizálják a CPU-oldali overheadokat, a gráf elindítása után szabadidőnk van arra, hogy elindítsuk és sorba állítsuk a következő köteg végrehajtását, amint az elérhetővé válik. A függőben lévő köteg eredményeit egy LazyTensor követi nyomon, amely lehetővé teszi egy Rustbeli aszinkron feladat számára, hogy blokkoljon, amíg az előző köteg befejezi a végrehajtást. A CUDA-gráfok segítik az alacsony latenciájú kiszolgálást azáltal, hogy biztosítják, hogy ne fogjanak vissza minket a kernelindítások költségei, és megkönnyítik a jobb ütemezést a nagy áteresztőképességű esetben, mivel felszabadítják a CPU-t, hogy hamarabb dolgozhasson a következő kötegen.
A CUDA-gráfokat minden egyes különálló konfigurációhoz rögzíteni kell, ami beágyazások esetén szekvencia- és tokenszám-kombinációnként egy gráfot jelent. Mivel ez a rács tágas, a tokenszámokat 64 vagy 256 többszöröseinek megfelelő bucketekbe paddingeljük (párnázuk). Ez még így is több ezer gráfot eredményez, amelyek rögzítése egy tipikus modell esetében több percet is igénybe vehet. A rögzítés költsége két forrásból származik: egy eager forward fázisból, amelyet le kell futtatni a kernelek lefordításához és a pufferek beállításához a rájuk utalt kernelek számára, amit a Python-kódot újra végrehajtó rögzítési (capture) futás követ.
Az indítási költségeket úgy mérsékeljük, hogy lusta módon (lazily) rögzítjük a CUDA-gráfokat, ahogy a motor kiszolgál. Nyomon követjük az egyes konfigurációkat, és biztosítjuk, hogy egy eager bemelegítési futáson menjenek keresztül, mielőtt elindítanánk a gráfrögzítést és -lejátszást a második találatkor. Ugyanazon gráfkonfiguráció összes további végrehajtása ezután a CUDA-gráf lejátszásán megy keresztül. A lusta gráfrögzítés hatással van az indulás közbeni p99 latenciákra; azonban értékes abban, hogy több percnyi eager munkát több órára terít szét. A gyorsabb indítási idők lehetővé teszik számunkra a beágyazási deployments jobb skálázását és kezelését.
Lusta tenzorok (Lazy Tensors)
A CUDA-n keresztül a GPU-munka aszinkron. Mivel egy kernel aszinkron indítása beejti azt egy streambe, a hoszt kódnak explicit módon szinkronizálnia kell az eredményül kapott vektorok kiolvasásához. A magasabb fokú párhuzamosság elősegítése, valamint az, hogy a leendő kötegek elindíthatók legyenek, amíg arra várunk, hogy az előző befejeződjön az eszközön, egy LazyTensor absztrakcióra támaszkodunk az értékek nyomon követéséhez.
A LazyTensor nyomon követ egy oldalra rögzített (page-locked) memóriában lévő hoszt puffert és egy cudaMemcpyAsync műveletet az eszközről adatokat másoló eseményen keresztül. Ez az előremenő fázis indítása után indul el ugyanazon a streamen. Mivel a másolási műveletnek meg kell várnia a streamen lévő összes korábbi kernel végrehajtását, a kapcsolódó esemény nyomon követi mind az előremenő fázis befejeződését, mind az eredmény elérhetőségét a CPU-n.
A LazyTensor-okat a ROSE enkódermotorunkban használjuk a GPU- és CPU-munkák átfedésére. Ahelyett, hogy minden egyes step() hívás futtatná a CUDA-gráfot és megvárná annak befejeződését, a step() egy LazyTensor-t ad vissza az eredmény aszinkron nyomon követéséhez. A CUDA-gráfokkal párosulva ez segít az alacsony latencia és a jobb áteresztőképesség elérésében.
ROSE
Adaptáltuk a ROSE motorunkat, amelyet eredetileg LLM-kiszolgáláshoz építettünk, hogy a beágyazási modellek futtatását is kezelje. A beágyazási modellek támogatásához szükséges erőfeszítések minimalizálása érdekében a ROSE agresszíven újrafelhasználja a kódot az LLM-ek és a beágyazások között. Például a pplx-embed kiszolgálás és a Qwen3.5 LLM dekódolás mind ugyanazon kerneleken megy keresztül. Ez a megosztás lehetővé teszi számunkra, hogy könnyen kiszolgáljunk egy olyan beágyazási modellt, amelyet eredetileg egy LLM-ből finomhangoltak prototípuskészítéshez, értékeléshez és gyártási inferenciához.
Sűrű rétegek esetén a beágyazás és az LLM-inferencia azonos, mivel a tokenvektorokat függetlenül dolgozzák fel. Az attention rétegekben a különbségeket a rongyos (ragged) bemenetek támogatásának hozzáadásával kezeljük, az LLM-ek által megkövetelt lapozott (paged) prefill és decode beállítások mellett. Egy beágyazási modell kiszolgálásakor nem példányosítunk KV-gyorsítótárat, és az attention kernelek azon változataira irányítjuk a kérést, amelyek támogatják a rongyos formátumot a padding elkerülése érdekében. A támogató konverziós és kalibrációs rutinok szintén meg vannak osztva az LLM-ekkel.
Ivy
Az Ivy, az inferencia HTTP-proxy rétegünk szintén fontos szerepet játszik a teljesítményben. Mivel a kérés-payloadok a gyártás során változnak, az egyes kérések egyes replikákra történő routingja terhelési egyensúlyhiányt okozhat. Az Ivy a nagy köteges kérések darabokra osztja, és terheléselosztást végez a replikák között, javítva a kihasználtságot és simítva a latenciát. Az on-premise unigram tokenizációval kapcsolatos, az Ivyban teljesen bevezetett legújabb munkánk drasztikusan javítja a latenciákat a készen kapott tokenizálókhoz képest.
...de a kernelek továbbra is számítanak
A ROSE különféle attention backendeket támogat. Különböző kernelek lehetnek alkalmasak specifikus problémaméretekhez. Idővel integráltuk a FlashInfer 2, FlashInfer 3 és FlashAttention 4 kerneleket a rongyos (ragged) attention megvalósításához.
Általában azt tapasztaljuk, hogy a FlashAttention 4 gyorsabb. A FlashInfer 3 azonban túlszárnyalja a Qwen-alapú modelleken, nagyon hosszú szekvenciahosszak esetén. Mivel a teljesítmény és a hangolás változhat az attention fejek számától és dimenziójától függően, több konfiguráció támogatását fenntartjuk, és eseti döntést hozunk a kiszolgálás során.
Benchmarkok
A vLLM v0.22.0 verziójához viszonyítva benchmarkolunk, BF16 pontossággal futtatva az inferenciát a tényleges modell súlyokon és az értékelési adatkészletekből származó bemeneteken. Minden időzítési futás előtt bemelegítési futások voltak, amelyek ellenőrizték, hogy a koszinusz-hasonlóság eltérése 0,1%-on belül van.
Alacsony latenciájú beágyazások (p50 / p90 / p99 / max ms)
Futtatási időket jelentünk az előre tokenizált 1-es kérés kötegelt mérethez, teljesen szekvenciális kérésekhez, 128, 512 és 4096 tokenes szekvenciahosszal.
Alacsony latenciájú pontozás (p50 / p90 / p99 / max ms)
Előre tokenizált kérés kötegelt méretei: 5, 25 és 50, 512 tokenes szekvenciahosszal.
Nagy áteresztőképességű beágyazások (emb/s)
Kérés kötegelt mérete: 100, négy párhuzamos folyamat küld kéréseket, 512, 1024 és 4096 tokenes szekvenciahosszal.
Nagy egyidejűségű beágyazások (p50 / p90 / p99 / max ms)
512-es szekvenciahossz, 1-es kötegméret, de 1, 2, 4, 8 és 16 párhuzamos kérést küldünk. Ez a benchmark magában foglalja az Ivy általi tokenizálási költségeket is, valamint az Ivy és a Tulip közötti hálózati overheadet.
Következtetés és jövőbeli munka
Az Ivyból, a Tulipből és a ROSE-ból álló kiszolgálási infrastruktúra lehetővé teszi számunkra, hogy alacsonyabb latenciával és jobb áteresztőképességgel szolgáljunk ki beágyazásokat a Perplexity számára, ami pontosabb keresést eredményez csökkentett költségek mellett a kész megoldásokhoz képest.
Azáltal, hogy specifikus modellekre fókuszálunk és az egész stacket felügyeljük, megkapjuk azt a szabadságot, amely szükséges a teljesítmény és a rugalmasság hatékony egyensúlyának megteremtéséhez, összekeverve a nagymértékben újrafelhasználható és teljesítményképes Rust primitíveket az általánosabb Python modellezési kóddal. Számos nyílt forráskódú inferencia-motor, mint például a vLLM, a SGLang és a TokenSpeed, beépíti a Rust és C++ nyelveket a stackjébe. Az elmúlt két évben a Rustba fektettünk, és nagyszerű eredményeket arattunk mind a teljesítmény, mind a karbantarthatóság terén. Azáltal, hogy a beágyazási megvalósítás nagy részét megosztjuk az LLM-kiszolgáló stackünkkel, áteresztőképesség-növekedést is elérünk anélkül, hogy jelentős mérnöki erőfeszítést kellene fordítanunk a beágyazási modellek karbantartására.
Ahogy a modellek fejlődnek, továbbra is javítani fogjuk a stackünk minden rétegét, hogy csökkentsük mind a CPU-, mind a GPU-kötött latenciákat. Az Ivy és a Tulip belsejében lévő egyedi gRPC-alapú protokolljaink lehetővé teszik a kommunikáció finomhangolását a hálózati latenciák csökkentése érdekében, míg a ROSE alapot biztosít a számítási áteresztőképesség javításához. Emellett, ahogy az ökoszisztémában növekszik a free-threaded Python támogatása, képesek leszünk tovább javítani a Python-Rust interoperabilitást az overheadok csökkentése érdekében.