Embedding veloci su GPU
La ricerca veloce e accurata è vitale per tutta Perplexity, da Search e Computer alla nostra API Platform. Dietro le quinte, il lavoro pesante è svolto da modelli di embedding e ranking, che aiutano i nostri sistemi a identificare i risultati più pertinenti per una determinata query. Otteniamo qualità e latenza allo stato dell'
La ricerca veloce e accurata è vitale per tutta Perplexity, da Search e Computer alla nostra API Platform. Dietro le quinte, il lavoro pesante è svolto dai modelli di embedding e ranking, che aiutano i nostri sistemi a identificare i risultati più pertinenti per una determinata query. Otteniamo qualità e latenza allo stato dell'arte addestrando e servendo i nostri modelli, come pplx-embed.
Questo articolo presenta una vista sotto il cofano dell'infrastruttura di serving di Perplexity per questa classe speciale di modelli. Discutiamo le nostre tecniche per affrontare in modo efficiente le esigenze di inferenza della ricerca nativa per l'intelligenza artificiale, consentendo una rapida prototipazione e valutazione dei modelli mentre alimentiamo il nostro indice di ricerca su scala exabyte. Queste tecniche espandono collettivamente la frontiera di Pareto della qualità e dell'efficienza della ricerca, consentendoci di servire agenti e utenti con i migliori risultati possibili al minor costo e alla minor latenza.
Embedding per la ricerca
In una tipica configurazione di ricerca, i documenti indicizzati vengono mappati in uno spazio vettoriale ad alta dimensionalità utilizzando un modello di embedding e memorizzati in un database vettoriale. Incorporando una query utilizzando lo stesso modello, è possibile individuare documenti simili trovando i vettori più vicini a quello della query. Ciò dà origine a due diversi modelli di traffico che un motore di inferenza deve servire:
- Embedding in batch: durante la creazione, l'espansione o la reindicizzazione del database, i documenti in blocco devono essere incorporati nello spazio vettoriale, massimizzando il throughput per ridurre al minimo i costi.
A seguito della ricerca vettoriale, grandi batch di documenti devono essere valutati, trovando un equilibrio tra throughput e latenza.
- Embedding online: durante l'interrogazione del database, una query breve deve essere incorporata per le ricerche, riducendo al minimo la latenza.
Abbiamo costruito la nostra infrastruttura di inferenza per sfruttare quanti più componenti comuni possibile tra i vari casi d'uso. Poiché tipicamente utilizziamo piccoli modelli Transformer per produrre embedding, condividiamo la parte principale dell'implementazione con il nostro codice di inferenza LLM: gli embedding in batch sono simili al prefill vincolato dal calcolo, mentre gli embedding online, che spesso vengono eseguiti su pochi token, sono computazionalmente simili al decode vincolato dalla memoria. Riutilizziamo quindi i nostri kernel di prefill e decode ottimizzati per servire i modelli di embedding. Di conseguenza, possiamo ottenere un enorme throughput di inferenza in batch con un lavoro ingegneristico aggiuntivo minimo, preservando al contempo una bassa latenza per i carichi di lavoro di embedding online.
Tulip, Rose e un po' di Ivy
Esponiamo l'inferenza tramite API standardizzate, sia internamente che esternamente attraverso la nostra API Platform. Sotto il cofano, più servizi sono coinvolti nell'elaborazione di una richiesta di embedding:
- Ivy è un gateway HTTP Rust che i servizi di Perplexity chiamano.
Gestisce il lavoro lato CPU per richieste come l'analisi JSON, la tokenizzazione, la creazione di modelli di input e la suddivisione dei batch, traducendo le richieste in un protocollo gRPC personalizzato per i server downstream. Questa separazione ci consente di configurare determinati parametri relativi alla tokenizzazione e alla formattazione dell'input senza dover toccare le istanze di inferenza più pesanti.
- Tulip è l'interfaccia del server di inferenza.
È un server gRPC implementato con Rust, tokio e `tonic`. Tulip riceve richieste di inferenza gRPC, gestendo la pianificazione e il batching. Invia quindi i batch al motore ROSE, restituendo le risposte completate ai client.
- **ROSE** (Runtime-Optimized Serving Engine) implementa l'inferenza del modello.
È definito principalmente in Python, fornendo kernel, livelli e definizioni per un'ampia varietà di modelli. ROSE implementa i passaggi in avanti attraverso i modelli, fornendo anche la gestione dei grafi CUDA specializzata per gli embedding. È collegato a Tulip tramite una funzione step(), che prende un batch e restituisce un riferimento al calcolo che esegue sull'acceleratore.
Prestare attenzione oltre il kernel
Sia i modelli basati su Transformer che le architetture sottostanti Hopper/Blackwell sono tecnologie mature, quindi l'inferenza di embedding sul lato GPU è converso verso un'implementazione ampiamente ottimale su vari motori di inferenza. Ciononostante, abbiamo scoperto ulteriori opportunità di miglioramento nei runtime e nei moduli di interfaccia che espongono i modelli end-to-end a un client. In particolare, abbiamo scoperto di poter migliorare le latenze gestendo con attenzione i grafi CUDA e costruendo un'astrazione LazyTensor per tracciare in modo asincrono un risultato sul lato GPU nel motore Rust nativo. Abbiamo implementato queste funzionalità in Tulip, in modo che potesse interfacciarsi efficacemente con le implementazioni dei modelli di ROSE.
Tulip
Abbiamo progettato Tulip per essere un'interfaccia quanto più leggera possibile sul nostro model serving. Gestisce le richieste in arrivo in attività asincrone di Tokio, mantenendo un pool di richieste che traccia e da cui pianifica i batch da inviare all'acceleratore. Il meccanismo di pianificazione in Tulip è molto semplice: le richieste si accumulano mentre Tulip sta inviando lavoro o è in attesa di risultati. Dalle richieste accumulate, le sequenze vengono scelte secondo il criterio del primo arrivato, primo servito per essere eseguite attraverso il modello.
Il semplice meccanismo di pianificazione è motivato da un'osservazione sulle prestazioni del modello. Per i piccoli modelli di embedding, alle lunghezze di sequenza che serviamo, abbiamo notato che il costo lineare dei livelli densi è dominante rispetto al costo quadratico dell'attenzione. Pertanto, la latenza è per lo più proporzionale al numero di token, non al numero di sequenze. Di conseguenza, una volta che un batch è abbastanza grande da saturare la GPU, che è di circa 512 token su un modello sotto un miliardo di parametri, inserire più sequenze al suo interno non migliora l'efficienza.
Per interfacciarsi efficacemente con il modello, Tulip si basa sui grafi CUDA e sul tracciamento dei risultati pigri per sovrapporre il lavoro di GPU e CPU e utilizzare appieno le risorse disponibili.
Gestione dei grafi CUDA
L'esecuzione del passaggio in avanti di un modello comporta lavoro sia sul lato CPU che sul lato GPU. La CPU è responsabile della pianificazione dei batch e dell'avvio dei kernel con i parametri appropriati, mentre la GPU esegue la relativa moltiplicazione di matrici, attenzione, norma o i kernel di attivazione. Per carichi di lavoro ad alto throughput come l'addestramento e la reindicizzazione, gli overhead sul lato CPU sono trascurabili perché le dimensioni dei batch e la latenza sul lato GPU sono entrambe grandi. Tuttavia, su dimensioni di batch più piccole, il lavoro sul lato CPU può superare il lavoro sul lato GPU.
Per mitigare gli overhead, invece di lanciare kernel indipendenti, è possibile costruire un grafo CUDA per catturare i metadati necessari per lanciare tutti i kernel di un passaggio in avanti con una singola chiamata al driver CUDA. Ciò elimina la necessità di rieseguire costoso codice Python e PyTorch per le configurazioni per cui è possibile catturare i grafi CUDA.
Su ciascun modello, tracciamo un punto di flesso, determinando il numero minimo di token in cui l'esecuzione su GPU è più costosa rispetto all'avvio del kernel sul lato CPU. Poiché i modelli di embedding sono piccoli, osserviamo che questo punto di flesso arriva a batch di migliaia di token e decine di sequenze. Alcune implementazioni di attenzione si basano su input dinamici lato host per configurare gli avvii dei kernel, impedendo grafi CUDA densi/di prefill a modello intero. Abbiamo eseguito l'upstream delle modifiche ai kernel pertinenti per abilitarli nel nostro motore di inferenza.
Per affrontare gli overhead, costruiamo grafi CUDA a modello intero per tutti i modelli di embedding e sovrapponiamo il lavoro della CPU con il lavoro della GPU. Poiché i grafi CUDA riducono al minimo gli overhead sul lato CPU, una volta lanciato un grafo, abbiamo tempo libero per avviare e mettere in coda l'esecuzione del batch successivo ogni volta che è disponibile. I risultati del batch in sospeso vengono tracciati con un LazyTensor, che consente a un'attività asincrona in Rust di bloccarsi fino al completamento dell'esecuzione del batch precedente. I grafi CUDA aiutano il serving a bassa latenza assicurando di non essere frenati dal costo degli avvii dei kernel e facilitano una migliore gestione della pianificazione nel caso di alto throughput poiché liberano la CPU per lavorare prima sul batch successivo.
I grafi CUDA devono essere catturati per ogni configurazione distinta, il che per gli embedding significa un grafo per combinazione di conteggio delle sequenze e conteggio dei token. Poiché questa griglia è estesa, eseguiamo il padding dei conteggi dei token a bucket che sono multipli di 64 o 256. Ciò si traduce comunque in migliaia di grafi che potrebbero richiedere diversi minuti per essere catturati per un modello tipico. Il costo della cattura deriva da due fonti: un passaggio in avanti eager che deve essere eseguito per compilare i kernel e configurare i buffer per vari kernel che ne hanno bisogno, seguito dall'esecuzione di cattura che riesegue il codice Python.
Mitighiamo i costi di avvio catturando i grafi CUDA in modo pigro man mano che il motore serve. Teniamo traccia di ciascuna configurazione e ci assicuriamo che passi attraverso un'esecuzione di riscaldamento eager prima di attivare la cattura e la riproduzione del grafo alla seconda richiesta. Tutte le esecuzioni successive della stessa configurazione di grafo passano quindi attraverso la riproduzione del grafo CUDA. La cattura pigra del grafo ha un impatto sulle latenze p99 durante l'avvio; tuttavia, è preziosa per distribuire diversi minuti di lavoro eager su più ore. Tempi di avvio più rapidi ci consentono di scalare e gestire meglio le distribuzioni di embedding.
Tensor pigri
Tramite CUDA, il lavoro della GPU è asincrono. Poiché il lancio di un kernel in modo asincrono lo mette in coda su un flusso, il codice dell'host deve sincronizzarsi esplicitamente per leggere i vettori risultanti. Per facilitare un maggiore grado di parallelismo ed essere in grado di avviare batch futuri in attesa che il precedente completi l'esecuzione sul dispositivo, ci affidiamo a un'astrazione LazyTensor per tracciare i valori.
Il LazyTensor tiene traccia di un buffer host in memoria bloccata in pagina e di un'operazione cudaMemcpyAsync tramite un evento che copia i dati dal dispositivo. Viene avviato dopo il lancio del passaggio in avanti sullo stesso flusso. Poiché l'operazione di copia deve attendere l'esecuzione di tutti i kernel precedenti sul flusso, l'evento associato tiene traccia sia del completamento del passaggio in avanti sia della disponibilità del risultato sulla CPU.
Sfruttiamo i LazyTensor nel nostro motore di codifica ROSE per sovrapporre il lavoro di GPU e CPU. Invece di ogni chiamata a step() che esegue il grafo CUDA e attende il suo completamento, step() restituisce un LazyTensor per tracciarne asincronamente il risultato. Abbinato ai grafi CUDA, questo ci aiuta a ottenere basse latenze e un migliore throughput.
ROSE
Abbiamo adattato il nostro motore ROSE, originariamente costruito per il serving di LLM, per gestire anche l'esecuzione di modelli di embedding. Per ridurre al minimo lo sforzo necessario per supportare i modelli di embedding, ROSE riutilizza aggressivamente il codice tra LLM e embedding. Ad esempio, il serving di pplx-embed e il decoding di LLM Qwen3.5 passano tutti attraverso gli stessi kernel. Questa condivisione ci consente di servire facilmente un modello di embedding originariamente sottoposto a fine-tuning da un LLM per la prototipazione, la valutazione e l'inferenza di produzione.
Per i livelli densi, l'inferenza di embedding e LLM è identica poiché i vettori di token vengono elaborati in modo indipendente. Nei livelli di attenzione, le differenze vengono gestite aggiungendo il supporto per input irregolari, insieme alle configurazioni di prefill e decode con paginazione richieste dagli LLM. Quando si serve un modello di embedding, non istanziamo una cache KV e ci indirizziamo a variazioni di kernel di attenzione che supportano il formato irregolare per evitare il padding. Le routine di conversione e calibrazione di supporto sono anch'esse condivise con gli LLM.
Ivy
Ivy, il nostro livello di proxy HTTP per l'inferenza, svolge anch'esso un ruolo importante nelle prestazioni. Poiché i payload delle richieste variano in produzione, il routing delle singole richieste a singole repliche può causare uno squilibrio del carico. Ivy suddivide le richieste di grandi batch in blocchi e ne bilancia il carico tra le repliche, migliorando l'utilizzo e attenuando la latenza. Il nostro recente lavoro sulla tokenizzazione unigram interna, completamente distribuita in Ivy, migliora drasticamente le latenze rispetto ai tokenizer pronti all'uso.
...ma i kernel contano ancora
ROSE supporta una varietà di backend di attenzione. Diversi kernel possono essere adatti a dimensioni di problema specifiche. Nel tempo, abbiamo integrato i kernel FlashInfer 2, FlashInfer 3 e FlashAttention 4 per implementare l'attenzione irregolare.
In generale, osserviamo che FlashAttention 4 è più veloce. Tuttavia, FlashInfer 3 lo supera sui modelli basati su Qwen a lunghezze di sequenza molto elevate. Poiché le prestazioni e la messa a punto possono variare in base al numero e alla dimensione delle teste di attenzione, manteniamo il supporto per più configurazioni e prendiamo una decisione caso per caso durante il serving.
Benchmark
Eseguiamo il benchmark rispetto a vLLM v0.22.0, eseguendo l'inferenza su precisione BF16 su pesi di modelli reali e input derivati da dataset di valutazione. Tutte le esecuzioni di tempistica sono state precedute da esecuzioni di riscaldamento che hanno verificato che la divergenza nella somiglianza del coseno sia entro lo 0,1%.
Embedding a bassa latenza (p50 / p90 / p99 / max ms)
Riportiamo i runtime per dimensione del batch di richieste pre-tokenizzate 1, richieste completamente sequenziali, lunghezze delle sequenze di 128, 512 e 4096 token.
Scoring a bassa latenza (p50 / p90 / p99 / max ms)
Dimensioni del batch di richieste pre-tokenizzate 5, 25 e 50, lunghezza della sequenza di 512 token.
Embedding ad alto throughput (emb/s)
Dimensione del batch di richieste 100, quattro processi concorrenti che inviano richieste, lunghezze delle sequenze di 512, 1024 e 4096 token.
Embedding ad alta concorrenza (p50 / p90 / p99 / max ms)
Lunghezza della sequenza 512, dimensione del batch 1, ma inviamo 1, 2, 4, 8 e 16 richieste concorrenti. Questo benchmark include anche i costi di tokenizzazione tramite Ivy, insieme all'overhead di rete tra Ivy e Tulip.
Conclusione e lavoro futuro
L'infrastruttura di serving composta da Ivy, Tulip e ROSE ci consente di servire gli embedding per Perplexity con una latenza inferiore e un migliore throughput, risultando in una ricerca più accurata a un costo ridotto rispetto alle soluzioni pronte all'uso.
Concentrandoci su modelli specifici e prendendo il controllo dell'intero stack, otteniamo la libertà necessaria per trovare un equilibrio efficace tra prestazioni e flessibilità, mescolando primitive Rust altamente riutilizzabili e performanti insieme a codice di modellazione Python più generico. Molti motori di inferenza open source, come vLLM, SGLang e TokenSpeed, stanno integrando linguaggi come Rust e C++ nel loro stack. Abbiamo investito in Rust negli ultimi due anni e abbiamo raccolto grandi frutti sia in termini di prestazioni che di manutenibilità. Condividendo la maggior parte dell'implementazione degli embedding con il nostro stack di serving LLM, otteniamo guadagni anche in termini di throughput, senza richiedere un significativo sforzo ingegneristico da dedicare alla manutenzione dei modelli di embedding.
Man mano che i modelli si evolvono, continueremo a migliorare ogni livello del nostro stack per ridurre sia le latenze vincolate dalla CPU che quelle vincolate dalla GPU. I nostri protocolli personalizzati basati su gRPC all'interno di Ivy e Tulip ci consentono di ottimizzare la comunicazione per ridurre le latenze di rete, mentre ROSE fornisce una base per migliorare il throughput computazionale. Inoltre, man mano che il supporto per Python senza GIL cresce in tutto l'ecosistema, saremo in grado di migliorare ulteriormente l'interoperabilità Python-Rust per ridurre gli overhead.