Prefill e decode desagregados

A separação de prefill e decode entre dispositivos acelera a inferência de LLM e melhora o desempenho

AutoresPerplexity Research

Para gerar tokens de saída a partir de um prompt de entrada, a inferência de LLM é dividida em dois estágios: prefill e decode. O prefill é executado nos tokens de entrada, preenchendo os caches KV, antes de entrar no estágio de decode, que gera tokens um a um. Embora uma única etapa de decode seja executada tipicamente em dezenas de milissegundos, o prefill leva muito mais tempo. Se executado nos mesmos dispositivos, misturar prefill com decode degrada o desempenho do decode. Neste artigo, exploramos uma solução estabelecida na forma de prefill e decode desagregados, executando-os em dispositivos separados para maximizar tanto a taxa de transferência do prefill quanto as latências do decode.

Desempenho de Prefill vs Decode

Em um mecanismo de atendimento de LLM típico, o agendador de lotes seleciona solicitações para processar em cada etapa de execução de um modelo. Ao executar em um único dispositivo ou nó, as solicitações de prefill e decode são agrupadas em lotes. O custo da atenção, que se agrega ao longo do comprimento da sequência, cresce tanto para o prefill quanto para o decode, proporcionalmente ao comprimento das entradas no cache KV (kv_len). As solicitações de decode geralmente encaminham um único token (qo_len=1), a um custo mínimo por meio de outras camadas que operam independentemente nos tokens de uma sequência. As solicitações de prefill encaminham milhares ou dezenas de milhares de tokens a um custo significativo por meio de camadas densas (qo_len grande).

A latência de uma passagem de encaminhamento é mais fortemente influenciada pelo número de tokens independentes passados por camadas densas (qo_len) do que pelo número de tokens recuperados do cache KV durante a atenção (kv_len). A atenção pode ser paralelizada tanto pelo número de solicitações quanto pelo kv_len proporcional aos comprimentos das sequência, alcançando uma boa utilização. O prefill é limitado por computação: com qo_len alto, os kernels GEMM podem alocar blocos suficientes ao longo da dimensão M para utilizar totalmente as capacidades computacionais das GPUs modernas. O decode é limitado por memória: devido a tamanhos de lote tipicamente baixos, o número de entradas ao longo de M geralmente é pequeno, suficiente para apenas um bloco. Embora os kernels Split-K GEMM possam melhorar a utilização da SM para tamanhos de lote de tokens baixos, os caches e as unidades de multiplicação de matriz geralmente permanecem subutilizados.

Batching de requisições por meio de uma camada transformer, mostrando projeções de query e key/value, atenção e atualizações do cache KV.

Quando misturados, os lotes que contêm solicitações de prefill incorrem em latências mais altas na passagem de encaminhamento, afetando negativamente a taxa de transferência de decode de toda a instância. Embora misturar solicitações de prefill com solicitações de decode ou empregar prefill em partes possa melhorar ligeiramente o desempenho do decode, é difícil manter uma taxa de transferência de prefill suficiente para processar solicitações suficientes em uma instância para maximizar a taxa de transferência de decode. No caso de grandes modelos, com comprimentos de saída típicos, para manter um tamanho de lote grande para o decode, o prefill deve ser realizado com frequência suficiente para degradar significativamente a latência média e causar engasgos na saída.

Agendamento misto de prefill–decode para três requisições, mostrando intervalos de prefill e decode em uma linha do tempo compartilhada.
Agendamento desagregado de prefill–decode, com um prefiller dedicado alimentando o trabalho de decode para três requisições.

Esses problemas podem ser resolvidos usando um conjunto separado de nós para realizar o prefill e o decode. Ao associar um nó prefiller a vários nós decodificadores, solicitações suficientes podem ser agendadas para prefill para maximizar a taxa de transferência e manter um número suficientemente grande de solicitações concorrentes nos nós decodificadores para também maximizar a taxa de transferência de decode. Os nós prefiller preenchem os caches KV, que são então transferidos para os nós decodificadores. Como os decodificadores não precisam mais parar para o prefill, as latências tornam-se muito mais determinísticas, já que o impacto geral do aumento do kv_len das solicitações ativas é muito menos pronunciado. O custo é pago em um aumento no Tempo para o Primeiro Token (TTFT), pois a transferência de caches KV pela rede pode levar dezenas a centenas de milissegundos.

KV Messenger

Na Perplexity, nossa implementação para prefill e decode desagregados é construída em torno de um KV messenger que interage com o mecanismo LLM para orquestrar transferências de cache KV de nós prefiller para os nós decodificadores através de uma rede. No lado do prefiller, o messenger aceita solicitações de nós decodificadores, repassando-as para o agendador de lotes e monitorando a execução da passagem de encaminhamento para despachar caches KV com o mínimo de latência possível. No lado do decodificador, após a alocação de páginas não passiveis de remoção, o messenger bloqueia a solicitação de ser agendada para decode até que seja notificado da conclusão do cache KV e das transferências de contexto do decodificador.

Um ciclo de vida de requisição desagregado, mostrando a coordenação entre o decoder e o prefiller e a transferência de dados do cache KV.

A desagregação do prefill exige conexões de alta taxa de transferência e baixa latência, portanto, nossa implementação é adaptada para RDMA, suportando controladores de interface de rede (NICs) EFA e ConnectX. O KV Messenger é construído sobre libfabric, usando nossos wrappers fabric-lib para fornecer abstrações de alto nível e baixa latência sobre as primitivas de Acesso Direto à Memória Remota (RDMA), implementando transferências eficientes de páginas e metadados, juntamente com sinalização de baixa latência. Em segundo plano, o fabric-lib coordena uma GPU e suas NICs diretamente conectadas para copiar dados do nó prefiller para o nó decodificador.

Após o recebimento, o nó prefiller aloca um conjunto correspondente de páginas KV de origem e agenda a solicitação de prefill usando seu mecanismo local. Para minimizar a latência, as transferências não esperam pela passagem de encaminhamento: em vez disso, as cópias de páginas KV são iniciadas assim que o modelo termina de anexar as entradas de cache KV ao cache KV para camadas individuais. Como as solicitações de prefill podem ser divididas em partes, o agendador de lotes notifica o KV messenger sobre as partes atualmente agendadas antes da execução. Para suportar gráficos CUDA enquanto é capaz de rastrear camadas, o messenger mantém um tópico dedicado pesquisando um contador incrementado após a projeção de saída da atenção. O contador é mantido apenas no nó principal em um ambiente fragmentado: embora as entradas de cache KV sejam válidas após a anexação e antes da atenção, a projeção de saída é reduzida entre as classificações, sincronizando-as implicitamente. Assim que uma alteração no contador é observada, o messenger é notificado e chama fabric-lib para iniciar a transferência de uma camada.

Transferências de cache KV camada por camada, mostrando gravações paginadas e únicas e contadores de conclusão entre o prefiller e o decoder.

Após a conclusão da transferência da última parte, quaisquer metadados adicionais também são copiados: a decodificação especulativa ou MTP exigem que logits e estados ocultos sejam movidos para o decodificador. Essas cópias também são realizadas por meio de RDMA, para e de buffers pré-alocados.

Após a conclusão de todas as transferências pendentes do último bloco, o nó de preenchimento prévio desaloca as páginas KV e conclui a solicitação. O nó do decodificador não é notificado explicitamente: em vez disso, ele usa contadores imediatos para acompanhar o número de operações concluídas. O número de operações RDMA no lado do preenchimento prévio é proporcional ao número de páginas transferidas. Após a conclusão do número conhecido de cópias de páginas e contexto, fabric-lib chama o mensageiro KV para indicar que uma solicitação está pronta para decodificação. O mensageiro desaloca qualquer contexto e encaminha a solicitação ao mecanismo de LLM.

Transferências de Cache KV Fragmentado

Se o prefiller e o decodificador dependerem de Paralelismo de Tensor (TP) e fragmentarem ou replicarem os caches KV de forma idêntica, um único mecanismo de transferência coordena vários dispositivos para enviar e receber as páginas de todas as réplicas. Para poder usar um único messenger e mecanismo de transferência, apesar do fato de o executor do modelo ser replicado em vários dispositivos e processos, cuMem e cuMemImportFromShareableHandle são usados para alocar a memória do dispositivo que dá suporte aos caches KV e mapeá-la para o processo principal. O mecanismo de transferência inspeciona a topologia do nó para encontrar as NICs e as CPUs no nó NUMA mais próximo para usar nas transferências de cada uma das fatias de cache KV.

Se a origem e o destino forem fragmentados de forma idêntica, as transferências serão triviais, pois há um mapeamento de um para um entre os dispositivos e as páginas da origem e do destino. Nessa situação, a fragmentação ajuda implicitamente nas latências de transferência: ao usar mais GPUs, mais NICs associadas podem ser empregadas, aproximando-se da utilização total da largura de banda. No entanto, se houver uma incompatibilidade, o mecanismo de transferência deve dividir ou reconstruir páginas dependendo da proporção entre as fatias de origem e destino.

Topologia de transferência de cache KV entre um prefiller TP/EP e decoders DP/EP, mostrando as conexões NIC.

Se o prefiller dividir o cache KV em mais dispositivos, páginas completas serão reconstruídas no decodificador enviando as metades correspondentes a partir dos dispositivos prefiller. Se o decodificador tiver mais fragmentos, ele receberá páginas de várias origens. O decodificador precisa conhecer o esquema de fragmentação do prefiller para poder calcular o número de gravações RDMA que espera receber. Se a replicação estiver envolvida, o prefiller agrupa os dispositivos em conjuntos de réplicas que replicam o cache KV completo dentro de si mesmos. Os conjuntos de réplicas de destino recebem aleatoriamente um dos conjuntos de origem para usar todos os dispositivos disponíveis para iniciar gravações RDMA.

Layouts de cache KV NHD e HND, contrastando a ordem de tokens e heads em cada camada e página.

As transferências fragmentadas exigem um pequeno ajuste nos caches KV. Por padrão, o FlashInfer depende do layout NHD, que ordena os tokens dentro de uma página dentro dos cabeçotes. Como os caches provavelmente estão fragmentados ao longo do número de cabeçotes de atenção, isso cria descontinuidade dentro do cabeçote. As transferências RDMA não suportam implicitamente gravações com saltos, exigindo uma operação por cabeçote para realizar a transferência. Em vez disso, para reduzir o número de interações com o libfabric, organizamos os caches KV usando o layout HND, que coloca a dimensão do cabeçote antes do número de tokens. Isso garante a continuidade, permitindo que uma página seja copiada com uma única gravação.

Decodificação Especulativa

A decodificação especulativa exige pequenos ajustes no prefill-decode desagregado. Em nossa implementação, os nós prefiller não têm permissão para amostrar tokens. Como os modelos Sonar da Perplexity suportam saída estruturada, não queremos incorrer na complexidade de sincronizar as implementações do processador de esquema entre prefillers e decodificadores. Nos mecanismos MTP e de decodificação especulativa, o prefill do modelo de rascunho até o último token envolve a amostragem de tokens do modelo de destino.

Decodificação especulativa entre o prefiller e o decoder, mostrando o comprimento do cache KV, logits e amostragem de tokens.

Para contornar esses problemas, o prefill não inclui o último token da sequência de entrada. Em vez disso, os estados ocultos ou logits do prefill anteriores ao último token são transferidos e tratados como um token de decode na próxima etapa do decodificador. Embora isso aumente ligeiramente as latências, já que uma etapa de decode completa deve ser realizada após o prefill para emitir o primeiro token, a complexidade da implementação é bastante reduzida.

Implantações Desagregadas

Implantamos ou experimentamos várias configurações desagregadas com diferentes modelos, para dar suporte ao tráfego de produção ou a cargas de trabalho de avaliação interna. Com base no tamanho e no mecanismo de atenção dos modelos, escolhemos esquemas de fragmentação adequados para os nós prefiller e decodificador para utilizar melhor as GPUs.

DeepSeek-R1

Com o DeepSeek, consideramos implantações de Paralelismo de Tensor (TP) e Paralelismo de Dados (DP). Conforme discutido em postagens de blog anteriores, as implantações de TP fornecem melhor latência ao custo de menor taxa de transferência, exigindo mais GPUs para atender a tráfego pesado. As implantações de DP escalam muito melhor com a carga, no entanto, sua taxa de transferência de pico é menor devido ao custo de comunicação entre dispositivos ou entre nós.

O DeepSeek depende da Atenção Latente Multi-Head, compactando os caches KV. Como todos os cabeçotes KV são compactados em um único vetor latente, o TP não pode fragmentar os caches KV, pois deve replicar os vetores latentes em todas as classificações. A fragmentação ocorre após a descompactação, pois cada classificação pode extrair diferentes cabeçotes da mesma representação latente. Consequentemente, todos os fragmentos de cache KV são idênticos em ambos os fragmentos de prefiller e decodificador.

Com uma configuração de TP intranode, tanto os prefillers quanto os decodificadores são fragmentados de forma idêntica. As transferências são despachadas de todas as classificações para utilizar totalmente todas as NICs disponíveis. No entanto, com uma implantação de DP, onde o tamanho da classificação de TP é menor ou cada classificação de DP é atribuída a uma única GPU, qualquer dispositivo prefiller que possua uma cópia replicada do cache KV pode despachá-lo. Para equilibrar as solicitações em todas as NICs disponíveis, selecionamos aleatoriamente uma GPU e uma NIC para enviar o cache KV do prefiller para o decodificador.

Com o prefill-decode misto, nossa implantação do R1 estava lutando para exceder consistentemente 50 TPS devido a interrupções frequentes de prefill na ordem de centenas de milissegundos. Em contrapartida, ao separar o prefill, incorremos em uma penalidade de cerca de 100ms no TTFT para cada solicitação, mas um único nó prefiller podia manter tamanhos de lote consistentes em 3 nós decodificadores, entregando uma taxa de transferência superior a 90 TPS enquanto lidava com uma carga de cerca de 1 QPS por nó decodificador. Com implantações paralelas de dados, o TPS foi ligeiramente menor, em torno de 50, no entanto, as instâncias podiam lidar com uma carga de 1 QPS por classificação, com 8 classificações para um único nó.

Qwen3-Coder

Este modelo de 480B usa Atenção de Consulta Agrupada (GQA), para que a atenção possa ser facilmente fragmentada e possa se beneficiar do paralelismo de tensores sem sacrificar a memória para os caches KV. Consequentemente, pudemos fragmentar o modelo em 8 GPUs para prefill e decode, pareando cerca de 3 nós decodificadores com um único nó prefiller. Como a atenção é fragmentada, confiamos no layout de cache KV HND para fragmentar os caches KV do prefiller e do decodificador, pareando as classificações do prefiller com as classificações do decodificador e utilizando totalmente todas as NICs para transferir fatias em paralelo.

Citação

BibTeX
@online{perplexity2025DisaggregatedPrefillAndDecode,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Kevin Hu and Erin McKnight and Phi Hoang},
  title = {{Disaggregated prefill and decode}},
  date = {2025-08-01},
  url = {https://perplexity.coolkeypoint.com/hub/blog/disaggregated-prefill-and-decode}
}