Chunked Prefill

O pré-preenchimento em blocos é uma técnica de serviço de modelos que divide uma mensagem de entrada longa em grupos mais pequenos de tokens e os processa ao longo de várias etapas, intercalando-os com a geração de tokens de outros pedidos. Sem esta técnica, um motor de inferência que aceita uma mensagem longa passa uma etapa inteira a processá-la, e todas as respostas que já estão a ser transmitidas a outros utilizadores ficam em pausa até terminar. O artigo SARATHI, de 2023, introduziu o pré-preenchimento em blocos, e o trabalho seguinte, Sarathi-Serve, criou com base nele um escalonamento «sem pausas». O Dynamic SplitFuse do DeepSpeed-FastGen usa a mesma ideia, e o pré-preenchimento em blocos está agora ativado por predefinição no vLLM e disponível no SGLang e no TensorRT-LLM, motores que servem muitos modelos de linguagem de grande escala (LLMs) com pesos abertos.

Pré-preenchimento, descodificação e pausas na geração#

A inferência de Transformer autorregressivos tem duas fases:

• Pré-preenchimento: O modelo processa em paralelo todos os tokens da mensagem, calcula as respetivas chaves e valores e armazena-os na cache KV. O pré-preenchimento executa muitas operações por byte de pesos lido, pelo que costuma ser limitado pelo processamento.
• Descodificação: O modelo gera um token por etapa para cada pedido, lendo toda a cache em cada etapa. A descodificação realiza poucas operações aritméticas por byte lido, pelo que costuma ser limitada pela largura de banda da memória.

Uma etapa da GPU que inclui um pré-preenchimento de 20 000 tokens demora muito mais do que uma etapa de descodificação, pelo que todos os pedidos do lote ficam à espera. O Sarathi-Serve chama a estas pausas «pausas de geração», que se manifestam como picos na latência entre tokens para utilizadores cujas próprias mensagens são curtas.

Como funciona o pré-preenchimento segmentado#

O escalonador atribui a cada etapa um orçamento de tokens, definido no vLLM por max_num_batched_tokens. Primeiro, acrescenta um token para cada descodificação em curso e, em seguida, preenche o orçamento restante com blocos das mensagens em espera. Uma mensagem que não caiba prossegue na etapa seguinte, expandindo a sua cache KV bloco a bloco até que toda a mensagem seja processada e comece a descodificação. A combinação das duas fases numa só etapa também permite utilizar melhor a GPU, pois as descodificações deixam capacidade de processamento livre que os blocos de pré-preenchimento podem ocupar:

O processamento em blocos altera apenas o momento em que os tokens da mensagem são processados. Cada bloco presta atenção às chaves e aos valores em cache de todos os blocos anteriores, bem como aos seus próprios tokens, aplicando uma máscara causal. Assim, o resultado é igual ao de um pré-preenchimento feito numa só passagem, como mostra uma camada de atenção em PyTorch:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
prompt_len, chunk = 1000, 256
q, k, v = torch.randn(3, 1, 8, prompt_len, 64).unbind(0)  # one layer: 8 heads, 64 dims each

# One-shot prefill: all prompt tokens attend causally in a single pass
full = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Chunked prefill: each chunk attends to the cached keys and values of earlier chunks plus itself
outputs = []
for start in range(0, prompt_len, chunk):
    end = min(start + chunk, prompt_len)
    mask = torch.ones(end - start, end, dtype=torch.bool).tril(diagonal=start)  # causal over absolute positions
    q_chunk, k_cached, v_cached = q[..., start:end, :], k[..., :end, :], v[..., :end, :]
    outputs.append(F.scaled_dot_product_attention(q_chunk, k_cached, v_cached, attn_mask=mask))
chunked = torch.cat(outputs, dim=2)

print(f"{len(outputs)} chunks, same result as one pass: {torch.allclose(full, chunked, atol=1e-5)}")

Imprime 4 chunks, same result as one pass: True: a mensagem de 1 000 tokens é processada em blocos de 256, 256, 256 e 232 tokens, e os resultados correspondem aos da passagem única.

Técnicas de serviço relacionadas#

• Agrupamento contínuo: Adiciona e remove pedidos em cada etapa, em vez de esperar que um lote inteiro termine. O pré-preenchimento em blocos permite escalonar mensagens longas dentro desse ciclo.
• PagedAttention: Armazena a cache KV em páginas de tamanho fixo, permitindo que uma mensagem pré-preenchida ao longo de várias etapas aumente a cache sem reservar antecipadamente memória para todo o seu comprimento.
• Armazenamento em cache de mensagens: Reutiliza as chaves e os valores em cache de um prefixo de mensagem repetido, eliminando por completo essa parte do pré-preenchimento.
• Separação entre pré-preenchimento e descodificação: Executa as duas fases em conjuntos distintos de GPU e transfere a cache KV entre eles, permitindo que cada conjunto seja dimensionado e ajustado para o seu próprio estrangulamento.
• Descodificação especulativa: Acelera a fase de descodificação ao propor vários tokens e verificá-los numa só passagem.

Aplicações#

O pré-preenchimento em blocos é particularmente importante quando um servidor recebe mensagens longas e curtas. Um assistente de geração aumentada por recuperação (RAG) que preenche a sua janela de contexto com um relatório de 50 páginas pode pré-preenchê-la em blocos enquanto as conversas de outros utilizadores continuam a ser transmitidas. Os modelos de visão e linguagem (VLMs) acrescentam centenas ou milhares de tokens de imagem a cada mensagem, pelo que os pipelines de câmaras que enviam fotogramas para descrição criam exatamente essa combinação. A interface LLM da Ultralytics envia texto e imagens para qualquer ponto final compatível com OpenAI, incluindo um servidor vLLM alojado localmente, onde os pedidos de um pipeline YOLO26 para descrever os fotogramas em que deteta pessoas são pré-preenchidos em blocos juntamente com o tráfego dos outros utilizadores.

Ajustar o orçamento de tokens#

O orçamento envolve uma troca entre duas métricas de latência. Um orçamento pequeno deixa pouco espaço para o pré-preenchimento em cada etapa, o que mantém baixa a latência entre tokens para os utilizadores de transmissão contínua, mas aumenta o tempo até ao primeiro token dos novos pedidos; um orçamento grande produz o efeito inverso. Cada bloco posterior de uma mensagem longa também volta a ler a cache KV crescente dos blocos anteriores, pelo que blocos muito pequenos aumentam o tráfego de memória. Os operadores ajustam o orçamento com base na mediana e nos valores de cauda da latência de inferência, no tempo até ao primeiro token e no débito, usando comprimentos de mensagem realistas. O guia de otimização do vLLM indica 2 048 tokens como um orçamento que favorece a latência entre tokens e recomenda mais de 8 192 tokens para maximizar o débito com modelos mais pequenos em GPU de grande capacidade.