Papers
Efficient Memory Management for Large Language Model Serving with PagedAttention
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica · SOSP 2023, 2023
O problema: o KV cache limita o tamanho do batch
Servir um modelo de linguagem grande (LLM) com alto throughput exige processar muitas requisições juntas em um batch. O gargalo, segundo o paper, é a memória da GPU. Um LLM autorregressivo gera um token por vez, e para cada token já visto guarda os vetores de chave e valor da atenção, o chamado KV cache, para não recalculá-los a cada passo. No exemplo do paper, uma GPU NVIDIA A100 de 40 GB servindo um modelo de 13 bilhões de parâmetros usa cerca de 65% da memória para os pesos, que ficam fixos, e mais de 30% para o KV cache; o restante é pouco e serve para ativações temporárias. Como os pesos são constantes, é a gestão do KV cache que decide quantas requisições cabem no batch. O tamanho é grande: no OPT-13B, um único token ocupa cerca de 800 KB de KV cache, então uma sequência de até 2048 tokens chega a cerca de 1,6 GB. Além disso, o KV cache cresce e encolhe dinamicamente, e o comprimento final da saída não é conhecido de antemão. Os autores argumentam que o ganho de computação das GPUs cresce mais rápido que o de memória (da A100 para a H100 os FLOPS sobem mais de 2x, mas a memória máxima continua em 80 GB), então a memória tende a ser um gargalo cada vez mais relevante.
Desperdício por fragmentação nos sistemas existentes
Os sistemas anteriores citados pelos autores, como FasterTransformer e Orca, guardam o KV cache de cada requisição em um bloco contíguo de memória, porque os frameworks de deep learning da época esperam tensores contíguos. Como o tamanho final da saída é imprevisível, eles reservam antecipadamente espaço para o comprimento máximo possível da requisição (por exemplo, 2048 tokens). O paper identifica três fontes de desperdício. A primeira é a reserva: posições destinadas a tokens futuros que ainda não foram gerados, mas que ficam ocupadas e impedem outras requisições de usar esse espaço. A segunda é a fragmentação interna: o espaço superdimensionado que sobra quando a saída real é bem menor que o máximo. A terceira é a fragmentação externa, causada pelo alocador de memória (o paper cita o alocador buddy) quando cada requisição reserva um tamanho diferente. Na medição dos autores com variantes do Orca (Max, Pow2 e Oracle), apenas 20,4% a 38,2% da memória do KV cache continha estados de tokens de fato; no vLLM, o sistema proposto, esse valor foi de 96,3%. Os autores também notam que compactar a memória para reduzir a fragmentação seria impraticável, dado o volume de dados do KV cache.
A ideia central: PagedAttention e blocos de KV cache
O paper propõe o PagedAttention, um algoritmo de atenção inspirado na memória virtual com paginação dos sistemas operacionais. O KV cache de cada sequência é dividido em blocos de tamanho fixo, cada um guardando chaves e valores de um número fixo de tokens; a analogia é que blocos são páginas, tokens são bytes e requisições são processos. Os blocos físicos não precisam ser contíguos na memória da GPU: o kernel de atenção busca cada bloco separadamente. O vLLM, motor de serviço construído sobre essa ideia, mantém para cada requisição uma block table que mapeia blocos lógicos (preenchidos da esquerda para a direita conforme os tokens surgem) em blocos físicos, registrando também quantas posições de cada bloco já estão preenchidas. Um bloco físico novo só é alocado quando os anteriores estão cheios, sob demanda, sem reservar o máximo de antemão. Assim, o desperdício de uma requisição fica restrito ao último bloco, e a fragmentação externa deixa de existir porque todos os blocos têm o mesmo tamanho. O tamanho do bloco é um trade-off: blocos muito pequenos aproveitam mal o paralelismo da GPU, e blocos grandes aumentam a fragmentação interna e reduzem a chance de compartilhamento. O vLLM usa 16 por padrão, escolha que os autores justificam com a avaliação de diferentes tamanhos.
Compartilhamento, batching contínuo e preempção
Como o mapeamento passa por uma tabela, vários blocos lógicos podem apontar para o mesmo bloco físico, controlado por um contador de referências. Em amostragem paralela, em que uma requisição gera várias saídas a partir do mesmo prompt, o KV cache do prompt é guardado uma só vez. Quando uma sequência precisa escrever em um bloco compartilhado, o vLLM aplica copy-on-write no nível do bloco: aloca um novo bloco, copia o conteúdo e reduz o contador do original. No beam search o compartilhamento é maior e muda ao longo da decodificação; no prefixo compartilhado (por exemplo, uma instrução comum a várias requisições), o provedor pode manter blocos do prefixo já calculados. Esse contexto se soma ao batching em nível de iteração, descrito no paper como técnica anterior (como no Orca): ao contrário dos métodos tradicionais, que trabalham em nível de requisição, ele opera em nível de iteração; após cada iteração, requisições concluídas saem do batch e novas entram, e uma requisição nova espera só uma iteração, não o batch inteiro terminar. Quando faltam blocos físicos na GPU, o vLLM usa política FCFS e preempta as requisições mais recentes primeiro, com política de tudo ou nada: remove todos os blocos de uma sequência, e sequências de um mesmo grupo são preemptadas juntas. Para recuperar os blocos, há duas opções: swapping para a memória da CPU ou recomputação do KV cache.
O que a avaliação mede e os resultados reportados
Os autores avaliam o vLLM com modelos OPT de 13B, 66B e 175B parâmetros e LLaMA de 13B, em GPUs A100 na Google Cloud, usando cargas sintetizadas a partir dos datasets ShareGPT e Alpaca, com chegadas modeladas por distribuição de Poisson. A métrica principal é a latência normalizada, a média, entre as requisições, da latência ponta a ponta de cada uma dividida pelo seu comprimento de saída (como no Orca), medida em função da taxa de requisições: um sistema de alto throughput mantém essa latência baixa mesmo com taxas altas. Os baselines são o FasterTransformer e o Orca, este último reimplementado pelos autores em três versões (Oracle, Pow2 e Max) conforme quanto espaço reservam. Com amostragem básica no ShareGPT, o vLLM sustentou taxas de requisições 1,7x a 2,7x maiores que o Orca (Oracle) e 2,7x a 8x maiores que o Orca (Max), com latências semelhantes; frente ao FasterTransformer, até 22x. No resumo, o paper fala em 2 a 4x de ganho de throughput com a mesma latência. Os ganhos crescem com mais compartilhamento: no beam search, a economia de memória no Alpaca foi de 37,6% a 55,2% (conforme a largura testada), e com prefixo compartilhado de 5 exemplos o throughput foi 3,58x o do Orca (Oracle). O custo: os kernels de atenção ficaram 20% a 26% mais lentos que os do FasterTransformer; os autores consideram o custo pequeno porque ele afeta só o operador de atenção, e não os demais, como as camadas lineares, e o desempenho ponta a ponta do vLLM ainda foi superior.
Veja as ideias em ação
Pratique na escola de System Design os desenhos que usam este paper.