Incrustacions ràpides a les GPU
Una cerca ràpida i precisa és vital per a tot Perplexity, des de Search i Computer fins a la nostra API Platform. Darrere de les escenes, la feina pesada es fa mitjançant models d'incrustació i classificació, que ajuden els nostres sistemes a identificar els resultats més rellevants per a una consulta donada. Aconseguim l'estat de la tècnica
Una cerca ràpida i precisa és vital per a tot Perplexity, des de Search i Computer fins a la nostra API Platform. Entre bastidors, la feina pesada la fan els models d'incrustació i classificació, que ajuden els nostres sistemes a identificar els resultats més rellevants per a una consulta donada. Aconseguim una qualitat i una latència d'última generació entrenant i servint els nostres propis models, com ara pplx-embed.
Aquest article ofereix una visió interna de la infraestructura de serveis de Perplexity per a aquesta classe especial de models. Discutim les nostres tècniques per abordar eficientment les necessitats d'inferència de la cerca nativa d'IA, permetent la creació ràpida de prototips i l'avaluació de models alhora que potenciem el nostre índex de cerca a escala d'exaoctets. Aquestes tècniques conjuntament amplien la frontera de Pareto de la qualitat i l'eficiència de la cerca, permetent-nos servir agents i usuaris amb els millors resultats possibles al menor cost i latència.
Incrustacions per a la cerca
En una configuració de cerca típica, els documents indexats es mapegen a un espai vectorial d'alta dimensió utilitzant un model d'incrustació i s'emmagatzemen en una base de dades vectorial. Mitjançant la incrustació d'una consulta utilitzant el mateix model, es poden localitzar documents similars trobant els vectors més propers als de la consulta. Això dóna lloc a dos patrons de trànsit diferents perquè un motor d'inferència serveixi:
- Incrustació per lots: en construir, ampliar o reindexar la base de dades, els documents massius s'han d'incrustar a l'espai vectorial, maximitzant el rendiment per minimitzar el cost.
Després de la cerca de vectors, s'han de puntuar grans lots de documents, trobant un equilibri entre rendiment i latència.
- Incrustació en línia: en consultar la base de dades, s'ha d'incrustar una consulta curta per a les cerques, minimitzant la latència.
Vam construir la nostra infraestructura d'inferència per aprofitar tants components comuns com sigui possible en diferents casos d'ús. Atès que normalment utilitzem models Transformer petits per produir incrustacions, compartim la major part de la implementació amb el nostre codi d'inferència LLM: les incrustacions per lots són similars a la precàrrega limitada pel càlcul, mentre que les incrustacions en línia, que sovint s'executen en uns quants tokens, són computacionalment similars a la descodificació limitada per la memòria. Per tant, reutlitzem els nostres kernels de precàrrega i descodificació optimitzats per servir models d'incrustació. Com a resultat, podem aconseguir un rendiment massiu d'inferència per lots amb un treball d'enginyeria addicional mínim, alhora que preservem una baixa latència per a les càrregues de treball d'incrustacions en línia.
Tulipes, roses i una mica d'Ivy
Exposem la inferència a través d'API estandarditzades, tant internament com externament a través de la nostra API Platform. Sota el capó, hi ha múltiples serveis implicats en el processament d'una sol·licitud d'incrustació:
- Ivy és una porta d'enllaç HTTP de Rust a la qual truquen els serveis de Perplexity.
S'encarrega del treball del costat de la CPU per a sol·licituds com ara l'anàlisi JSON, la tokenització, la creació de plantilles d'entrada i la divisió de lots, traduint les sol·licituds a un protocol gRPC personalitzat per als servidors aigües avall. Aquesta separació ens permet configurar determinats paràmetres relacionats amb la tokenització i el format d'entrada sense haver de tocar les instàncies d'inferència més pesades.
- Tulip és la interfície del servidor d'inferència.
És un servidor gRPC implementat amb Rust, tokio i `tonic`. Tulip rep sol·licituds d'inferència gRPC, gestionant la planificació i l'agrupació per lots. A continuació, envia els lots al motor ROSE, retornant les respostes completades als clients.
- **ROSE** (Runtime-Optimized Serving Engine) implementa la inferència de models.
Està definit principalment en Python, proporcionant kernels, capes i definicions per a una àmplia varietat de models. ROSE implementa els passos endavant a través dels models, proporcionant també la gestió de grafs CUDA especialitzada per a encastaments. Es connecta a Tulip mitjançant una funció step(), que pren un lot i retorna una referència al càlcul que realitza a l'accelerador.
Prestant atenció més enllà del kernel
Tant els models basats en Transformers com les arquitectures subjacents Hopper/Blackwell són tecnologies madures, de manera que la inferència incrustada al costat de la GPU ha convergit cap a una implementació majoritàriament òptima en diversos motors d'inferència. Tot i així, vam descobrir oportunitats addicionals de millora en els entorns d'execució i les infraestructures que exposen els models d'extrem a extrem a un client. En particular, vam trobar que podem millorar les latències gestionant acuradament els grafs CUDA i construint una abstracció LazyTensor per fer un seguiment asíncron d'un resultat al costat de la GPU en el motor nadiu de Rust. Vam implementar aquestes característiques a Tulip perquè pogués connectar-se eficaçment amb les implementacions de models de ROSE.
Tulip
Hem dissenyat Tulip perquè sigui una interfície tan lleugera com sigui possible sobre el nostre servei de models. Gestiona les sol·licituds entrants en tasques asíncrones de Tokio, mantenint un grup de sol·licituds de les quals fa un seguiment i planifica lots per enviar a l'accelerador. El mecanisme de planificació de Tulip és molt senzill: les sol·licituds s'acumulen mentre Tulip està enviant feina o esperant resultats. De les sol·licituds acumulades, les seqüències es trien per ordre d'arribada per executar-les a través del model.
El mecanisme de planificació senzill està motivat per una observació sobre el rendiment del model. Per a models d'incrustació petits, a les longituds de seqüència per a les quals servim, vam notar que el cost lineal de les capes denses és dominant sobre el cost quadràtic de l'atenció. Per tant, la latència és majoritàriament proporcional al nombre de tokens, no al nombre de seqüències. En conseqüència, un cop un lot és prou gran com per saturar la GPU, que és al voltant de 512 tokens en un model de menys de mil milions de paràmetres, afegir més seqüències no millora l'eficiència.
Per connectar-se eficaçment amb el model, Tulip es basa en els grafs CUDA i el seguiment de resultats mandrós per solapar el treball de la GPU i la CPU i utilitzar plenament els recursos disponibles.
Gestió de grafs CUDA
L'execució del pas endavant d'un model implica tant el treball de la CPU com de la GPU. La CPU s'encarrega de planificar els lots i lullançar kernels amb els paràmetres adequats, mentre que la GPU executa els kernels de multiplicació de matrius, atenció, norma o activació rellevants. Per a càrregues de treball d'alt rendiment com l'entrenament i la reindexació, les sobrecàrregues del costat de la CPU són negligibles perquè les mides dels lots i la latència del costat de la GPU són grans. No obstant això, en mides de lot més petites, el treball de la CPU pot superar el treball de la GPU.
Per mitigar les sobrecàrregues, en lloc de llançar kernels independents, es pot construir un graf CUDA per capturar les metadades necessàries per llançar tots els kernels d'un pas endavant amb una sola crida al controlador CUDA. Això elimina la necessitat de tornar a executar codi costós de Python i PyTorch per a les configuracions per a les quals es poden capturar grafs CUDA.
En cada model, fem un seguiment d'un punt d'inflexió, determinant el nombre mínim de tokens en què l'execució de la GPU és més cara que el llançament de kernel al costat de la CPU. Com que els models d'incrustació són petits, observem que aquest punt d'inflexió es produeix en lots de milers de tokens i desenes de seqüències. Algunes implementacions d'atenció depenen d'entrades dinàmiques del costat de l'amfitrió per configurar els llançaments de kernels, evitant els grafs CUDA de precàrrega/densitat de tot el model. Vam integrar a aigües amunt els canvis als kernels rellevants per habilitar-los al nostre motor d'inferència.
Per abordar les sobrecàrregues, construïm grafs CUDA de tot el model per a tots els models d'incrustació i solapem el treball de la CPU amb el treball de la GPU. Com que els grafs CUDA minimitzen les sobrecàrregues del costat de la CPU, un costat es llança un graf, tenim temps lliure per iniciar i encadenar l'execució del següent lot sempre que estigui disponible. Els resultats del lot pendent es fan un seguiment amb un LazyTensor, la qual cosa permet que una tasca asíncrona a Rust es bloquegi fins que el lot anterior acabi l'execució. Els grafs CUDA ajuden al servei de baixa latència assegurant que no ens veiem frenats pel cost dels llançaments de kernels i faciliten una planificació millorada en el cas d'alt rendiment, ja que alliberen la CPU perquè pugui treballar en el lot següent més ràpidament.
Els grafs CUDA s'han de capturar per a cada configuració diferent, la qual cosa per a les incrustacions significa un graf per combinació de recompte de seqüències i recompte de tokens. Com que aquesta graella és àmplia, farcim els recomptes de tokens a cubells que són múltiples de 64 o 256. Això encara dóna lloc a milers de grafs que poden trigar uns quants minuts a capturar-se per a un model típic. El cost de la captura prové de dues fonts: un pas endavant eager que s'ha d'executar per compilar kernels i configurar memòries cau per a diversos kernels que els necessiten, seguit de l'execució de captura que torna a executar el codi de Python.
Mitiguem els costos d'inici capturant grafs CUDA mandrosament a mesura que el motor serveix. Fem un seguiment de cada configuració i ens assegurem que passa per una execució d'escalfament eager abans de desencadenar la captura i reproducció del graf en el segon encert. Totes les execucions posteriors de la mateixa configuració de graf passen llavors per la reproducció del graf CUDA. La captura de grafs mandrosa té un impacte en les latències p99 durant l'inici; tanmateix, és valuosa per repartir diversos minuts de treball eager al llarg de diverses hores. Uns temps d'inici més ràpids ens permeten escalar i gestionar millor els desplegaments d'incrustacions.
Tensors mandrosos
A través de CUDA, el treball de la GPU és asíncron. Com que el llançament d'un kernel de manera asíncrona l'encadena en un flux, el codi de l'amfitrió s'ha de sincronitzar explícitament per llegir els vectors resultants. Per facilitar un major grau de paral·lelisme i poder iniciar lots futurs mentre s'espera que l'anterior es completi al dispositiu, confiem en una abstracció LazyTensor per fer un seguiment dels valors.
El LazyTensor fa un seguiment d'una memòria cau d'amfitrió en memòria bloquejada a la pàgina i d'una operació cudaMemcpyAsync mitjançant un esdeveniment que copia dades del dispositiu. S'inicia després del llançament del pas endavant al mateix flux. Com que l'operació de còpia ha d'esperar que s'executin tots els kernels anteriors del flux, l'esdeveniment associat fa un seguiment tant de la finalització del pas endavant com de la disponibilitat del resultat a la CPU.
Aprofitem els LazyTensors al nostre motor de codificació ROSE per solapar el treball de la GPU i la CPU. En lloc de cada crida a step() executant el graf CUDA i esperant que acabi, step() retorna un LazyTensor per fer un seguiment asíncron del seu resultat. Juntament amb els grafs CUDA, això ens ajuda a aconseguir latències baixes i un millor rendiment.
ROSE
Hem adaptat el nostre motor ROSE, que originalment vam construir per al servei LLM, per gestionar també l'execució de models d'incrustació. Per minimitzar l'esforç necessari per donar suport als models d'incrustació, ROSE reutilitza agressivament el codi entre LLM i incrustacions. Per exemple, el servei pplx-embed i la descodificació LLM Qwen3.5 passen pels mateixos kernels. Aquesta compartició ens permet servir fàcilment un model d'incrustació que es va ajustar originalment a partir d'un LLM per a la creació de prototips, l'avaluació i la inferència de producció.
Per a les capes denses, la inferència d'incrustació i LLM són idèntiques, ja que els vectors de tokens es processen de manera independent. A les capes d'atenció, les diferències es gestionen afegint suport per a entrades irregulars, juntament amb les configuracions de precàrrega i descodificació paginades requerides pels LLM. Quan servim un model d'incrustació, no instanciem una memòria cau KV i enviem a variacions de kernels d'atenció que admeten el format irregular per evitar el farciment. Les rutines de conversió i calibratge de suport també es comparteixen amb els LLM.
Ivy
Ivy, la nostra capa de servidor intermediari HTTP d'inferència, també juga un paper important en el rendiment. Com que les càrregues útils de les sol·licituds varien en producció, enrutar sol·licituds individuals a rèpliques individuals pot provocar un desequilibri de càrrega. Ivy divideix les sol·licituds de lots grans en fragments i equilibra la càrrega entre les rèpliques, millorant la utilització i suavitzant la latència. El nostre treball recent sobre la tokenització de unigrames pròpia, completament implementada a Ivy, millora dràsticament les latències respecte als tokenitzadors comercials.
...però els kernels encara importen
ROSE admet una varietat de dorsals d'atenció. Diferents kernels poden ser adequats per a mides de problemes específiques. Amb el temps, vam integrar els kernels FlashInfer 2, FlashInfer 3 i FlashAttention 4 per implementar l'atenció irregular.
En general, observem que FlashAttention 4 és més ràpid. No obstant això, FlashInfer 3 el supera en models basats en Qwen amb longituds de seqüència molt llargues. Com que el rendiment i l'ajust poden variar segons el nombre i la dimensió de les capçaleres d'atenció, mantenim el suport per a múltiples configuracions i prenem una decisió cas per cas en servir.
Referències
Fem una comparativa amb vLLM v0.22.0, executant la inferència amb precisió BF16 sobre pesos de models reals i entrades derivades de conjunts de dades d'avaluació. Totes les execucions de temps van ser precedides per execucions d'escalfament que van verificar que la divergència en la similitud del cosinus està dins del 0,1%.
Incrustacions de baixa latència (p50 / p90 / p99 / màx. ms)
Informem dels temps d'execució per a una mida de lot de sol·licituds pretokenitzades d'1, sol·licituds completament seqüencials, longituds de seqüència de 128, 512 i 4096 tokens.
Puntuació de baixa latència (p50 / p90 / p99 / màx. ms)
Mides de lots de sol·licituds pretokenitzades de 5, 25 i 50, longitud de seqüència de 512 tokens.
Incrustacions d'alt rendiment (emb/s)
Midat del lot de sol·licituds 100, quatre processos concurrents que envien sol·licituds, longituds de seqüència de 512, 1024 i 4096 tokens.
Incrustacions d'alta concurrència (p50 / p90 / p99 / màx. ms)
Longitud de seqüència 512, mida de lot 1, però enviem 1, 2, 4, 8 i 16 sol·licituds concurrents. Aquesta referència també inclou els costos de tokenització a través d'Ivy, juntament amb la sobrecàrrega de xarxa entre Ivy i Tulip.
Conclusió i treball futur
La infraestructura de serveis composta per Ivy, Tulip i ROSE ens permet servir incrustacions per a Perplexity amb una latència més baixa i un millor rendiment, donant lloc a una cerca més precisa a un cost reduït en comparació amb les solucions comercials.
Centrant-nos en models específics i assumint la propietat de tota la pila, obtenim la llibertat necessària per trobar un equilibri eficaç entre rendiment i flexibilitat, barrejant primitives de Rust altament reutilitzables i d'alt rendiment juntament amb codi de modelatge de Python més genèric. Molts motors d'inferència de codi obert, com vLLM, SGLang i TokenSpeed, estan integrant llenguatges com Rust i C++ a la seva pila. Hem invertit en Rust durant els darrers dos anys i hem obtingut grans recompenses tant en rendiment com en mantenibilitat. En compartir la major part de la implementació d'incrustació amb la nostra pila de serveis LLM, també obtenim guanys de rendiment, sense necessitat de dedicar un esforç d'enginyeria significatiu al manteniment dels models d'incrustació.
A mesura que els models evolucionin, continuarem millorant cada capa de la nostra pila per reduir tant les latències limitades per la CPU com per la GPU. Els nostres protocols personalitzats basats en gRPC dins d'Ivy i Tulip ens permeten ajustar la comunicació per reduir les latències de la xarxa, mentre que ROSE ofereix una base per millorar el rendiment computacional. A més, a mesura que el suport per a Python de fil lliure creixi a tot l'ecosistema, podrem millorar encara més la interoperabilitat entre Python i Rust per reduir les sobrecàrregues.