La calidad de recuperación en un producto de búsqueda de IA está limitada por dos factores: la calidad del modelo de embedding y el costo de ejecución a través de un índice. Esta semana, el equipo de ingeniería de Perplexity publicó Fast Embeddings on GPUs, un relato técnico sobre el segundo punto: la infraestructura de serving detrás de pplx-embed y los modelos de ranking utilizados en Perplexity Search, Computer y la plataforma API.

El equipo de Perplexity afirma que la inferencia de embeddings en el lado de la GPU ha convergido en gran medida entre motores en hardware maduro como Hopper y Blackwell. Las ganancias residen en el runtime y el entorno que rodea al modelo: gestión de gráficos CUDA, una abstracción de seguimiento de resultados asíncrona y una ruta de solicitud en Rust.

Dos patrones de tráfico, un solo motor

Perplexity define el servicio de embeddings como dos cargas de trabajo distintas. El batch embedding ocurre al construir o reindexar la base de datos vectorial, donde el throughput minimiza el costo. El online embedding ocurre en el momento de la consulta, donde una búsqueda corta debe procesarse rápidamente. El scoring se sitúa en medio: tras la búsqueda vectorial, se clasifican grandes lotes de documentos, equilibrando ambas necesidades.

La decisión clave es que Perplexity no construyó un motor de embedding separado. Dado que los modelos de embedding son Transformers pequeños, el batch embedding se asemeja al prefill limitado por cómputo, y el online embedding, a menudo de pocos tokens, se asemeja al decode limitado por memoria. Por ello, el equipo de investigación reutiliza los prefill and decode kernels de su stack de LLM.

Ivy, Tulip y ROSE

Tres servicios gestionan cada solicitud:

  • Ivy es una puerta de enlace HTTP en Rust. Realiza el trabajo del lado de la CPU: análisis JSON, tokenización, plantillado de entrada, división de lotes y traduce las solicitudes a un protocolo gRPC personalizado. También divide las solicitudes de lotes grandes en fragmentos y los balancea entre réplicas, lo que corrige el desequilibrio de carga que surge cuando los payloads de producción varían en tamaño.
  • Tulip es la interfaz del servidor de inferencia: un servidor gRPC construido con Rust, tokio y tonic, que maneja la programación y el procesamiento por lotes antes de enviarlos al motor.
  • ROSE (Runtime-Optimized Serving Engine) implementa la inferencia del modelo. Es principalmente Python, proporciona kernels, capas y definiciones de modelos, gestiona gráficos CUDA y expone una función step() para Tulip.

¿Por qué el programador es deliberadamente simple?

Tulip elige secuencias bajo el principio de primero en entrar, primero en salir (FIFO) mientras las solicitudes se acumulan. Esta simplicidad se justifica con una medición: para modelos de embedding pequeños a las longitudes de secuencia que sirve Perplexity, el costo lineal de las capas densas domina el costo cuadrático de la atención. La latencia, por lo tanto, es aproximadamente proporcional al conteo de tokens, no al conteo de secuencias. Una vez que un lote satura la GPU, alrededor de 512 tokens en un modelo de menos de mil millones de parámetros, añadir más secuencias no mejora la eficiencia.

Gráficos CUDA y LazyTensors

En lotes pequeños, el lanzamiento de kernels desde la CPU puede superar la ejecución en la GPU. Perplexity construye gráficos CUDA de modelo completo para todos los modelos de embedding, capturando cada lanzamiento en una sola llamada al driver. Debido a que los modelos de embedding son pequeños, el punto de inflexión donde el trabajo de la GPU excede el costo de lanzamiento llega con lotes de miles de tokens y decenas de secuencias. Algunas implementaciones de atención bloquean los gráficos de modelo completo al depender de entradas dinámicas del host; Perplexity propuso cambios a FlashInfer para permitir la captura.

Los gráficos deben capturarse por configuración, por lo que los conteos de tokens se rellenan a buckets que son múltiplos de 64 o 256. Esto genera miles de gráficos y varios minutos de captura por modelo. La solución es la captura perezosa (lazy capture): cada configuración recibe una ejecución de calentamiento (warmup) ansiosa, luego activa la captura y repetición en su segundo hit. Esto afecta la latencia p99 al inicio, pero distribuye minutos de trabajo ansioso a lo largo de horas.

El segundo componente es el LazyTensor, que rastrea un buffer de host bloqueado en página (page-locked) más un cudaMemcpyAsync y un evento CUDA. En lugar de que step() se bloquee en el dispositivo, devuelve un LazyTensor, permitiendo que una tarea asíncrona de Rust espere por el lote N mientras la CPU encola el N+1.

Los kernels siguen siendo fundamentales

ROSE admite múltiples backends de atención para entradas irregulares: FlashInfer 2, FlashInfer 3 y FlashAttention 4. El equipo de Perplexity informa que FlashAttention 4 es generalmente más rápido, pero FlashInfer 3 lo supera en modelos basados en Qwen en longitudes de secuencia muy largas, por lo que la selección del backend se hace caso por caso. Cabe destacar que, al servir un modelo de embedding, ROSE no instancia un KV cache y despacha a variantes de atención irregular para evitar el relleno (padding).

Benchmarks

Perplexity realiza benchmarks contra vLLM v0.22.0 en BF16 con pesos reales y entradas derivadas de evaluación, con ejecuciones de calentamiento que verifican que la divergencia de similitud de coseno esté dentro del 0.1%. Se grafican cuatro suites: embeddings de baja latencia (lote 1; 128/512/4096 tokens), scoring de baja latencia (lote 5/25/50 a 512 tokens), embeddings de alto throughput (lote 100, cuatro procesos concurrentes) y embeddings de alta concurrencia (1 a 16 solicitudes concurrentes, incluyendo la tokenización de Ivy y la sobrecarga de red).

Conclusiones clave

  • El stack de embeddings de Perplexity reutiliza sus kernels de prefill/decode de LLM en lugar de ejecutar un motor separado.
  • La latencia sigue el conteo de tokens, no el de secuencias; ~512 tokens saturan un modelo de menos de 1B de parámetros.
  • Los gráficos CUDA de modelo completo más la captura perezosa reducen la sobrecarga de lanzamiento sin arranques de minutos de duración.
  • LazyTensor solapa la preparación de lotes en CPU con el trabajo en GPU en vuelo en lugar de bloquearse en sincronización.
  • Ivy, Tulip y ROSE son internos; pplx-embed es accesible a través de la API de Embeddings de Perplexity.

Consulte los detalles técnicos. Vía MarkTechPost.