Incrustaciones rápidas en GPU
Una búsqueda rápida y precisa es vital para todo Perplexity, desde Search y Computer hasta nuestra plataforma de 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 la última generación de
Una búsqueda rápida y precisa es vital para todo Perplexity, desde Search y Computer hasta nuestra plataforma de 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 una latencia de vanguardia mediante el entrenamiento y la puesta en servicio 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, permitiendo 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 ofrecer a los agentes y usuarios los mejores resultados posibles al menor costo y latencia.
Incrustaciones para búsquedas
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 atienda:
- Incrustación por lotes (Batch Embedding): 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 (throughput) y la latencia.
- Incrustación en línea (Online Embedding): 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 entre los casos de uso. Como normalmente usamos modelos Transformer pequeños para producir incrustaciones (embeddings), compartimos la mayor parte de la implementación con nuestro código de inferencia de LLM: las incrustaciones por lotes son similares al llenado previo (prefill) limitado por el 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 la memoria. Por lo tanto, reutilizamos nuestros kernels optimizados de prefill y decodificación para servir a los modelos de incrustación. 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 incrustación en línea.
Tulips, Roses y algo de Ivy
Exponemos la inferencia a través de API estandarizadas, tanto internamente como externamente a través de nuestra plataforma de API. Bajo el capó, participan múltiples servicios en el procesamiento de una solicitud de incrustación:
- Ivy es una pasarela HTTP en Rust a la que llaman los servicios de Perplexity.
Maneja el trabajo del lado de la CPU para solicitudes como el análisis de JSON, la tokenización, la creación de plantillas de entrada y la 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 y devuelve 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 las pasadas hacia adelante (forward passes) a través de los modelos, y también proporciona gestión de grafos de CUDA especializada para incrustaciones (embeddings). 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 Transformers como las arquitecturas Hopper/Blackwell subyacentes son tecnologías maduras, por lo que la inferencia de incrustaciones (embeddings) 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 los entornos (harnesses) 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 de LazyTensor para rastrear de forma asíncrona un resultado en el lado de la GPU dentro del 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 sea 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 (pool) de solicitudes que rastrea y a partir del cual 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. A partir 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 la atención. Por lo tanto, la latencia es mayormente 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 seguimiento de resultados perezosos para solapar el trabajo de la GPU y la CPU y utilizar al máximo los recursos disponibles.
Gestión de grafos de CUDA
La ejecución de la pasada hacia adelante (forward pass) de un modelo implica trabajo tanto en el lado de la CPU como en 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 relevantes de multiplicación de matrices, atención, norma o activación. 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 los lotes 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 una pasada hacia adelante con una sola llamada al controlador de CUDA. Esto elimina la necesidad de volver a ejecutar costoso código de Python y PyTorch para las configuraciones de las que se pueden capturar grafos de CUDA.
En cada modelo, rastreamos un punto de inflexión, determinando el número mínimo de tokens en el que la ejecución en la GPU es más costosa que el lanzamiento de kernels en el lado de la CPU. Debido a que los modelos de incrustación 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 el uso de grafos de CUDA densos o de prefill de modelo completo. Integramos cambios en los kernels relevantes para habilitarlos en nuestro motor de inferencia.
Para abordar las sobrecargas, construimos grafos de CUDA de todo el modelo para todos los modelos de incrustación y solapamos 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 no nos veamos frenados por el costo de los lanzamientos de kernels y facilitan una mejor programación en el caso de alto rendimiento, ya que liberan a la CPU para que pueda trabajar 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 cada combinación de recuento de secuencias y recuento de tokens. Como esta cuadrícula es amplia, rellenamos los recuentos de tokens en depósitos (buckets) que son múltiplos de 64 o 256. Esto aún da como resultado miles de grafos cuya captura puede llevar varios minutos para un modelo típico. El costo de la captura proviene de dos fuentes: una pasada hacia adelante ansiosa (eager) que debe ejecutarse para compilar los kernels y configurar los búferes para los diversos kernels que los necesitan, seguida de la ejecución de captura que vuelve a ejecutar el código de Python.
Mitigamos los costos de inicio capturando los grafos de CUDA de forma perezosa a medida que el motor presta servicio. Realizamos un seguimiento de cada configuración y nos aseguramos de que pase por una ejecución de calentamiento ansiosa (eager warmup) antes de activar la captura y reproducción del grafo en el segundo intento. Todas las ejecuciones posteriores de la misma configuración de grafo pasan entonces por la reproducción del grafo de CUDA. La captura de grafos perezosa tiene un impacto en las latencias del percentil 99 (p99) durante el inicio; sin embargo, es valiosa para distribuir varios minutos de trabajo ansioso a lo largo de varias horas. Los tiempos de inicio más rápidos nos permiten escalar y gestionar mejor los despliegues de incrustaciones.
Tensores perezosos (Lazy Tensors)
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 una secuencia, el código del host debe sincronizarse explícitamente para leer los vectores resultantes. Para facilitar un mayor grado de poder de procesamiento paralelo y poder iniciar lotes futuros mientras se espera a que se complete el anterior en el dispositivo, confiamos en una abstracción de LazyTensor para rastrear los valores.
El LazyTensor rastrea un búfer de host en memoria bloqueada en página (page-locked) y una operación cudaMemcpyAsync a través de un evento que copia datos desde el dispositivo. Se inicia después del lanzamiento de la pasada hacia adelante en la misma secuencia (stream). Dado que la operación de copia debe esperar a que se ejecuten todos los kernels anteriores en la secuencia, el evento asociado rastrea tanto la finalización de la pasada hacia adelante como la disponibilidad del resultado en la CPU.
Aprovechamos los LazyTensors en nuestro motor de codificador 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 rastrear su resultado de forma asíncrona. Junto con los grafos de CUDA, esto nos ayuda a lograr latencias bajas y un mejor rendimiento.
ROSE
Adaptamos nuestro motor ROSE, que originalmente construimos para el servicio de LLM, para que también maneje la ejecución de modelos de incrustación. Para minimizar el esfuerzo necesario para dar soporte a los modelos de incrustación, 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. Este uso compartido nos permite servir fácilmente un modelo de incrustación que se ajustó originalmente (fine-tuned) a partir de un LLM para la creación de prototipos, la evaluación y la inferencia en producción.
Para las capas densas, la inferencia de incrustación y de LLM son idénticas, ya que los vectores de tokens se procesan de manera independiente. En las capas de atención, las diferencias se manejan agregando soporte para entradas irregulares (ragged), junto con las configuraciones de prefill y decodificación paginadas requeridas por los LLM. Al servir un modelo de incrustación, no instanciamos una caché KV y nos remitimos a variantes de kernels de atención que admiten el formato irregular para evitar el relleno (padding). 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, enrutar 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 las 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 (ragged attention).
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 grandes. Dado que el rendimiento y la optimización pueden variar según el número y la dimensión de las cabezas de atención, mantenemos soporte para múltiples configuraciones y tomamos una decisión caso por caso durante el servicio.
Pruebas comparativas (Benchmarks)
Realizamos pruebas comparativas frente a vLLM v0.22.0, ejecutando la inferencia en precisión BF16 sobre pesos de modelos reales y entradas derivadas de conjuntos de datos de evaluación. Todas las ejecuciones de cronometraje fueron precedidas por ejecuciones de calentamiento que verificaron que la divergencia en la similitud de coseno estuviera 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 pretokenizadas de 1, solicitudes totalmente secuenciales y 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 pretokenizadas de 5, 25 y 50, con una longitud de secuencia de 512 tokens.
Incrustaciones de alto rendimiento (emb/s)
Tamaño de lote de solicitudes de 100, cuatro procesos concurrentes enviando solicitudes, longitudes de secuencia de 512, 1024 y 4096 tokens.
Incrustaciones de alta concurrencia (p50 / p90 / p99 / máx. ms)
Longitud de secuencia de 512, tamaño de lote 1, pero enviamos 1, 2, 4, 8 y 16 solicitudes concurrentes. Este punto de referencia (benchmark) también incluye los costos 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 tecnológica, 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 motores de inferencia de código abierto, 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 el rendimiento, sin necesidad de dedicar un esfuerzo de ingeniería significativo al mantenimiento de modelos de incrustación.
A medida que los modelos evolucionen, continuaremos mejorando cada capa de nuestra pila para reducir tanto las latencias limitadas por la CPU como las limitadas por la GPU. Nuestros protocolos personalizados basados en gRPC 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 de hilos libres (free-threaded) en todo el ecosistema, podremos mejorar aún más la interoperabilidad entre Python y Rust para reducir las sobrecargas.