Prefill e decodificação desagregados

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

Desempenho de Prefill vs. Decodificação

Em um mecanismo de serviço de LLM típico, o agendador de lote 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 de decodificação são processadas em lote juntas. O custo da atenção, que se agrega ao longo do comprimento da sequência, cresce tanto para o prefill quanto para a decodificação, proporcionalmente ao comprimento das entradas no cache KV (kv_len). As solicitações de decodificação geralmente encaminham um único token (qo_len=1), a um custo mínimo através 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 através de camadas densas (grande qo_len).

A latência de uma passagem direta é influenciada mais fortemente pelo número de tokens independentes passados através de 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 paralelizar tanto em relação ao número de solicitações quanto ao kv_len proporcional aos comprimentos das sequências, alcançando boa utilização. O prefill é limitado pela 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. A decodificação é limitada pela memória: devido aos tamanhos de lote geralmente 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 de SM para tamanhos de lote de token baixos, os caches e as unidades de multiplicação de matrizes geralmente permanecem subutilizados.

Expandir

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

Expandir

Expandir

Essas questões podem ser abordadas usando um conjunto separado de nós para realizar o prefill e a decodificação. Ao associar um nó de prefill a múltiplos nós de decodificação, solicitações suficientes podem ser agendadas para prefill para maximizar o rendimento e manter um número grande o suficiente de solicitações simultâneas nos nós de decodificação para também maximizar o rendimento da decodificação. Os nós de prefill populam os caches KV, que são então transferidos para os nós de decodificação. Como os decodificadores não precisam mais fazer pausas para o prefill, as latências tornam-se muito mais determinísticas, uma vez que o impacto geral do crescimento do kv_len de solicitações ativas é muito menos pronunciado. O custo é pago em um aumento no Tempo para o Primeiro Token (TTFT), já que a transferência de caches KV pela rede pode levar de dezenas a centenas de milissegundos.

KV Messenger

Na Perplexity, nossa implementação para prefill e decodificação desagregados é construída em torno de um KV messenger que interage com o mecanismo de LLM para orquestrar as transferências de cache KV dos nós de prefill para os nós de decodificação por meio de uma rede. No lado do prefill, o messenger aceita solicitações dos nós de decodificação, encaminhando-as para o agendador de lote e acompanhando a execução da passagem direta para despachar os caches KV com a menor latência possível. No lado do decodificador, após a alocação de páginas não evacuáveis, o messenger bloqueia a solicitação de ser agendada para decodificação até que seja notificado da conclusão das transferências de cache KV e contexto do decodificador.

Expandir

A desagregação do prefill requer conexões de alto rendimento 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 baixo nível de latência sobre os primitivos de Acesso Direto à Memória Remota (RDMA), implementando transferências eficientes de página e metadados, juntamente com sinalização de baixa latência. Em segundo plano, o fabric-lib coordena uma GPU e seus NICs conectados diretamente para copiar dados do nó de prefill para o nó do decodificador.

Após o recebimento, o nó de prefill aloca um conjunto correspondente de páginas KV de origem e agenda a solicitação para prefill usando seu mecanismo local. Para minimizar a latência, as transferências não aguardam a passagem direta: em vez disso, as cópias de página KV são iniciadas assim que o modelo termina de anexar entradas de cache KV ao cache KV para camadas individuais. Como as solicitações de prefill podem ser fragmentadas, o agendador de lote notifica o KV messenger sobre os fragmentos agendados atualmente antes da execução. Para suportar gráficos CUDA enquanto consegue rastrear camadas, o messenger mantém um tópico dedicado sondando um contador incrementado após a projeção de saída da atenção. O contador é mantido apenas no nó líder em um ambiente fragmentado: embora as entradas de cache KV sejam válidas após o anexo e antes da atenção, a projeção de saída é reduzida entre as classificações, sincronizando-as implicitamente. Uma vez observada uma mudança no contador, o messenger é notificado e chama o fabric-lib para iniciar a transferência de uma camada.

Expandir

Após a conclusão da transferência do último fragmento, 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 através de RDMA, para e de buffers pré-alocados.

Após a conclusão de todas as transferências pendentes do último fragmento, o nó de prefill desaloca as páginas KV e conclui a solicitação. O nó do decodificador não é explicitamente notificado: 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 prefill é proporcional ao número de páginas transferidas. Após a conclusão do número conhecido de cópias de página e contexto, o fabric-lib chama o KV messenger para indicar que uma solicitação está pronta para decodificação. O messenger desaloca qualquer contexto e entrega a solicitação ao mecanismo de LLM.

Transferências de Cache KV Fragmentado

Se o prefill e o decodificador dependem de Paralelismo de Tensor (TP) e fragmentam ou replicam os caches KV identicamente, um único mecanismo de transferência coordena múltiplos 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 suporta os caches KV e mapeá-la para o processo principal. O mecanismo de transferência inspeciona a topologia do nó para encontrar os 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 fragmentarem identicamente, as transferências são triviais, pois há um mapeamento um-para-um dos dispositivos e páginas da origem e do destino. Nessa situação, a fragmentação ajuda implicitamente nas latências de transferência: usando mais GPUs, mais NICs associados podem ser empregados, 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 as páginas dependendo da proporção entre as fatias de origem e destino.

Expandir

Se o prefill dividir o cache KV em mais dispositivos, páginas inteiras são reconstruídas no decodificador enviando as metades correspondentes para fora dos dispositivos de prefill. Se o decodificador tiver mais fragmentos, ele recebe páginas de múltiplas fontes. O decodificador precisa conhecer o esquema de fragmentação do prefill para poder calcular o número de gravações RDMA que se espera que receba. Se a replicação estiver envolvida, o prefill agrupa os dispositivos em conjuntos de réplicas que replicam o cache KV completo dentro deles mesmos. Os conjuntos de réplicas de destino são atribuídos aleatoriamente a um dos conjuntos de origem para usar todos os dispositivos disponíveis para iniciar as gravações RDMA.

Expandir

As transferências fragmentadas exigem um leve 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çalhos. Como os caches são mais provavelmente fragmentados ao longo do número de cabeçalhos de atenção, isso cria descontinuidade dentro do cabeçalho. As transferências RDMA não suportam implicitamente gravações com saltos (strided writes), exigindo uma operação por cabeçalho para realizar a transferência. Em vez disso, para reduzir o número de interações com o libfabric, organizamos caches KV usando o layout HND, que coloca a dimensão do cabeçalho antes do número de tokens. Isso garante continuidade, permitindo que uma página seja copiada com uma única gravação.

Decodificação Especulativa

A decodificação especulativa requer ajustes leves no prefill-decodificação desagregado. Em nossa implementação, os nós de prefill 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 alvo.

Expandir

Para contornar esses problemas, o prefill não inclui o último token da sequência de entrada. Em vez disso, estados ocultos ou logits do prefill que precedem o último token são transferidos e ele é tratado como um token de decodificação na próxima etapa no decodificador. Embora isso aumente ligeiramente as latências, já que uma etapa de decodificação 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 múltiplas configurações desagregadas com diferentes modelos, para suportar o tráfego de produção ou 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 de prefill e decodificação para melhor utilizar as GPUs.

DeepSeek-R1

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

O DeepSeek depende da Atenção Latente de Cabeçalho Múltiplo (Multi-Head Latent Attention), comprimindo caches KV. Como todos os cabeçalhos KV são comprimidos em um único vetor latente, o TP não pode fragmentar os caches KV, pois deve, em vez disso, replicar os vetores latentes em todas as classificações. A fragmentação ocorre após a descompressão, pois cada classificação pode extrair diferentes cabeçalhos da mesma representação latente. Consequentemente, todos os fragmentos de cache KV são idênticos em ambos os fragmentos de prefill e decodificador.

Com uma configuração de TP intra-nó, tanto os prefillers quanto os decodificadores são fragmentados identicamente. As transferências são despachadas de todas as classificações para utilizar totalmente todos os 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 de prefill que mantenha uma cópia replicada do cache KV pode despachá-lo. Para equilibrar as solicitações em todos os NICs disponíveis, selecionamos aleatoriamente uma GPU e um NIC para enviar o cache KV do prefill para o decodificador.

Com prefill-decodificação misto, nossa implantação R1 estava lutando para exceder consistentemente 50 TPS devido a interrupções frequentes de prefill na ordem de centenas de milissegundos. Em contraste, ao separar o prefill, incorremos em uma penalidade de cerca de 100ms para o TTFT para cada solicitação, mas um único nó de prefill poderia manter tamanhos de lote consistentes em 3 nós de decodificador, entregando um rendimento superior a 90 TPS enquanto lidava com uma carga de cerca de 1 QPS por nó de decodificador. Com implantações de paralelismo 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 por nó.

Qwen3-Coder

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

Citação

@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}
}