Încorporări rapide pe GPU-uri

Căutarea rapidă și precisă este vitală pentru tot ceea ce ține de Perplexity, de la Search și Computer până la Platforma noastră API. În culise, munca grea este făcută de modelele de încorporare și clasificare, care ajută sistemele noastre să identifice cele mai relevante rezultate pentru o anumită interogare. Obținem starea de

AutoriPerplexity Engineering

Căutarea rapidă și precisă este vitală pentru tot ceea ce ține de Perplexity, de la Search și Computer până la Platforma noastră API. În culise, munca grea este făcută de modelele de încorporare și clasificare, care ajută sistemele noastre să identifice cele mai relevante rezultate pentru o anumită interogare. Obținem o calitate și o latență de ultimă generație prin antrenarea și servirea propriilor noastre modele, cum ar fi pplx-embed.

Acest articol prezintă o perspectivă din culise asupra infrastructurii de servire a Perplexity pentru această clasă specială de modele. Discutăm tehnicile noastre de a aborda eficient nevoile de inferență ale căutării native AI, permițând prototiparea și evaluarea rapidă a modelelor în timp ce alimentăm indexul nostru de căutare la scară de exabiți. Aceste tehnici extind în mod colectiv frontiera Pareto a calității și eficienței căutării, permițându-ne să deservim agenții și utilizatorii cu cele mai bune rezultate posibile la cel mai mic cost și latență.

Încorporări pentru căutare

Într-o configurație tipică de căutare, documentele indexate sunt mapate într-un spațiu vectorial de înaltă dimensiune utilizând un model de încorporare și stocate într-o bază de date vectorială. Prin încorporarea unei interogări folosind același model, documentele similare pot fi localizate găsind vectorii cei mai apropiați de cel al interogării. Acest lucru dă naștere la două modele de trafic diferite pentru un motor de inferență pe care să le deservească:

  • Încorporare în lot: la construirea, extinderea sau reindexarea bazei de date, documentele în masă trebuie încorporate în spațiul vectorial, maximizând debitul pentru a minimiza costul.

În urma căutării vectoriale, loturi mari de documente trebuie scorate, stabilind un echilibru între debit și latență.

  • Încorporare online: la interogarea bazei de date, o interogare scurtă trebuie încorporată pentru căutări, minimizând latența.

Ne-am construit infrastructura de inferență pentru a valorifica cât mai multe componente comune posibile între cazurile de utilizare. Deoarece folosim de obicei modele Transformer mici pentru a produce încorporări, împărțim cea mai mare parte a implementării cu codul nostru de inferență LLM: încorporările în lot sunt similare cu preumplerea legată de calcul, în timp ce încorporările online, care rulează adesea pe câteva tokenuri, sunt computațional similare cu decodarea legată de memorie. Prin urmare, reutilizăm kernel-urile noastre optimizate de preumplere și decodare pentru a servi modele de încorporare. Ca rezultat, putem obține un debit masiv de inferență în lot cu un efort de inginerie suplimentar minim, menținând în același timp o latență scăzută pentru sarcinile de lucru de încorporare online.

Tulipe, Roze și niște Iederă

Expunem inferența prin API-uri standardizate, atât intern, cât și extern prin Platforma noastră API. În culise, mai multe servicii sunt implicate în procesarea unei cereri de încorporare:

  • Ivy este un gateway HTTP Rust pe care serviciile Perplexity îl apelează.

Se ocupă de munca pe partea de CPU pentru cereri precum parsarea JSON, tokenizarea, crearea de șabloane de intrare și împărțirea loturilor, traducând cererile într-un protocol gRPC personalizat pentru serverele din aval. Această separare ne permite să configurăm anumit parametri referitori la tokenizare și formatarea intrărilor fără a fi nevoie să atingem instanțele de inferență mai grele.

  • Tulip este interfața serverului de inferență.

Este un server gRPC implementat cu Rust, tokio și `tonic`. Tulip primește cereri de inferență gRPC, gestionând programarea și loturile. Apoi trimite loturile către motorul ROSE, returnând răspunsurile finalizate către clienți.

  • **ROSE** (Runtime-Optimized Serving Engine) implementează inferența de modele.

Este definit în principal în Python, oferind kernel-uri, straturi și definiții pentru o mare varietate de modele. ROSE implementează trecerile înainte prin modele, oferind de asemenea o gestionare a grafurilor CUDA specializată pentru încorporări. Este conectat la Tulip prin intermediul unei funcții step(), care preia un lot și returnează o referință la calculul pe care îl efectuează pe accelerator.

Arhitectura de servire de la o cerere de încorporare prin Ivy către servere Tulip replicate

Acordarea atenției dincolo de kernel

Atât modelele bazate pe Transformer, cât și arhitecturile Hopper/Blackwell subiacente sunt tehnologii mature, astfel încât inferența încorporată pe partea de GPU a ajuns la o implementare în mare parte optimă pentru diverse motoare de inferență. Chiar și așa, am descoperit oportunități suplimentare de îmbunătățire a timpurilor de rulare și a interfețelor care expun modelele end-to-end către un client. În particular, am constatat că putem îmbunătăți latențele gestionând cu atenție grafurile CUDA și construind o abstracție LazyTensor pentru a urmări asincron un rezultat pe partea de GPU în motorul nativ Rust. Am implementat aceste caracteristici în Tulip, astfel încât să poată interacționa eficient cu implementările de modele ale ROSE.

Tulip

Am conceput Tulip pentru a fi o interfață cât mai ușoară posibil asupra servirii modelelor noastre. Gestionează cererile primite în sarcini asincrone Tokio, menținând un grup de cereri pe care le urmărește și din care programează loturi pentru a le expedia către accelerator. Mecanismul de programare din Tulip este foarte simplu: cererile se acumulează în timp ce Tulip expediază munca sau așteaptă rezultatele. Din cererile acumulate, secvențele sunt alese pe principiul primul venit, primul servit pentru a fi rulate prin model.

Mecanismul simplu de programare este motivat de o observație privind performanța modelului. Pentru modelele mici de încorporare, la lungimile de secvență pe care le deservim, am observat că costul liniar al straturilor dense este dominant față de costul pătratic al atenției. Astfel, latența este în mare parte proporțională cu numărul de tokenuri, nu cu numărul de secvențe. În consecință, odată ce un lot este suficient de mare pentru a satura GPU-ul, ceea ce înseamnă în jur de 512 tokenuri pe un model de sub un miliard de parametri, adăugarea de mai multe secvențe în acesta nu îmbunătățește eficiența.

Pentru a interacționa eficient cu modelul, Tulip se bazează pe grafurile CUDA și pe urmărirea leneșă a rezultatelor pentru a suprapune lucrul pe GPU și CPU și a utiliza pe deplin resursele disponibile.

Gestionarea grafurilor CUDA

Rularea trecerii înainte a unui model implică atât munca pe partea de CPU, cât și pe partea de GPU. CPU-ul este responsabil pentru programarea loturilor și lansarea kernel-urilor cu parametrii adecvați, în timp ce GPU-ul execută multiplicarea de matrice relevantă, atenția, norma sau kernel-urile de activare. Pentru sarcinile de lucru cu debit mare, cum ar fi antrenamentul și reindexarea, costurile generale pe partea de CPU sunt neglijabile deoarece dimensiunile loturilor și latența pe partea de GPU sunt ambele mari. Totuși, la dimensiuni de lot mai mici, munca pe partea de CPU poate depăși munca pe partea de GPU.

Trecere înainte nerăbdătoare: apeluri de gazdă intercalate cu kernel-uri de dispozitiv

Pentru a atenua costurile generale, în loc să lanseze kernel-uri independente, se poate construi un grafic CUDA pentru a captura metadatele necesare pentru a lansa toate kernel-urile unei treceri înainte cu un singur apel către driverul CUDA. Acest lucru elimină necesitatea de a rula din nou codul Python și PyTorch costisitor pentru configurațiile pentru care pot fi capturate grafurile CUDA.

Pentru fiecare model, urmărim un punct de inflexiune, determinând numărul minim de tokenuri la care execuția GPU este mai scumpă decât lansarea kernelului pe partea de CPU. Deoarece modelele de încorporare sunt mici, observăm că acest punct de inflexiune apare la loturi de mii de tokenuri și zeci de secvențe. Unele implementări de atenție se bazează pe intrări dinamice pe partea de gazdă pentru a configura lansările de kernel, prevenind grafurile CUDA dense/de preumplere pentru modelul complet. Am mutat upstream modificările aduse kernel-urilor relevante pentru a le activa în motorul nostru de inferență.

Pentru a aborda costurile generale, construim grafuri CUDA pentru modelul întreg pentru toate modelele de încorporare și suprapunem munca CPU cu munca GPU. Deoarece grafurile CUDA minimizează costurile generale pe partea de CPU, odată ce un grafic este lansat, avem timp liber pentru a iniția și a adăuga în coadă execuția următorului lot ori de câte ori este disponibil. Rezultatele lotului în așteptare sunt urmărite cu un LazyTensor, ceea ce permite unei sarcini asincrone în Rust să se blocheze până când lotul anterior își termină execuția. Grafurile CUDA ajută la servirea cu latență scăzută asigurându-se că nu suntem ținuți pe loc de costul lansărilor de kernel și facilitează o programare îmbunătățită în cazul unui debit mare, deoarece eliberează CPU-ul pentru a lucra la următorul lot mai devreme.

Trecere înainte Cudagraph

Grafurile CUDA trebuie capturate pentru fiecare configurație distinctă, ceea ce pentru încorporări înseamnă un grafic per combinație de număr de secvențe și număr de tokenuri. Deoarece această grilă este extinsă, completăm numărul de tokenuri la găleți care sunt multipli de 64 sau 256. Aceasta are totuși ca rezultat mii de grafice care pot dura câteva minute pentru a fi capturate pentru un model tipic. Costul capturii provine din două surse: o trecere înainte nerăbdătoare care trebuie executată pentru a compila kernel-uri și a configura buffere pentru diverse kernel-uri care au nevoie de ele, urmată de rularea capturii care reexecută codul Python.

Atenuăm costurile de pornire prin capturarea grafurilor CUDA în mod leneș pe măsură ce motorul servește. Ținem evidența fiecărei configurații și ne asigurăm că aceasta trece printr-o rulare de încălzire nerăbdătoare înainte de a declanșa capturarea și reluarea grafului la a doua accesare. Toate execuțiile ulterioare ale aceleiași configurații de grafic trec apoi prin reluarea grafului CUDA. Capturarea leneșă a graficelor are un impact asupra latențelor p99 în timpul pornirii; cu toate acestea, este valoroasă în răspândirea a câteva minute de lucru nerăbdător pe parcursul mai multor ore. Timpii de pornire mai rapidi ne permit să scalăm și să gestionăm mai bine implementările de încorporare.

Tensori leneși

Prin CUDA, munca GPU este asincronă. Deoarece lansarea unui kernel în mod asincron îl adaugă în coadă pe un flux, codul gazdă trebuie să se sincronizeze în mod explicit pentru a citi vectorii rezultati. Pentru a facilita un grad mai mare de paralelism și pentru a putea iniția loturi viitoare în timp ce îl așteptăm pe cel anterior să se finalizeze pe dispozitiv, ne bazăm pe o abstracție LazyTensor pentru a urmări valorile.

LazyTensor urmărește un buffer de gazdă în memorie blocată în pagină și o operație cudaMemcpyAsync printr-un eveniment care copiază datele de pe dispozitiv. Este lansat după pornirea trecerii înainte pe același flux. Deoarece operația de copiere trebuie să aștepte ca toate kernel-urile anterioare de pe flux să se execute, evenimentul asociat urmărește atât finalizarea trecerii înainte, cât și disponibilitatea rezultatului pe CPU.

Execuție LazyTensor

Folosim LazyTensor-uri în motorul nostru de codificare ROSE pentru a suprapune lucrul pe GPU și CPU. În loc ca fiecare apel step() să ruleze graful CUDA și să aștepte finalizarea acestuia, step() returnează un LazyTensor pentru a-i urmări rezultatul în mod asincron. Împreună cu grafurile CUDA, acest lucru ne ajută să obținem latențe scăzute și un debit mai bun.

Cronologie care arată pregătirea CPU și sincronizarea care suprapun loturi GPU succesive

ROSE

Ne-am adaptat motorul ROSE, pe care l-am construit inițial pentru servirea LLM, pentru a gestiona și execuția modelelor de încorporare. Pentru a minimiza efortul necesar pentru a sprijini modelele de încorporare, ROSE reutilizează în mod agresiv codul între LLM-uri și încorporări. De exemplu, servirea pplx-embed și decodarea LLM Qwen3.5 trec prin aceleași kernel-uri. Această partajare ne permite să servim cu ușurință un model de încorporare care a fost inițial fin-tuned dintr-un LLM pentru prototipare, evaluare și inferență în producție.

Pentru straturile dense, încorporarea și inferența LLM sunt identice, deoarece vectorii de tokeni sunt procesați independent. În straturile de atenție, diferențele sunt tratate prin adăugarea de suport pentru intrări zdrențuite (ragged), alături de configurațiile de preumplere și decodare paginată cerute de LLM-uri. Când servim un model de încorporare, nu instanțiem un cache KV și expediem la variații de kernel-uri de atenție care acceptă formatul ragged pentru a evita umplutura. Rutinele de conversie și calibrare suportate sunt, de asemenea, partajate cu LLM-urile.

Ivy

Ivy, stratul nostru de proxy HTTP pentru inferență, joacă, de asemenea, un rol important în performanță. Deoarece sarcinile utile ale cererilor variază în producție, rutarea cererilor individuale către replici individuale poate provoca un dezechilibru al sarcinii. Ivy împarte cererile cu loturi mari în bucăți și echilibrează sarcina între replici, îmbunătățind utilizarea și netezind latența. Munca noastră recentă privind tokenizarea unigram internă, implementată complet în Ivy, îmbunătățește drastic latențele față de tokenizatoarele gata Făcute.

...dar kernel-urile încă contează

ROSE acceptă o varietate de backend-uri de atenție. Diferite kernel-uri pot fi potrivite pentru dimensiuni specifice ale problemelor. În timp, am integrat kernel-urile FlashInfer 2, FlashInfer 3 și FlashAttention 4 pentru a implementa atenția ragged.

Performanța kernelului de atenție după forma modelului și dimensiunea problemei

În general, observăm că FlashAttention 4 este mai rapid. Totuși, FlashInfer 3 îl depășește pe modelele bazate pe Qwen la lungimi de secvență foarte mari. Deoarece performanța și reglarea pot varia în funcție de numărul și dimensiunea capetelor de atenție, menținem suportul pentru mai multe configurații și luăm o decizie de la caz la caz atunci când oferim servicii.

Benchmark-uri

Facem benchmark comparativ cu vLLM v0.22.0, rulând inferența pe precizie BF16 pe ponderi reale ale modelelor și intrări derivate din seturile de date de evaluare. Toate rulările de sincronizare au fost precedate de rulări de încălzire care au verificat că divergența în asemănarea cosinus este în termen de 0.1%.

Încorporări cu latență scăzută (p50 / p90 / p99 / max ms)

Raportăm timpi de rulare pentru dimensiunea lotului de cereri pre-tokenizate 1, cereri complet secvențiale, lungimi de secvență de 128, 512 și 4096 tokenuri.

Rezultatele benchmark-ului pentru încorporări cu latență scăzută

Scorare cu latență scăzută (p50 / p90 / p99 / max ms)

Dimensiuni de lot de cereri pre-tokenizate 5, 25 și 50, lungime a secvenței de 512 tokenuri.

Rezultatele benchmark-ului pentru scorare cu latență scăzută

Încorporări cu debit mare (emb/s)

Dimensiune lot cereri 100, patru procese concurente care trimit cereri, lungimi de secvență de 512, 1024 și 4096 tokenuri.

Rezultatele benchmark-ului pentru încorporări cu debit mare

Încorporări cu concurență ridicată (p50 / p90 / p99 / max ms)

Lungimea secvenței 512, dimensiunea lotului 1, dar trimitem 1, 2, 4, 8 și 16 cereri concurente. Acest benchmark include și costurile de tokenizare prin Ivy, alături de costurile generale de rețea dintre Ivy și Tulip.

Rezultatele benchmark-ului pentru încorporări cu concurență ridicată

Concluzii și lucrări viitoare

Infrastructura de servire compusă din Ivy, Tulip și ROSE ne permite să servim încorporări pentru Perplexity cu o latență mai mică și un debit mai bun, rezultând o căutare mai precisă la un cost redus în comparație cu soluțiile gata fäcute.

Concentrându-ne pe modele specifice și preluând controlul asupra întregii stive, obținem libertatea necesară pentru a găsi un echilibru eficient între performanță și flexibilitate, combinând primitive Rust extrem de reutilizabile și performante alături de cod de modelare Python mai generic. Multe motoare de inferență open-source, cum ar fi vLLM, SGLang și TokenSpeed, integrează limbi precum Rust și C++ în stiva lor. Am investit în Rust în ultimii doi ani și am cules roade mari atât în ceea ce privește performanța, cât și menținerea. Partajând cea mai mare parte a implementării de încorporare cu stiva noastră de servire LLM, obținem câștiguri și în ceea ce privește debitul, fără a necesita un efort de inginerie semnificativ cheltuit pentru întreținerea modelelor de încorporare.

Pe măsură ce modelele evoluează, vom continua să îmbunătățim fiecare strat al stivei noastre pentru a reduce atât latențele legate de CPU, cât și cele legate de GPU. Protocoalele noastre personalizate bazate pe gRPC din cadrul Ivy și Tulip ne permit să ajustăm comunicarea pentru a reduce latențele rețelei, în timp ce ROSE oferă o fundație pentru a îmbunătăți debitul computațional. În plus, pe măsură ce suportul pentru Python fără fire de execuție (free-threaded) crește în cadrul ecosistemului, vom putea îmbunătăți și mai mult interoperabilitatea Python-Rust pentru a reduce costurile generale.

Referințe