Incrustacións rápidas en GPUs
A busca rápida e precisa é fundamental para todo Perplexity, desde Search e Computer ata a nosa plataforma de API. Detrás das cámaras, o traballo pesado realízase mediante modelos de incrustación e ranking, que axudan aos nosos sistemas a identificar os resultados máis relevantes para unha consulta determinada. Conseguimos o estado da arte
A busca rápida e precisa é fundamental para todo Perplexity, desde Search e Computer ata a nosa plataforma de API. Detrás das cámaras, o traballo pesado realízase mediante modelos de incrustación e ranking, que axudan aos nosos sistemas a identificar os resultados máis relevantes para unha consulta determinada. Conseguimos calidade e latencia de vangarda mediante o adestramento e servizo dos nosos propios modelos, como pplx-embed.
Este artigo presenta unha visión interna da infraestrutura de servizo de Perplexity para esta clase especial de modelos. Discutimos as nosas técnicas para abordar de maneira eficiente as necesidades de inferencia da busca nativa de IA, permitindo a creación rápida de prototipos e a avaliación de modelos mentres se impulsa o noso índice de busca a escala de exabytes. Estas técnicas amplían colectivamente a fronteira de Pareto de calidade e eficiencia de busca, permitíndonos atender a axentes e usuarios cos mellores resultados posibles ao menor custo e latencia.
Incrustacións para a busca
Nunha configuración de busca típica, os documentos indexados mámpanse nun espazo vectorial de alta dimensionalidade utilizando un modelo de incrustación e almacénanse nunha base de datos vectorial. Ao incrustar unha consulta utilizando o mesmo modelo, pódense localizar documentos similares atopando os vectores máis achegados aos da consulta. Isto dá lugar a dous patróns de tráfico diferentes para que un motor de inferencia os atenda:
- Incrustación por lotes: ao construír, expandir ou reindexar a base de datos, os documentos masivos deben incrustarse no espazo vectorial, maximizando o rendemento para minimizar o custo.
Despois da busca vectorial, grandes lotes de documentos deben ser puntuados, logrando un equilibrio entre rendemento e latencia.
- Incrustación en liña: ao consultar a base de datos, debe incrustarse unha consulta curta para as buscas, minimizando a latencia.
Construímos a nosa infraestrutura de inferencia para aproveitar tantos compoñentes comúns como sexa posible entre casos de uso. Dado que habitualmente empregamos pequenos modelos Transformer para producir incrustacións, compartimos a maior parte da implementación co noso código de inferencia de LLM: as incrustacións por lotes son similares ao preenchido limitado por computación, mentres que as incrustacións en liña, que adoitan executarse en poucos tokens, son computacionalmente similares á descodificación limitada por memoria. Así, reutilizamos os nosos kernels optimizados de preenchido e descodificación para servir modelos de incrustación. Como resultado, podemos acadar un enorme rendemento de inferencia por lotes cun traballo de enxeñaría adicional mínimo, ao tempo que preservamos unha baixa latencia para as cargas de traballo de incrustacións en liña.
Tulips, Roses e algo de Ivy
Expoñemos a inferencia a través de API estandarizadas, tanto internamente como externamente a través da nosa plataforma de API. Detrás das cámaras, múltiples servizos están implicados no procesamento dunha solicitude de incrustación:
- Ivy é unha porta de enlace HTTP en Rust á que chaman os servizos de Perplexity.
Xestiona o traballo do lado da CPU para solicitudes como análise JSON, tokenización, plantillas de entrada e división de lotes, traducindo solicitudes a un protocolo gRPC personalizado para servidores posteriores. Esta separación permítenos configurar certos parámetros arredor da tokenización e o formato de entrada sen ter que tocar as instancias de inferencia máis pesadas.
- Tulip é a interface do servidor de inferencia.
É un servidor gRPC implementado con Rust, tokio e `tonic`. Tulip recibe solicitudes de inferencia gRPC, xestionando a programación e a formación de lotes. A continuación, envía os lotes ao motor ROSE, devolvendo respostas completas aos clientes.
- **ROSE** (Runtime-Optimized Serving Engine) implementa a inferencia de modelos.
Está definido principalmente en Python, proporcionando kernels, capas e definicións para unha ampla variedade de modelos. ROSE implementa os pasos cara adiante a través dos modelos, proporcionando tamén xestión de gráficos CUDA especializada para incrustacións. Conéctase a Tulip mediante unha función step(), que toma un lote e devolve unha referencia á computación que realiza no acelerador.
Prestando atención máis alá do kernel
Tanto os modelos baseados en Transformer como as arquitecturas subxacentes Hopper/Blackwell son tecnoloxías maduras, polo que a inferencia de incrustacións do lado da GPU converxeu nunha implementación amplamente óptima en varios motores de inferencia. Aínda así, descubrimos oportunidades adicionais de mellora nos tempos de execución e nos marcos que expoñen os modelos de extremo a extremo a un cliente. En particular, descubrimos que podemos mellorar as latencias xestionando coidadoamente os gráficos CUDA e construíndo unha abstracción LazyTensor para rastrexar de maneira asíncrona un resultado do lado da GPU no motor nativo en Rust. Implementamos estas características en Tulip, para que puidese interactuar eficazmente coas implementacións de modelos de ROSE.
Tulip
Deseñamos Tulip para que sexa unha interface o máis lixeira posible sobre o noso servizo de modelos. Xestiona as solicitudes entrantes en tarefas asíncronas de Tokio, mantendo un conxunto de solicitudes que rastrexa e desde o que programa lotes para enviar ao acelerador. O mecanismo de programación en Tulip é moi sinxelo: as solicitudes acumúlanse mentres Tulip está a enviar traballo ou agardando resultados. Das solicitudes acumuladas, as secuencias elíxense por orde de chegada para ser executadas a través do modelo.
O mecanismo de programación sinxelo está motivado por unha observación sobre o rendemento do modelo. Para pequenos modelos de incrustación, nas lonxitudes de secuencia que servimos, notamos que o custo lineal das capas densas é dominante sobre o custo cuadrático da atención. Por tanto, a latencia é na súa maioría proporcional ao número de tokens, non ao número de secuencias. En consecuencia, unha vez que un lote é o suficientemente grande como para saturar a GPU, que é ao redor de 512 tokens nun modelo de menos de mil millóns de parámetros, empaquetar máis secuencias nel non mellora a eficiencia.
Para interactuar eficazmente co modelo, Tulip confía en gráficos CUDA e rastreo de resultados preguiceiros para solapar o traballo da GPU e da CPU e utilizar plenamente os recursos dispoñibles.
Xestión de gráficos CUDA
A execución do paso cara adiante dun modelo implica traballo tanto no lado da CPU como no da GPU. A CPU encárgase de programar lotes e lanzar kernels cos parámetros axeitados, mentres que a GPU executa a multiplicación de matrices, a atención, a norma ou os kernels de activación relevantes. Para cargas de traballo de alto rendemento como o adestramento e a reindexación, as sobrecargas do lado da CPU son despreciables porque os tamaños dos lotes e a latencia do lado da GPU son grandes. Con todo, en tamaños de lote máis pequenos, o traballo do lado da CPU pode superar ao do lado da GPU.
Para mitigar as sobrecargas, en lugar de lanzar kernels independentes, pódese construír un gráfico CUDA para capturar os metadatos necesarios para lanzar todos os kernels dun paso cara adiante cunha soa chamada ao controlador CUDA. Isto elimina a necesidade de volver executar código costoso de Python e PyTorch para as configuracións para as que se poden capturar gráficos CUDA.
En cada modelo, rastrexamos un punto de inflexión, determinando o número mínimo de tokens no que a execución na GPU é máis cara que o lanzamento de kernels do lado da CPU. Dado que os modelos de incrustación son pequenos, observamos que este punto de inflexión prodúcese en lotes de miles de tokens e decenas de secuencias. Algunhas implementacións de atención confían en entradas dinámicas do lado do host para configurar os lanzamentos de kernels, o que impide os gráficos CUDA densos/de preenchido de modelo completo. Integramos cambios nos kernels relevantes para habelos dispoñibles no noso motor de inferencia.
Para abordar as sobrecargas, construímos gráficos CUDA de modelo completo para todos os modelos de incrustación e solapamos o traballo da CPU co traballo da GPU. Dado que os gráficos CUDA minimizan as sobrecargas do lado da CPU, unha vez que se lanza un gráfico, temos tempo libre para iniciar e poñer en cola a execución do seguinte lote sempre que estea dispoñible. Os resultados do lote pendente rastrexanse cun LazyTensor, o que permite que unha tarefa asíncrona en Rust se bloquee ata que o lote anterior remate a execución. Os gráficos CUDA axudan ao servizo de baixa latencia ao garantir que non nos veñamos frenados polo custo dos lanzamentos de kernels e facilitan unha mellor programación no caso de alto rendemento, xa que liberan a CPU para facer traballo no seguinte lote antes.
Os gráficos CUDA deben capturarse para cada configuración distinta, o que para as incrustacións significa un gráfico por combinación de recuento de secuencias e recuento de tokens. Dado que esta grella é ampla, recheamos os recuentos de tokens a cubos que son múltiplos de 64 ou 256. Isto aínda dá como resultado miles de gráficos que poden levar varios minutos capturar para un modelo típico. O custo da captura provén de dúas fontes: un paso cara adiante ansioso que debe executarse para compilar kernels e configurar búferes para varios kernels que os necesitan, seguido da execución de captura que volve executar código Python.
Mitigamos os custos de inicio capturando gráficos CUDA de forma preguiceira a medida que o motor serve. Facemos un seguimento de cada configuración e asegurámonos de que pase por unha execución de quecemento ansiosa antes de activar a captura e reprodución do gráfico no segundo acerto. Todas as execucións posteriores da mesma configuración de gráfico pasan entón pola reprodución do gráfico CUDA. A captura de gráficos preguiceira ten un impacto nas latencias p99 durante o inicio; con todo, é valiosa para espallar varios minutos de traballo ansioso en varias horas. Uns tempos de inicio máis rápidos permítennos escalar e xestionar mellor os desprazamentos de incrustacións.
Tensores preguiceiros
A través de CUDA, o traballo da GPU é asíncrono. Dado que o lanzamento dun kernel de forma asíncrona ponno en cola nun fluxo, o código do host debe sincronizarse explícitamente para ler os vectores resultantes. Para facilitar un maior grao de paralelismo e poder iniciar futuros lotes mentres se agarda a que o anterior remate no dispositivo, confiamos nunha abstracción LazyTensor para rastrexar os valores.
O LazyTensor rastrexa un búfer de host en memoria bloqueada en páxina e unha operación cudaMemcpyAsync mediante un evento que copia datos desde o dispositivo. Iníciase despois do lanzamento do paso cara adiante no mesmo fluxo. Dado que a operación de copia debe agardar a que se executen todos os kernels anteriores no fluxo, o evento asociado rastrexa tanto a finalización do paso cara adiante como a dispoñibilidade do resultado na CPU.
Aproveitamos os LazyTensor no noso motor de codificación ROSE para solapar o traballo da GPU e da CPU. En lugar de que cada chamada a step() execute o gráfico CUDA e agarde a que remate, step() devolve un LazyTensor para rastrexar o seu resultado de forma asíncrona. Xunto cos gráficos CUDA, isto axúdanos a acadar latencias baixas e un mellor rendemento.
ROSE
Adaptamos o noso motor ROSE, construído orixinalmente para o servizo de LLM, para xestionar tamén a execución de modelos de incrustación. Para minimizar o esforzo necesario para admitir modelos de incrustación, ROSE reutiliza agresivamente código entre LLMs e incrustacións. Por exemplo, o servizo pplx-embed e a descodificación de LLM Qwen3.5 pasan polos mesmos kernels. Esta compartición permítenos servir facilmente un modelo de incrustación que foi axustado finamente orixinalmente desde un LLM para prototipos, avaliación e inferencia de produción.
Para capas densas, a inferencia de incrustacións e LLM son idénticas xa que os vectores de tokens procésanse de forma independente. nas capas de atención, as diferenzas xestiónanse engadindo compatibilidade con entradas irregulares, xunto coas configuracións de preenchido e descodificación paginadas que requiren os LLM. Ao servir un modelo de incrustación, non instanciamos unha caché KV e enviamos a variacións de kernels de atención que admiten o formato irregular para evitar o recheo. As rutinas de conversión e calibración de apoio tamén se comparten cos LLM.
Ivy
Ivy, a nosa capa de proxy HTTP de inferencia, tamén xoga un papel importante no rendemento. Debido a que as cargas útiles das solicitudes varían na produción, rotear solicitudes individuais a réplicas individuais pode causar desequilibrio de carga. Ivy divide as solicitudes de lotes grandes en anacos e equilibra a carga entre as réplicas, mellorando a utilización e suavizando a latencia. O noso traballo recente en tokenización de unigramas propia, implementada por completo en Ivy, mellora drasticamente as latencias en comparación cos tokenizadores pre-deseñados.
...pero os kernels seguen importando
ROSE admite unha variedade de backends de atención. Diferentes kernels poden ser axeitados para tamaños de problemas específicos. Co tempo, integramos os kernels FlashInfer 2, FlashInfer 3 e FlashAttention 4 para implementar a atención irregular.
En xeral, observamos que FlashAttention 4 é máis rápido. Con todo, FlashInfer 3 superáseo en modelos baseados en Qwen con lonxitudes de secuencia moi longas. Dado que o rendemento e a optimización poden variar co número e a dimensión das cabezas de atención, mantemos a compatibilidade con varias configuracións e tomamos unha decisión caso por caso ao servir.
Probas de rendemento
Facemos unha comparativa fronte a vLLM v0.22.0, executando a inferencia en precisión BF16 en pesos de modelos reais e entradas derivadas de conxuntos de datos de avaliación. Todas as execucións de tempos foron precedidas por execucións de quecemento que verificaron que a diverxencia na similitud de coseno está dentro do 0,1%.
Incrustacións de baixa latencia (p50 / p90 / p99 / máximo ms)
Informamos dos tempos de execución para un tamaño de lote de solicitudes pre-tokenizadas de 1, solicitudes totalmente secuencias, lonxitudes de secuencia de 128, 512 e 4096 tokens.
Puntuación de baixa latencia (p50 / p90 / p99 / máximo ms)
Tamaños de lote de solicitudes pre-tokenizadas 5, 25 e 50, lonxitude de secuencia de 512 tokens.
Incrustacións de alto rendemento (emb/s)
Tamaño de lote de solicitudes 100, catro procesos concorrentes enviando solicitudes, lonxitudes de secuencia de 512, 1024 e 4096 tokens.
Incrustacións de alta concorrencia (p50 / p90 / p99 / máximo ms)
Lonxitude de secuencia 512, tamaño de lote 1, pero enviamos 1, 2, 4, 8 e 16 solicitudes concorrentes. Esta proba de rendemento tamén inclúe os custos de tokenización a través de Ivy, xunto coa sobrecarga de rede entre Ivy e Tulip.
Conclusión e traballo futuro
A infraestrutura de servizo composta por Ivy, Tulip e ROSE permiten servir incrustacións para Perplexity cunha menor latencia e mellor rendemento, o que resulta nunha busca máis precisa a un custo reducido en comparación coas solucións pre-deseñadas.
Ao centrarnos en modelos específicos e asumir a propiedade de toda a pila, obtemos a liberdade necesaria para lograr un equilibrio efectivo entre rendemento e flexibilidade, mesturando primitivas en Rust altamente reutilizables e de alto rendemento xunto con código de modelado en Python máis xenérico. Moitos motores de inferencia de código aberto, como vLLM, SGLang e TokenSpeed, están a integrar linguaxes como Rust e C++ na súa pila. Investimos en Rust durante os últimos dous anos e obtivemos grandes beneficios tanto en rendemento como en mantemento. Ao compartir a maior parte da implementación de incrustacións coa nosa pila de servizo de LLM, tamén logramos ganancias de rendemento, sen requirir un esforzo de enxeñaría significativo dedicado ao mantemento de modelos de incrustación.
A medida que os modelos evolucionen, seguiremos mellorando cada capa da nosa pila para reducir tanto as latencias limitadas pola CPU como as limitadas pola GPU. Os nosos protocolos baseados en gRPC personalizados dentro de Ivy e Tulip permítenos axustar a comunicación para reducir as latencias de rede, mentres que ROSE proporciona unha base para mellorar o rendemento computacional. Ademais, a medida que o soporte para Python sen fíos libres creza en todo o ecosistema, poderemos mellorar aínda máis a interoperabilidade entre Python e Rust para reducir as sobrecargas.