Pré-preenchimento 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: pré-preenchimento e decodificação. O pré-preenchimento ocorre nos tokens de entrada, populando caches KV, antes de entrar no estágio de decodificação que gera tokens um por um.

Enquanto um único passo de decodificação normalmente demora dezenas de milissegundos, o pré-preenchimento leva substancialmente mais tempo. Se executado nos mesmos dispositivos, misturar pré-preenchimento com decodificação degrada o desempenho da decodificação. Neste artigo, exploramos uma solução estabelecida na forma de pré-preenchimento e decodificação desagregados, executando-os em dispositivos separados para maximizar tanto a taxa de transferência do pré-preenchimento quanto as latências de decodificação.

Desempenho de Pré-preenchimento vs Decodificação

Em um motor típico de serviço LLM, 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ó, tanto as solicitações de pré-preenchimento quanto de decodificação são agrupadas juntas. O custo da atenção, que se agrega ao longo do comprimento da sequência, cresce tanto para pré-preenchimento quanto para decodificação, proporcionalmente ao comprimento das entradas na 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 pré-preenchimento encaminham milhares ou dezenas de milhares de tokens a um custo significativo por meio de camadas densas (grande qo_len).

A latência de uma passagem adiante é mais fortemente influenciada pelo número de tokens independentes que passam através das camadas densas (qo_len) do que pelo número de tokens recuperados da cache KV durante a atenção (kv_len). A atenção pode ser paralelizada tanto ao longo do número de solicitações quanto do kv_len proporcional aos comprimentos das sequências, alcançando boa utilização. O pré-preenchimento é limitado por computação: um qo_len alto permite que núcleos GEMM aloque blocos suficientes ao longo da dimensão M para utilizar completamente as capacidades de computação das GPUs modernas. A decodificação é limitada pela memória: devido aos tamanhos de lotes tipicamente baixos, o número de entradas ao longo de M geralmente é pequeno, suficiente para apenas um bloco. Enquanto núcleos Split-K GEMM podem melhorar a utilização de SM para tamanhos de lotes de token baixos, as caches e as unidades de multiplicação de matrizes geralmente permanecem subutilizadas.

Quando misturados, os lotes contendo solicitações para pré-preenchimento incorrer em latências mais altas através da passagem à frente, afetando negativamente a taxa de transferência de decodificação de toda a instância. Embora misturar solicitações de pré-preenchimento com solicitações de decodificação ou empregar pré-preenchimento em fragmentos possa melhorar ligeiramente o desempenho da decodificação, é difícil manter uma taxa de transferência de pré-preenchimento suficiente para processar solicitações suficientes em uma instância para maximizar a taxa de transferência de 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 pré-preenchimento deve ser realizado com frequência suficiente para que degrade significativamente a latência média e cause trêmulos na saída.

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

Messenger KV

Na Perplexity, nossa implementação para pré-preenchimento e decodificação desagregados é construída em torno de um messenger KV que interage com o motor LLM para orquestrar transferências de cache KV de nós de pré-preenchimento para os nós de decodificação através de uma rede. No lado do pré-preenchimento, o messenger aceita solicitações de nós de decodificação, entregando-as ao agendador de lotes e acompanhando a execução da passagem adiante para despachar caches KV com o mínimo de latência possível. No lado do decodificação, após as páginas não-ejetáveis serem alocadas, o messenger bloqueia a solicitação de ser agendada para decodificação até ser notificado da conclusão das transferências de cache KV e contexto de decodificação.

Desagregar o pré-preenchimento requer conexões de alta taxa de transferência e baixa latência, assim nossa implementação é adaptada para RDMA, suportando os Controladores de Interface de Rede EFA e ConnectX (NICs). O Messenger KV é construído sobre libfabric, usando nossos wrappers fabric-lib para fornecer abstrações de baixa latência de nível superior sobre os primitivos de Acesso Remoto à Memória Direta (RDMA), implementando transferências de páginas e metadados eficientes, juntamente com sinalização de baixa latência. Em segundo plano, fabric-lib coordena uma GPU e seus NICs diretamente conectados para copiar dados do nó de pré-preenchimento para o nó de decodificação.

Após o recebimento, o nó de pré-preenchimento aloca um conjunto correspondente de páginas de origem KV e agenda a solicitação para pré-preenchimento usando seu motor local. Para minimizar a latência, as transferências não esperam pela passagem adiante: em vez disso, as cópias de páginas KV são iniciadas assim que o modelo termina de anexar entradas de cache KV à cache KV para camadas individuais. Como as solicitações de pré-preenchimento podem ser em fragmentos, o agendador de lotes notifica o messenger KV dos fragmentos atualmente agendados antes da execução. Para dar suporte a gráficos CUDA enquanto é capaz de rastrear camadas, o messenger mantém uma thread dedicada a fazer polling em 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 de shard: mesmo que 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 através das posições, sincronizando-as implicitamente. Uma vez que uma mudança no contador é observada, o messenger é notificado e chama fabric-lib para iniciar a transferência de uma camada.

Após a transferência do último fragmento ser concluída, qualquer metadado adicional também é copiado: a decodificação especulativa ou MTP requer que logits e estados ocultos sejam movidos para o decodificador. Essas cópias também são realizadas através de RDMA, de e para buffers pré-alocados.

Após a conclusão de todas as transferências pendentes do último fragmento, o nó de pré-preenchimento desaloca as páginas KV e conclui a solicitação. O nó de decodificação não é explicitamente notificado: em vez disso, usa contadores imediatos para acompanhar o número de operações concluídas. O número de operações RDMA no lado do pré-preenchimento é 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, fabric-lib chama o messenger KV para indicar que uma solicitação está pronta para decodificação. O messenger desaloca qualquer contexto e entrega a solicitação ao motor LLM.

Transferências de Cache KV Fragmentadas

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

Se a origem e o destino forem fragmentados identicamente, as transferências são triviais, pois há um mapeamento um para um dos dispositivos e páginas da origem e do destino. Nesta situação, a fragmentação ajuda implicitamente nas latências de transferência: usando mais GPUs, mais NICs associados podem ser empregados, chegando mais perto da utilização total da largura de banda. Porém, se houver uma discrepância, o motor de transferência deve dividir ou reconstruir as páginas dependendo da relação entre as fatias de origem e destino.

Se o pré-preenchimento dividir a cache KV em mais dispositivos, páginas completas são reconstruídas no decodificador enviando as metades correspondentes dos dispositivos de pré-preenchimento. Se o decodificador tiver mais fragmentos, ele recebe páginas de várias fontes. O decodificador precisa conhecer o esquema de fragmentação do pré-preenchimento para ser capaz de calcular o número de escritas RDMA que se espera receber. Se houver replicação envolvida, o pré-preenchimento agrupa os dispositivos em conjuntos de réplicas que replicam a cache KV completa dentro de si. 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 escritas RDMA.

As transferências fragmentadas requerem um pequeno ajuste nas caches KV. Por padrão, o FlashInfer depende do layout NHD, que ordena os tokens dentro de uma página dentro das cabeças. Como as caches provavelmente são fragmentadas ao longo da quantidade de cabeças de atenção, isso cria descontinuidade dentro da cabeça. As transferências RDMA não suportam implicitamente escritas com passos, exigindo uma operação por cabeça para realizar a transferência. Em vez disso, para reduzir o número de interações com libfabric, organizamos as caches KV usando o layout HND que coloca a dimensão da cabeça antes do número de tokens. Isso assegura continuidade, permitindo que uma página seja copiada com uma única escrita.

Decodificação Especulativa

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

Para contornar essas questões, o pré-preenchimento não inclui o último token da sequência de entrada. Em vez disso, estados ocultos ou logits do pré-preenchimento que precedem o último token são transferidos e tratados como um token de decodificação na próxima etapa no decodificador. Embora isso aumente ligeiramente as latências, pois um passo completo de decodificação deve ser realizado após o pré-preenchimento para emitir o primeiro token, a complexidade da implementação é significativamente reduzida.

Implantações Desagregadas

Nós implantamos ou experimentamos com múltiplas configurações desagregadas com diferentes modelos, para suportar tráfego de produção ou cargas de trabalho de avaliação interna. Com base no tamanho e mecanismo de atenção dos modelos, escolhemos esquemas de fragmentação adequados para os nós de pré-preenchimento e decodificador para melhor utilizar as GPUs.

DeepSeek-R1

Com o DeepSeek, consideramos implantações tanto Tesnor-Paralelas (TP) quanto Dados-Paralelas (DP). Como discutido em postagens de blog anteriores, as implantações TP oferecem melhor latência ao custo de menor taxa de transferência, exigindo mais GPUs para atender ao grande tráfego. Implementações DP escalam muito melhor com carga, no entanto, sua taxa de transferência máxima é menor devido ao custo de comunicação entre dispositivos ou entre nós.

DeepSeek depende de Multi-Head Latent Attention, comprimindo caches KV. Como todas as cabeças KV estão comprimidas em um único vetor latente, TP não pode fragmentar as caches KV, pois deve em vez disso, replicar os vetores latentes em todas as posições. A fragmentação ocorre após a descompressão, pois cada posição pode extrair diferentes cabeças da mesma representação latente. Consequentemente, todos os fragmentos de caches KV são idênticos entre os fragmentos de pré-preenchimento e decodificação.

Com uma configuração TP intra-no, tanto pré-preenchedores quanto decodificadores são fragmentados de forma idêntica. As transferências são despachadas de todas as posições para utilizar completamente todos os NICs disponíveis. No entanto, com uma implementação DP, onde o tamanho da posição TP é menor ou cada posição DP é atribuída a uma única GPU, qualquer dispositivo de pré-preenchimento que mantenha uma cópia replicada do cache KV pode despachá-lo. Para balancear solicitações entre todos os NICs disponíveis, selecionamos aleatoriamente uma GPU e um NIC para enviar a cache KV do pré-preenchedor para o decodificador.

Com pré-preenchimento-decodificação misturados, nossa implantação R1 estava lutando para consistentemente exceder 50 TPS devido a frequentes interrupções de pré-preenchimento na ordem de centenas de milissegundos. Em contraste, ao separar o pré-preenchimento, sofremos uma penalidade de cerca de 100ms para TTFT para cada solicitação, mas um único nó de pré-preenchimento poderia manter tamanhos de lote consistentes em 3 nós de decodificação, entregando uma taxa de transferência superior a 90 TPS enquanto lidava com uma carga de cerca de 1 QPS por nó de decodificação. Com implantações de dados-paralelos, o TPS foi ligeiramente menor, em torno de 50, no entanto, as instâncias podiam lidar com uma carga de 1 QPS por posição, com 8 posições para um único nó.

Qwen3-Coder

Este modelo de 480B usa Atenção de Consulta Agrupada (GQA), então a atenção pode ser facilmente fragmentada e pode ser beneficiada com o paralelismo de tensor sem sacrificar a memória para caches KV. Consequentemente, pudemos fragmentar o modelo em 8 GPUs tanto para pré-preenchimento quanto para decodificação, emparelhando cerca de 3 nós de decodificação com um único nó de pré-preenchimento. Como a atenção é fragmentada, contamos com o layout de cache KV HND para fragmentar caches KV de pré-preenchimento e decodificação, emparelhando posições de pré-preenchimento com posições de decodificação e utilizando plenamente todos os NICs para transferir fragmentos em paralelo.