Incrustaciones rápidas en GPU
Una búsqueda rápida y precisa es vital para todo Perplexity, desde Search and Computer hasta nuestra plataforma API. Detrás de escena, el trabajo pesado lo realizan los modelos de incrustación y clasificación, que ayudan a nuestros sistemas a identificar los resultados más relevantes para una consulta dada. Logramos el estado del arte
Una búsqueda rápida y precisa es vital para todo Perplexity, desde Search and Computer hasta nuestra plataforma API. Detrás de escena, el trabajo pesado lo realizan los modelos de incrustación y clasificación, que ayudan a nuestros sistemas a identificar los resultados más relevantes para una consulta determinada. Logramos una calidad y latencia de última generación mediante el entrenamiento y la atención de nuestros propios modelos, como pplx-embed.
Este artículo presenta una visión interna de la infraestructura de servicio de Perplexity para esta clase especial de modelos. Analizamos nuestras técnicas para abordar de manera eficiente las necesidades de inferencia de la búsqueda nativa de IA, lo que permite la creación rápida de prototipos y la evaluación de modelos mientras potenciamos nuestro índice de búsqueda a escala de exabytes. Estas técnicas amplían colectivamente la frontera de Pareto de la calidad y la eficiencia de la búsqueda, lo que nos permite servir a los agentes y usuarios con los mejores resultados posibles al menor costo y latencia.
Incrustaciones para la búsqueda
En una configuración de búsqueda típica, los documentos indexados se asignan a un espacio vectorial de alta dimensión utilizando un modelo de incrustación y se almacenan en una base de datos vectorial. Al incrustar una consulta utilizando el mismo modelo, se pueden localizar documentos similares encontrando los vectores más cercanos al de la consulta. Esto da lugar a dos patrones de tráfico diferentes para que un motor de inferencia los sirva:
- Incrustación por lotes: al construir, expandir o reindexar la base de datos, los documentos masivos deben incrustarse en el espacio vectorial, maximizando el rendimiento para minimizar el costo.
Después de la búsqueda vectorial, se deben puntuar grandes lotes de documentos, logrando un equilibrio entre el rendimiento y la latencia.
- Incrustación en línea: al consultar la base de datos, se debe incrustar una consulta corta para las búsquedas, minimizando la latencia.
Construimos nuestra infraestructura de inferencia para aprovechar tantos componentes comunes como fuera posible en todos los casos de uso. Dado que normalmente utilizamos modelos Transformer pequeños para producir incrustaciones, compartimos la mayor parte de la implementación con nuestro código de inferencia de LLM: las incrustaciones por lotes son similares al prellenado limitado por cómputo, mientras que las incrustaciones en línea, que a menudo se ejecutan en unos pocos tokens, son computacionalmente similares a la decodificación limitada por memoria. Por lo tanto, reutilizamos nuestros kernels de prellenado y decodificación optimizados para servir modelos de incrustaciones. Como resultado, podemos lograr un rendimiento masivo de inferencia por lotes con un trabajo de ingeniería adicional mínimo, al tiempo que preservamos una baja latencia para las cargas de trabajo de incrustaciones en línea.
Tulips, Roses y un poco de Ivy
Exponemos la inferencia a través de API estandarizadas, tanto internamente como externamente a través de nuestra plataforma API. Detrás de escena, múltiples servicios participan en el procesamiento de una solicitud de incrustación:
- Ivy es una pasarela HTTP de Rust a la que llaman los servicios de Perplexity.
Maneja el trabajo del lado de la CPU para solicitudes como análisis de JSON, tokenización, plantillas de entrada y división de lotes, traduciendo las solicitudes a un protocolo gRPC personalizado para los servidores secundarios. Esta separación nos permite configurar ciertos parámetros relacionados con la tokenización y el formato de entrada sin tener que tocar las instancias de inferencia más pesadas.
- Tulip es la interfaz del servidor de inferencia.
Es un servidor gRPC implementado con Rust, tokio y `tonic`. Tulip recibe solicitudes de inferencia gRPC, gestionando la programación y el procesamiento por lotes. Luego envía los lotes al motor ROSE, devolviendo las respuestas completadas a los clientes.
- **ROSE** (Runtime-Optimized Serving Engine) implementa la inferencia de modelos.
Está definido principalmente en Python, y proporciona kernels, capas y definiciones para una amplia variedad de modelos. ROSE implementa los pasos hacia adelante a través de los modelos, proporcionando también gestión de grafos de CUDA especializada para incrustaciones. Se conecta a Tulip mediante una función step(), que toma un lote y devuelve una referencia al cálculo que realiza en el acelerador.
Prestando atención más allá del kernel
Tanto los modelos basados en Transformer como las arquitecturas subyacentes Hopper/Blackwell son tecnologías maduras, por lo que la inferencia de incrustaciones en el lado de la GPU ha convergido hacia una implementación en gran medida óptima en varios motores de inferencia. Aun así, descubrimos oportunidades adicionales de mejora en los tiempos de ejecución y las interfaces que exponen los modelos de extremo a extremo a un cliente. En particular, descubrimos que podemos mejorar las latencias gestionando cuidadosamente los grafos de CUDA y construyendo una abstracción LazyTensor para realizar un seguimiento asíncrono de un resultado en el lado de la GPU en el motor nativo de Rust. Implementamos estas características en Tulip para que pudiera interactuar eficazmente con las implementaciones de modelos de ROSE.
Tulip
Diseñamos Tulip para que fuera una interfaz lo más ligera posible sobre nuestro servicio de modelos. Maneja las solicitudes entrantes en tareas asíncronas de Tokio, manteniendo un grupo de solicitudes de las que realiza un seguimiento y programa lotes para enviar al acelerador. El mecanismo de programación en Tulip es muy simple: las solicitudes se acumulan mientras Tulip está enviando trabajo o esperando resultados. De las solicitudes acumuladas, las secuencias se eligen por orden de llegada para ejecutarse a través del modelo.
El mecanismo de programación simple está motivado por una observación sobre el rendimiento del modelo. Para los modelos de incrustación pequeños, en las longitudes de secuencia que servimos, notamos que el costo lineal de las capas densas es dominante sobre el costo cuadrático de atención. Por lo tanto, la latencia es principalmente proporcional al número de tokens, no al número de secuencias. En consecuencia, una vez que un lote es lo suficientemente grande como para saturar la GPU, lo que equivale a unos 512 tokens en un modelo de menos de mil millones de parámetros, empaquetar más secuencias en él no mejora la eficiencia.
Para interactuar eficazmente con el modelo, Tulip se basa en grafos de CUDA y el seguimiento de resultados perezosos para superponer el trabajo de la GPU y la CPU y utilizar plenamente los recursos disponibles.
Gestión de grafos de CUDA
La ejecución del paso hacia adelante de un modelo implica tanto el trabajo del lado de la CPU como el de la GPU. La CPU es responsable de programar los lotes y lanzar los kernels con los parámetros apropiados, mientras que la GPU ejecuta los kernels de multiplicación de matrices, atención, norma o activación relevantes. Para cargas de trabajo de alto rendimiento como el entrenamiento y la reindexación, las sobrecargas del lado de la CPU son insignificantes porque los tamaños de lote y la latencia del lado de la GPU son grandes. Sin embargo, en tamaños de lote más pequeños, el trabajo del lado de la CPU puede superar al trabajo del lado de la GPU.
Para mitigar las sobrecargas, en lugar de lanzar kernels independientes, se puede construir un grafo de CUDA para capturar los metadatos necesarios para lanzar todos los kernels de un paso hacia adelante con una sola llamada al controlador de CUDA. Esto elimina la necesidad de volver a ejecutar código costoso de Python y PyTorch para las configuraciones de las que se pueden capturar los grafos de CUDA.
En cada modelo, realizamos un seguimiento de un punto de inflexión, determinando el número mínimo de tokens en el que la ejecución en GPU es más costosa que el lanzamiento de kernels en el lado de la CPU. Debido a que los modelos de incrustaciones son pequeños, observamos que este punto de inflexión se produce en lotes de miles de tokens y decenas de secuencias. Algunas implementaciones de atención se basan en entradas dinámicas del lado del host para configurar los lanzamientos de kernels, lo que impide los grafos de CUDA densos o de prellenado de modelos completos. Hemos incorporado upstream cambios en los kernels relevantes para habilitarlos en nuestro motor de inferencia.
Para abordar las sobrecargas, construimos grafos de CUDA de modelo completo para todos los modelos de incrustaciones y superponemos el trabajo de la CPU con el trabajo de la GPU. Dado que los grafos de CUDA minimizan las sobrecargas del lado de la CPU, una vez que se lanza un grafo, tenemos tiempo libre para iniciar y poner en cola la ejecución del siguiente lote siempre que esté disponible. Los resultados del lote pendiente se rastrean con un LazyTensor, lo que permite que una tarea asíncrona en Rust se bloquee hasta que el lote anterior termine de ejecutarse. Los grafos de CUDA ayudan al servicio de baja latencia al garantizar que el costo de los lanzamientos de kernels no nos retrase y facilitan una mejor programación en el caso de alto rendimiento, ya que liberan a la CPU para que trabaje en el siguiente lote antes.
Los grafos de CUDA deben capturarse para cada configuración distinta, lo que para las incrustaciones significa un grafo por combinación de recuento de secuencias y recuento de tokens. Dado que esta cuadrícula es amplia, rellenamos los recuentos de tokens en depósitos que son múltiplos de 64 o 256. Esto aún da como resultado miles de grafos cuya captura puede tomar varios minutos para un modelo típico. El costo de la captura proviene de dos fuentes: un paso hacia adelante ansioso que debe ejecutarse para compilar kernels y configurar búferes para varios kernels que los necesitan, seguido de la ejecución de captura que vuelve a ejecutar el código de Python.
Mitigamos los costos de inicio capturando grafos de CUDA de manera perezosa a medida que el motor sirve. Realizamos un seguimiento de cada configuración y nos aseguramos de que pase por una ejecución de calentamiento ansiosa antes de activar la captura y reproducción del grafo en el segundo acierto. Todas las ejecuciones posteriores de la misma configuración de grafo pasan por la reproducción del grafo de CUDA. La captura de grafos perezosa tiene un impacto en las latencias p99 durante el inicio; sin embargo, es valiosa para distribuir varios minutos de trabajo ansioso en varias horas. Los tiempos de inicio más rápidos nos permiten escalar y gestionar mejor las implementaciones de incrustaciones.
Tensores perezosos
A través de CUDA, el trabajo de la GPU es asíncrono. Dado que el lanzamiento de un kernel de forma asíncrona lo pone en cola en un flujo, el código del host debe sincronizarse explícitamente para leer los vectores resultantes. Para facilitar un mayor grado de paralelismo y poder iniciar lotes futuros mientras se espera que el anterior se complete en el dispositivo, nos basamos en una abstracción LazyTensor para realizar un seguimiento de los valores.
El LazyTensor realiza un seguimiento de un búfer de host en memoria bloqueada en página y una operación cudaMemcpyAsync a través de un evento que copia datos desde el dispositivo. Se inicia después del lanzamiento del paso hacia adelante en el mismo flujo. Dado que la operación de copia debe esperar a que se ejecuten todos los kernels anteriores en el flujo, el evento asociado realiza un seguimiento tanto de la finalización del paso hacia adelante como de la disponibilidad del resultado en la CPU.
Aprovechamos los LazyTensors en nuestro motor de codificación ROSE para solapar el trabajo de la GPU y la CPU. En lugar de que cada llamada a step() ejecute el grafo de CUDA y espere a que termine, step() devuelve un LazyTensor para realizar un seguimiento asíncrono de su resultado. Junto con los grafos de CUDA, esto nos ayuda a lograr latencias bajas y un mejor rendimiento.
ROSE
Adaptamos nuestro motor ROSE, que construimos originalmente para el servicio de LLM, para manejar también la ejecución de modelos de incrustaciones. Para minimizar el esfuerzo necesario para admitir modelos de incrustaciones, ROSE reutiliza agresivamente el código entre los LLM y las incrustaciones. Por ejemplo, el servicio pplx-embed y la decodificación de LLM Qwen3.5 pasan por los mismos kernels. Esta compartición nos permite servir fácilmente un modelo de incrustación que se ajustó originalmente a partir de un LLM para la creación de prototipos, la evaluación y la inferencia de producción.
Para las capas densas, la inferencia de incrustaciones y de LLM es idéntica, ya que los vectores de tokens se procesan de forma independiente. En las capas de atención, las diferencias se manejan agregando soporte para entradas irregulares, junto con las configuraciones de prellenado y decodificación paginadas requeridas por los LLM. Al servir un modelo de incrustación, no instanciamos una caché KV y enviamos a variaciones de los kernels de atención que admiten el formato irregular para evitar el relleno. Las rutinas de conversión y calibración de soporte también se comparten con los LLM.
Ivy
Ivy, nuestra capa de proxy HTTP de inferencia, también desempeña un papel importante en el rendimiento. Debido a que las cargas útiles de las solicitudes varían en producción, el enrutamiento de solicitudes individuales a réplicas individuales puede causar un desequilibrio de carga. Ivy divide las solicitudes de lotes grandes en fragmentos y equilibra la carga entre réplicas, mejorando la utilización y suavizando la latencia. Nuestro trabajo reciente sobre la tokenización de unigramas propia, implementada por completo en Ivy, mejora drásticamente las latencias en comparación con los tokenizadores estándar.
...pero los kernels siguen importando
ROSE admite una variedad de backends de atención. Diferentes kernels pueden ser adecuados para tamaños de problemas específicos. Con el tiempo, integramos los kernels FlashInfer 2, FlashInfer 3 y FlashAttention 4 para implementar la atención irregular.
En general, observamos que FlashAttention 4 es más rápido. Sin embargo, FlashInfer 3 lo supera en modelos basados en Qwen con longitudes de secuencia muy largas. Dado que el rendimiento y la optimización pueden variar según el número y la dimensión de los cabezales de atención, mantenemos soporte para múltiples configuraciones y tomamos una decisión caso por caso al servir.
Pruebas comparativas
Realizamos pruebas comparativas frente a vLLM v0.22.0, ejecutando la inferencia en precisión BF16 en pesos de modelos reales y entradas derivadas de conjuntos de datos de evaluación. Todas las ejecuciones de temporización fueron precedidas por ejecuciones de calentamiento que verificaron que la divergencia en la similitud de coseno está dentro del 0,1%.
Incrustaciones de baja latencia (p50 / p90 / p99 / máx ms)
Informamos los tiempos de ejecución para un tamaño de lote de solicitudes pre-tokenizadas de 1, solicitudes totalmente secuenciales, longitudes de secuencia de 128, 512 y 4096 tokens.
Puntuación de baja latencia (p50 / p90 / p99 / máx ms)
Tamaños de lote de solicitudes pre-tokenizadas de 5, 25 y 50, longitud de secuencia de 512 tokens.
Incrustaciones de alto rendimiento (emb/s)
Tamaño de lote de solicitudes 100, cuatro procesos simultáneos que envían solicitudes, longitudes de secuencia de 512, 1024 y 4096 tokens.
Incrustaciones de alta concurrencia (p50 / p90 / p99 / máx ms)
Longitud de secuencia 512, tamaño de lote 1, pero enviamos 1, 2, 4, 8 y 16 solicitudes simultáneas. Esta prueba comparativa también incluye los costes de tokenización a través de Ivy, junto con la sobrecarga de red entre Ivy y Tulip.
Conclusión y trabajo futuro
La infraestructura de servicio compuesta por Ivy, Tulip y ROSE nos permite servir incrustaciones para Perplexity con menor latencia y mejor rendimiento, lo que resulta en una búsqueda más precisa a un costo reducido en comparación con las soluciones estándar.
Al centrarnos en modelos específicos y hacernos cargo de toda la pila, obtenemos la libertad necesaria para lograr un equilibrio eficaz entre rendimiento y flexibilidad, combinando primitivas de Rust altamente reutilizables y de alto rendimiento junto con código de modelado de Python más genérico. Muchos código abierto motores de inferencia, como vLLM, SGLang y TokenSpeed, están integrando lenguajes como Rust y C++ en su pila. Hemos invertido en Rust durante los últimos dos años y hemos obtenido grandes recompensas tanto en rendimiento como en capacidad de mantenimiento. Al compartir la mayor parte de la implementación de incrustaciones con nuestra pila de servicio de LLM, también obtenemos ganancias en rendimiento, sin necesidad de dedicar un esfuerzo de ingeniería significativo al mantenimiento de modelos de incrustaciones.
A medida que los modelos evolucionen, continuaremos mejorando cada capa de nuestra pila para reducir las latencias limitadas tanto por la CPU como por la GPU. Nuestros protocolos basados en gRPC personalizados dentro de Ivy y Tulip nos permiten ajustar la comunicación para reducir las latencias de red, mientras que ROSE proporciona una base para mejorar el rendimiento computacional. Además, a medida que crezca el soporte para Python sin bloqueo de subprocesos en todo el ecosistema, podremos mejorar aún más la interoperabilidad entre Python y Rust para reducir las sobrecargas.