Embeddings Rápidos em GPUs
A busca rápida e precisa é vital para toda a Perplexity, desde Search e Computer até nossa Plataforma de API. Nos bastidores, o trabalho pesado é feito por modelos de embedding e ranqueamento, que ajudam nossos sistemas a identificar os resultados mais relevantes para uma determinada consulta. Atingimos qualidade e latência de ponta
A busca rápida e precisa é vital para toda a Perplexity, desde Search e Computer até nossa Plataforma de API. Nos bastidores, o trabalho pesado é feito por modelos de embedding e ranqueamento, que ajudam nossos sistemas a identificar os resultados mais relevantes para uma determinada consulta. Atingimos qualidade e latência de ponta treinando e atendendo nossos próprios modelos, como o pplx-embed.
Este artigo apresenta uma visão dos bastidores da infraestrutura de atendimento da Perplexity para esta classe especial de modelos. Discutimos nossas técnicas para atender com eficiência às necessidades de inferência da busca nativa por IA, permitindo a prototipagem rápida e a avaliação de modelos enquanto alimentamos nosso índice de busca em escala de exabites. Essas técnicas expandem coletivamente a fronteira de Pareto de qualidade e eficiência de busca, permitindo-nos atender agentes e usuários com os melhores resultados possíveis ao menor custo e latência.
Embeddings para Busca
Em uma configuração de busca típica, os documentos indexados são mapeados para um espaço vetorial de alta dimensionalidade usando um modelo de embedding e armazenados em um banco de dados vetorial. Ao incorporar (embed) uma consulta usando o mesmo modelo, documentos semelhantes podem ser localizados encontrando os vetores mais próximos aos da consulta. Isso dá origem a dois padrões de tráfego diferentes para um mecanismo de inferência atender:
- Embedding em Lote: ao construir, expandir ou reindexar o banco de dados, documentos em massa devem ser incorporados no espaço vetorial, maximizando a taxa de transferência para minimizar o custo.
Após a busca vetorial, grandes lotes de documentos devem ser pontuados, encontrando um equilíbrio entre taxa de transferência e latência.
- Embedding Online: ao consultar o banco de dados, uma consulta curta deve ser incorporada para pesquisas, minimizando a latência.
Construímos nossa infraestrutura de inferência para aproveitar o máximo possível de componentes comuns em diferentes casos de uso. Como normalmente usamos pequenos modelos Transformer para produzir embeddings, compartilhamos a maior parte da implementação com nosso código de inferência de LLM: embeddings em lote são semelhantes ao preenchimento (prefill) limitado por computação, enquanto embeddings online, que geralmente são executados em alguns tokens, são computacionalmente semelhantes à decodificação limitada por memória. Portanto, reutilizamos nossos kernels otimizados de prefill e decode para atender a modelos de embedding. Como resultado, podemos alcançar uma enorme taxa de transferência de inferência em lote com o mínimo de trabalho de engenharia adicional, preservando a baixa latência para cargas de trabalho de embedding online.
Tulipas, Rosas e um pouco de Hera (Ivy)
Exibimos a inferência por meio de APIs padronizadas, tanto internamente quanto externamente através de nossa Plataforma de API. Nos bastidores, vários serviços estão envolvidos no processamento de uma solicitação de embedding:
- Ivy é um gateway HTTP em Rust que os serviços da Perplexity chamam.
Ele lida com o trabalho no lado da CPU para solicitações como análise de JSON, tokenização, criação de modelos de entrada (templating) e divisão de lotes, traduzindo solicitações para um protocolo gRPC personalizado para servidores subsequentes. Essa separação nos permite configurar determinados parâmetros relacionados à tokenização e à formatação de entrada sem precisar mexer nas instâncias de inferência mais pesadas.
- Tulip é a interface do servidor de inferência.
É um servidor gRPC implementado com Rust, tokio e `tonic`. O Tulip recebe solicitações de inferência gRPC, gerenciando o agendamento e o agrupamento em lotes (batching). Em seguida, ele envia os lotes para o mecanismo ROSE, retornando as respostas concluídas aos clientes.
- **ROSE** (Runtime-Optimized Serving Engine) implementa a inferência de modelos.
Ele é definido principalmente em Python, fornecendo kernels, camadas e definições para uma ampla variedade de modelos. O ROSE implementa as passagens diretas (forward passes) por meio de modelos, fornecendo também gerenciamento de grafos CUDA especializado para embeddings. Ele é conectado ao Tulip por meio de uma função step(), que aceita um lote e retorna uma referência à computação que ele realiza no acelerador.
Prestando Atenção Além do Kernel
Tanto os modelos baseados em Transformer quanto as arquiteturas Hopper/Blackwell subjacentes são tecnologias maduras, de modo que a inferência de incorporação (embedding) no lado da GPU convergiu para uma implementação amplamente ideal em vários mecanismos de inferência. Mesmo assim, descobrimos oportunidades adicionais de melhoria em tempos de execução e estruturas (harnesses) que expõem os modelos de ponta a ponta a um cliente. Em particular, descobrimos que podemos melhorar as latências gerenciando cuidadosamente os grafos CUDA e construindo uma abstração LazyTensor para rastrear de forma assíncrona um resultado no lado da GPU no mecanismo nativo Rust. Implementamos esses recursos no Tulip, para que ele pudesse fazer a interface de forma eficaz com as implementações de modelos do ROSE.
Tulip
Projetamos o Tulip para ser uma interface o mais leve possível sobre o nosso atendimento de modelos. Ele lida com solicitações recebidas em tarefas assíncronas do Tokio, mantendo um pool de solicitações que ele rastreia e a partir do qual agenda lotes para despachar para o acelerador. O mecanismo de agendamento no Tulip é muito simples: as solicitações se acumulam enquanto o Tulip está despachando trabalho ou aguardando resultados. A partir das solicitações acumuladas, as sequências são escolhidas por ordem de chegada para serem executadas através do modelo.
O mecanismo de agendamento simples é motivado por uma observação sobre o desempenho do modelo. Para pequenos modelos de embedding, nos comprimentos de sequência para os quais oferecemos atendimento, percebemos que o custo linear das camadas densas é dominante sobre o custo quadrático da atenção. Assim, a latência é principalmente proporcional ao número de tokens, e não ao número de sequências. Consequentemente, assim que um lote é grande o suficiente para saturar a GPU, o que gira em torno de 512 tokens em um modelo com menos de um bilhão de parâmetros, empacotar mais sequências nele não melhora a eficiência.
Para fazer a interface com o modelo de forma eficaz, o Tulip depende de grafos CUDA e rastreamento de resultados com preguiça (lazy) para sobrepor o trabalho da GPU e da CPU e utilizar totalmente os recursos disponíveis.
Gerenciamento de Grafos CUDA
A execução da passagem direta de um modelo envolve trabalho tanto no lado da CPU quanto no da GPU. A CPU é responsável por agendar lotes e iniciar kernels com os parâmetros apropriados, enquanto a GPU executa os kernels relevantes de multiplicação de matrizes, atenção, normalização (norm) ou ativação. Para cargas de trabalho de alta taxa de transferência, como treinamento e reindexação, as sobrecargas no lado da CPU são negligenciíveis porque os tamanhos dos lotes e a latência no lado da GPU são grandes. No entanto, em tamanhos de lote menores, o trabalho no lado da CPU pode superar o trabalho no lado da GPU.
Para mitigar as sobrecargas, em vez de lançar kernels independentes, um grafo CUDA pode ser construído para capturar os metadados necessários para lançar todos os kernels de uma passagem direta com uma única chamada ao driver CUDA. Isso elimina a necessidade de reexecutar código Python e PyTorch caro para as configurações para as quais os grafos CUDA podem ser capturados.
Em cada modelo, rastreamos um ponto de inflexão, determinando o número mínimo de tokens no qual a execução na GPU se torna mais cara do que a inicialização de kernel no lado da CPU. Como os modelos de embedding são pequenos, observamos que esse ponto de inflexão ocorre em lotes de milhares de tokens e dezenas de sequências. Algumas implementações de atenção dependem de entradas dinâmicas no lado do host para configurar inicializações de kernel, o que impede o uso de grafos CUDA densos ou de prefill de modelo completo. Fizemos o upstream de alterações nos kernels relevantes para habilitá-los em nosso mecanismo de inferência.
Para lidar com as sobrecargas, construímos grafos CUDA de modelo completo para todos os modelos de embedding e sobrepomos o trabalho da CPU com o trabalho da GPU. Como os grafos CUDA minimizam as sobrecargas no lado da CPU, assim que um grafo é lançado, temos tempo livre para iniciar e enfileirar a execução do próximo lote sempre que ele estiver disponível. Os resultados do lote pendente são rastreados com um LazyTensor, o que permite que uma tarefa assíncrona em Rust bloqueie até que o lote anterior termine a execução. Os grafos CUDA ajudam o atendimento de baixa latência, garantindo que não sejamos retidos pelo custo de inicializações de kernel, e facilitam um agendamento aprimorado no caso de alta taxa de transferência, pois liberam a CPU para trabalhar no próximo lote mais cedo.
Os grafos CUDA devem ser capturados para cada configuração distinta, o que para embeddings significa um grafo por combinação de contagem de sequências e contagem de tokens. Como essa grade é ampla, preenchemos as contagens de tokens para buckets que são múltiplos de 64 ou 256. Isso ainda resulta em milhares de grafos que podem levar vários minutos para serem capturados para um modelo típico. O custo da captura vem de duas fontes: uma passagem direta eager que deve ser executada para compilar os kernels e configurar buffers para vários kernels que precisam deles, seguida pela execução de captura que reexecuta o código Python.
Mitigamos os custos de inicialização capturando grafos CUDA com preguiça (lazily) à medida que o mecanismo atende às solicitações. Mantemos o controle de cada configuração e garantimos que ela passe por uma execução de aquecimento (warmup) eager antes de acionar a captura e a reprodução (replay) do grafo no segundo acesso. Todas as execuções subsequentes da mesma configuração de grafo passam então pela reprodução do grafo CUDA. A captura de grafo preguiçosa tem um impacto nas latências do p99 durante a inicialização; no entanto, ela é valiosa para distribuir vários minutos de trabalho eager ao longo de várias horas. Tempos de inicialização mais rápidos nos permitem dimensionar e gerenciar melhor as implantações de embedding.
Tensores com Preguiça (Lazy Tensors)
Através do CUDA, o trabalho da GPU é assíncrono. Como o lançamento de um kernel o enfileira assincronamente em um fluxo (stream), o código do host deve se sincronizar explicitamente para ler os vetores resultantes. Para facilitar um maior grau de paralelismo e poder iniciar lotes futuros enquanto aguarda o anterior concluir no dispositivo, contamos com uma abstração LazyTensor para rastrear valores.
O LazyTensor rastreia um buffer de host em memória travada de página (page-locked memory) e uma operação cudaMemcpyAsync por meio de um evento que copia dados do dispositivo. Ele é iniciado após o lançamento da passagem direta no mesmo fluxo (stream). Como a operação de cópia deve aguardar a execução de todos os kernels anteriores no fluxo, o evento associado rastreia tanto a conclusão da passagem direta quanto a disponibilidade do resultado na CPU.
Aproveitamos os LazyTensors em nosso mecanismo de codificador ROSE para sobrepor o trabalho da GPU e da CPU. Em vez de cada chamada step() executar o grafo CUDA e aguardar sua conclusão, step() retorna um LazyTensor para rastrear seu resultado de forma assíncrona. Combinado com os grafos CUDA, isso nos ajuda a alcançar baixas latências e melhor taxa de transferência.
ROSE
Adaptamos nosso mecanismo ROSE, construído originalmente para atendimento de LLM, para lidar também com a execução de modelos de embedding. Para minimizar o esforço necessário para dar suporte a modelos de embedding, o ROSE reutiliza código agressivamente entre LLMs e embeddings. Por exemplo, o atendimento pplx-embed e a decodificação de LLM Qwen3.5 passam pelos mesmos kernels. Esse compartilhamento nos permite atender facilmente a um modelo de embedding que foi originalmente ajustado (fine-tuned) a partir de um LLM para prototipagem, avaliação e inferência em produção.
Para camadas densas, a inferência de embedding e de LLM é idêntica, pois os vetores de token são processados de forma independente. Nas camadas de atenção, as diferenças são tratadas adicionando suporte a entradas irregulares (ragged), juntamente com as configurações de prefill e decode paginadas exigidas pelos LLMs. Ao atender a um modelo de embedding, não instanciamos um cache KV e despachamos para variações de kernels de atenção que dão suporte ao formato irregular para evitar o preenchimento (padding). As rotinas de conversão e calibração de suporte também são compartilhadas com os LLMs.
Ivy
O Ivy, nossa camada de proxy HTTP de inferência, também desempenha um papel importante no desempenho. Como as cargas úteis (payloads) de solicitação variam em produção, o direcionamento de solicitações individuais para réplicas individuais pode causar desbalanceamento de carga. O Ivy divide solicitações de grandes lotes em pedaços (chunks) e faz o balanceamento de carga delas entre as réplicas, melhorando a utilização e suavizando a latência. Nosso trabalho recente em tokenização de unigramas interna, totalmente implementada no Ivy, melhora drasticamente as latências em comparação com os tokenizadores prontos para uso.
...mas os Kernels Ainda Importam
O ROSE dá suporte a uma variedade de backends de atenção. Diferentes kernels podem ser adequados para tamanhos de problemas específicos. Com o tempo, integramos os kernels FlashInfer 2, FlashInfer 3 e FlashAttention 4 para implementar a atenção irregular (ragged attention).
Em geral, observamos que o FlashAttention 4 é mais rápido. No entanto, o FlashInfer 3 supera o desempenho dele em modelos baseados no Qwen com comprimentos de sequência muito longos. Como o desempenho e o ajuste podem variar de acordo com o número e a dimensão das cabeças de atenção, mantemos o suporte a várias configurações e tomamos uma decisão caso a caso durante o atendimento.
Benchmarks
Comparamos o desempenho com o vLLM v0.22.0, executando a inferência em precisão BF16 em pesos de modelo reais e entradas derivadas de conjuntos de dados de avaliação. Todas as execuções de cronometragem foram precedidas por execuções de aquecimento que verificaram se a divergência na similaridade de cosseno está dentro de 0,1%.
Embeddings de Baixa Latência (p50 / p90 / p99 / máx. em ms)
Relatamos tempos de execução para tamanho de lote de solicitação pré-tokenizada 1, solicitações totalmente sequenciais, comprimentos de sequência de 128, 512 e 4096 tokens.
Pontuação de Baixa Latência (p50 / p90 / p99 / máx. em ms)
Tamanhos de lote de solicitação pré-tokenizados de 5, 25 e 50, comprimento de sequência de 512 tokens.
Embeddings de Alta Taxa de Transferência (emb/s)
Tamanho do lote de solicitações de 100, quatro processos concorrentes enviando solicitações, comprimentos de sequência de 512, 1024 e 4096 tokens.
Embeddings de Alta Concorrência (p50 / p90 / p99 / máx. em ms)
Comprimento de sequência 512, tamanho de lote 1, mas enviamos 1, 2, 4, 8 e 16 solicitações concorrentes. Este benchmark também inclui os custos de tokenização por meio do Ivy, juntamente com a sobrecarga de rede entre o Ivy e o Tulip.
Conclusão e Trabalho Futuro
A infraestrutura de atendimento composta por Ivy, Tulip e ROSE nos permite fornecer embeddings para a Perplexity com menor latência e melhor taxa de transferência, resultando em uma busca mais precisa a um custo reduzido em comparação com soluções prontas para uso.
Ao focar em modelos específicos e assumir a responsabilidade por toda a pilha (stack), obtemos a liberdade necessária para encontrar um equilíbrio eficaz entre desempenho e flexibilidade, misturando primitivas em Rust altamente reutilizáveis e de alto desempenho com código de modelagem em Python mais genérico. Muitos mecanismos de inferência de código aberto, como vLLM, SGLang e TokenSpeed, estão integrando linguagens como Rust e C++ em sua pilha. Investimos em Rust nos últimos dois anos e colhemos grandes recompensas tanto em desempenho quanto em manutenibilidade. Ao compartilhar a maior parte da implementação de embedding com nossa pilha de atendimento de LLM, também obtemos ganhos na taxa de transferência, sem exigir que um esforço de engenharia significativo seja gasto na manutenção de modelos de embedding.
À medida que os modelos evoluem, continuaremos a melhorar cada camada de nossa pilha para reduzir as latências limitadas tanto por CPU quanto por GPU. Nossos protocolos personalizados baseados em gRPC dentro do Ivy e do Tulip nos permitem ajustar a comunicação para reduzir as latências de rede, enquanto o ROSE fornece uma base para melhorar a taxa de transferência computacional. Além disso, à medida que o suporte ao Python com threads livres (free-threaded) cresce em todo o ecossistema, poderemos melhorar ainda mais a interoperabilidade entre Python e Rust para reduzir as sobrecargas.