vLLM / PagedAttention (SOSP 2023)

01 · O PROBLEMA

O cache KV domina a memória de servir LLMs

Servir um modelo de linguagem grande é caro: segundo estimativas recentes citadas no paper, processar um pedido de LLM pode custar dez vezes mais que uma consulta tradicional por palavras-chave. Por isso, aumentar o throughput, e com ele reduzir o custo por pedido, vem se tornando mais importante.

É possível elevar o throughput agrupando vários pedidos num lote. Só que o número de pedidos que cabem juntos no lote continua limitado pela memória da GPU, sobretudo pelo espaço do cache KV, e por isso o throughput do sistema é limitado por memória.

No exemplo do paper, um LLM de 13B parâmetros numa GPU NVIDIA A100 com 40GB reserva cerca de 65% da memória para os pesos, que ficam estáticos durante o serviço, e perto de 30% para os estados dinâmicos dos pedidos. Para um Transformer, esses estados são as chaves e os valores da atenção, o cache KV.

O cache KV por token é grande. No OPT de 13B ele ocupa 800 KB por token, e como o modelo gera sequências de até 2048 tokens, um único pedido pode precisar de até 1,6 GB. Mesmo dedicando toda a memória disponível ao cache KV, caberiam só algumas dezenas de pedidos.

Duas características tornam esse cache difícil de gerenciar: ele cresce e encolhe dinamicamente conforme o modelo gera tokens, e seu tempo de vida e seu tamanho não são conhecidos de antemão.

02 · COMO ERA ANTES

Reserva contígua e agendamento por iteração

Os sistemas anteriores guardam o cache KV de cada pedido num tensor contíguo, porque a maioria dos frameworks de deep learning exige tensores em memória contígua. Como a saída é imprevisível, eles alocam estaticamente um pedaço com base no comprimento máximo possível da sequência, seja qual for o tamanho real da entrada ou da saída.

Isso gera três fontes de desperdício: slots reservados para tokens futuros, fragmentação interna por superdimensionar para o comprimento máximo e fragmentação externa vinda do alocador de memória, como o buddy allocator. O pedaço inteiro fica reservado durante a vida do pedido, e pedidos mais curtos não aproveitam a parte sem uso.

Nas medições do paper, só entre 20,4% e 38,2% da memória do cache KV guardava estados reais de tokens nos sistemas existentes.

Há também técnicas de agendamento de granularidade fina, como o agendamento em nível de iteração, operam por iteração: depois de cada uma, os pedidos concluídos saem do lote e novos entram. Assim, um pedido novo espera só uma iteração, e não o lote inteiro. O paper também descreve que o vLLM seleciona as sequências candidatas ao lote a cada iteração de decodificação.

1. O pedido chega pelo frontend ao escalonador. Com agendamento em nível de iteração, pedidos concluídos saem do lote e novos entram a cada iteração.

2. Para cada pedido, o sistema reserva na memória da GPU um pedaço contíguo dimensionado pelo comprimento máximo possível da sequência, seja qual for o tamanho real.

3. Enquanto o pedido vive, o pedaço inteiro continua reservado e pedidos mais curtos não aproveitam a parte que ainda não foi usada.

03 · PAGEDATTENTION

Memória virtual para o cache KV

O PagedAttention divide o cache KV de cada pedido em blocos, e cada bloco guarda as chaves e os valores de um número fixo de tokens. Esses blocos podem ficar em memória física não contígua. A analogia do paper é a memória virtual dos sistemas operacionais: blocos fazem o papel de páginas, tokens o de bytes e pedidos o de processos.

O gerenciador de cache KV mantém tabelas de blocos, que ligam os blocos lógicos de cada pedido aos blocos físicos. Os blocos lógicos são preenchidos da esquerda para a direita, e separar lógico de físico permite crescer o cache sob demanda, sem reservar todas as posições de antemão. Segundo o paper, isso elimina a maior parte do desperdício dos sistemas existentes.

O vLLM é um sistema de serving completo, com um frontend FastAPI e um motor de inferência em GPU. Seu escalonador centralizado coordena a execução dos trabalhadores de GPU distribuídos, e o gerenciador de cache KV cuida da memória física pelas instruções que recebe do escalonador.

1. No prefill, o gerenciador reserva só os blocos necessários ao prompt, e não o comprimento máximo. No exemplo do paper, um prompt de sete tokens usa dois blocos lógicos, mapeados a dois blocos físicos.

2. No primeiro passo de decodificação, o cache do novo token cabe no último bloco lógico, que ainda tinha uma posição livre.

3. No passo seguinte o último bloco está cheio: o gerenciador aloca um novo bloco físico e registra o mapeamento na tabela de blocos.

4. A cada passo, o escalonador prepara uma mensagem com os IDs dos tokens e a tabela de blocos de cada pedido e a transmite aos trabalhadores, que leem o cache KV conforme essa tabela.

Como um novo bloco físico só é alocado quando todos os anteriores estão cheios, o vLLM limita o desperdício de cada pedido a um bloco. Quando um pedido termina, seus blocos são liberados para outros pedidos. O paper acrescenta que blocos relativamente pequenos alocados sob demanda aliviam a fragmentação interna, e que blocos do mesmo tamanho eliminam a externa.

O tamanho do bloco é uma escolha de projeto. Se for pequeno demais, o vLLM pode não aproveitar todo o paralelismo da GPU ao ler e processar o cache KV; se for grande demais, a fragmentação interna aumenta e a chance de compartilhamento diminui. O paper adota 16 como padrão, por achar esse valor grande o bastante para usar bem a GPU e pequeno o bastante para evitar fragmentação interna significativa na maioria das cargas.

04 · COMPARTILHAR BLOCOS

Contagem de referência e copy-on-write

Na amostragem paralela, um pedido gera várias saídas a partir do mesmo prompt, então o cache KV do prompt pode ser compartilhado. No exemplo do paper com duas saídas, os blocos lógicos do prompt de ambas são mapeados aos mesmos blocos físicos.

Como um bloco físico pode ser mapeado a vários blocos lógicos, o vLLM mantém uma contagem de referência por bloco físico. Para os blocos que várias sequências precisam modificar, usa copy-on-write em granularidade de bloco, semelhante à técnica de copy-on-write da memória virtual dos sistemas operacionais.

No beam search o compartilhamento vai além do prompt: candidatos compartilham também outros blocos, e o padrão muda conforme a decodificação avança, como numa árvore de processos criada por forks. Sistemas anteriores exigem cópias frequentes do cache KV entre candidatos; no vLLM, a maioria dos blocos pode ser compartilhada, e o copy-on-write envolve copiar apenas um bloco.

Para prefixos compartilhados, como um prompt de sistema longo com instruções e exemplos, o provedor pode guardar antecipadamente o cache KV do prefixo. No vLLM, isso se faz reservando blocos físicos para um conjunto de prefixos predefinidos pelo provedor; o prompt do usuário mapeia seus blocos lógicos para esses blocos em cache, com o último marcado como copy-on-write, e a fase de prompt só roda sobre a entrada da tarefa.

1. O provedor reserva blocos físicos para um conjunto de prefixos predefinidos, e o cache KV do prefixo fica guardado de antemão.

2. Um pedido com o prefixo mapeia seus blocos lógicos para os blocos físicos em cache, e a fase de prompt só precisa executar sobre a entrada da tarefa.

3. Como um bloco físico passa a servir a mais de um bloco lógico, o gerenciador mantém uma contagem de referência por bloco físico.

4. Quando uma sequência precisa escrever num bloco compartilhado, o gerenciador reconhece que a contagem é maior que um, aloca um novo bloco físico, manda copiar o conteúdo e reduz a contagem do bloco original.

Uma camada de mapeamento comum, dos blocos lógicos para os físicos, esconde esses padrões de compartilhamento: o modelo e seu kernel veem apenas uma lista de IDs de blocos físicos por sequência. Por isso o vLLM processa ao mesmo tempo pedidos com métodos de decodificação diferentes, o que sistemas existentes não fazem com eficiência, segundo o paper.

05 · ESCALONAMENTO E ESCALA

Preempção, troca para a CPU e várias GPUs

Quando o tráfego passa da capacidade do sistema, o vLLM precisa priorizar pedidos e adota a política FCFS (primeiro a chegar, primeiro a ser servido) para todos os pedidos, o que garante justiça e evita starvation. Ao preemptar, os pedidos mais antigos são servidos primeiro e os mais recentes saem primeiro. À medida que pedidos e saídas crescem, o vLLM pode ficar sem blocos físicos na GPU para o novo cache KV.

Quando os blocos acabam, é preciso decidir quais blocos expulsar e como recuperá-los. Como todos os blocos de uma sequência são acessados juntos, o vLLM usa expulsão tudo-ou-nada: ou saem todos os blocos de uma sequência, ou nenhum. Sequências de um mesmo pedido, como os candidatos de um beam search, são agendadas em grupo e preemptadas ou reagendadas juntas.

O paper afirma que o desempenho de cada técnica depende da banda entre a RAM da CPU e a memória da GPU e do poder de computação da GPU. No swapping, o espaço de troca na RAM da CPU fica limitado pela memória da GPU alocada ao cache KV.

Para modelos grandes demais para uma GPU, o vLLM suporta o paralelismo de tensor no estilo Megatron-LM. Há um único gerenciador de cache KV dentro do escalonador centralizado, e os trabalhadores compartilham esse gerenciador e o mapeamento de blocos lógicos para físicos. Cada trabalhador tem os mesmos IDs de bloco físico, mas guarda só a parte do cache KV das suas cabeças de atenção. Assim, os trabalhadores não precisam sincronizar o gerenciamento de memória: recebem essa informação no início de cada iteração de decodificação, junto com as entradas do passo.

06 · RESULTADOS E LIMITES

O que o paper mediu, e o que não mede

O paper afirma que o vLLM melhora o throughput de LLMs populares em 2 a 4 vezes, com o mesmo nível de latência, em comparação com sistemas como FasterTransformer e Orca. A melhora é maior com sequências mais longas, modelos maiores e algoritmos de decodificação mais complexos.

No conjunto ShareGPT, o vLLM sustentou taxas de pedidos de 1,7× a 2,7× maiores que o Orca (Oracle) e de 2,7× a 8× maiores que o Orca (Max), com latências semelhantes. No OPT-13B, processou 2,2× mais pedidos ao mesmo tempo que o Orca (Oracle) e 4,3× mais que o Orca (Max).

Orca (Oracle) supõe conhecer o comprimento real das saídas, um limite superior inviável na prática; Orca (Max) sempre reserva o espaço até o comprimento máximo de sequência do modelo.

Há limites a considerar. Primeiro, o Orca não estava disponível publicamente, então os autores implementaram a própria versão, e as variantes diferem em quanto reservam a mais para a saída. Segundo, no OPT de 175B com o conjunto Alpaca, de sequências curtas, a vantagem sobre Orca (Oracle) e Orca (Pow2) foi menos pronunciada: sobrou memória para o cache KV, e o desempenho passou a ser limitado por computação, e não por memória. Terceiro, no teste de chatbot o cache KV não é guardado entre rodadas de conversa.