PagedAttention aplica memória virtual ao cache KV de LLMs

O problema do cache KV

O cache KV armazena as chaves e os valores de atenção de cada requisição para evitar que o modelo os recalcule a cada etapa de geração. Como ele cresce linearmente com o tamanho da sequência, pode consumir mais memória da GPU do que os próprios pesos do modelo, especialmente em contextos longos. A abordagem ingênua reserva antecipadamente o limite máximo de contexto para cada requisição, mesmo quando a maior parte desse espaço não será usada.

Essa estratégia causa fragmentação interna: uma requisição curta mantém milhares de posições vazias que não podem ser aproveitadas por outras. Também produz fragmentação externa, porque as requisições entram e terminam em momentos diferentes, deixando espaços livres descontínuos. O texto cita medições do trabalho original do vLLM segundo as quais a utilização do cache KV em cargas de produção ficava em torno de 20% a 40%, o que representava 60% a 80% de memória desperdiçada.

A analogia com memória virtual

A proposta do PagedAttention é aplicar ao cache KV a mesma ideia da memória virtual dos sistemas operacionais. Em vez de manter uma área contígua por requisição, a GPU usa um conjunto compartilhado de blocos físicos de tamanho fixo, geralmente com 16 tokens. Cada requisição mantém uma sequência lógica de blocos e uma block table, que associa cada bloco lógico ao bloco físico correspondente no pool.

Assim, os blocos podem ficar espalhados na memória física, mas continuam ordenados do ponto de vista da requisição. Eles são alocados sob demanda e devolvidos ao pool quando deixam de ser necessários. A memória usada passa a ser o tamanho efetivo da sequência, arredondado para o limite do bloco, em vez do max_context inteiro. No exemplo apresentado, três requisições de 25, 47 e 12 tokens usam 96 posições em blocos paginados, contra 6.144 posições reservadas no modelo contíguo. O texto observa que, nesse caso ilustrativo, a diferença é de 64 vezes.

Quando a atenção é executada, o núcleo percorre a tabela e busca chaves e valores nos blocos físicos espalhados. Como os blocos permanecem contíguos internamente e os núcleos de atenção já processam dados em partes, a mudança exige apenas indireção adicional e não uma reformulação completa da operação. Há um custo de acesso, descrito como situado na faixa de poucos pontos percentuais, mas a utilização do cache teria chegado a aproximadamente 96% em cargas semelhantes às de produção.

Compartilhamento e compromissos

O mecanismo também permite copy-on-write. Requisições que compartilham um prefixo, como um mesmo prompt de sistema, podem apontar para os mesmos blocos físicos. Quando uma delas precisa escrever em um bloco compartilhado, somente esse bloco é clonado; os demais continuam compartilhados. O texto afirma que, em produção, prompts de sistema reutilizados tendem a gerar economia maior que casos como busca em feixe e amostragem paralela. O SGLang amplia essa ideia com uma árvore de prefixos para localizar blocos compartilháveis.

O tamanho do bloco envolve um compromisso: blocos menores reduzem o preenchimento desperdiçado, mas aumentam consultas à tabela e acessos mais dispersos; blocos maiores reduzem essa sobrecarga, porém elevam a fragmentação interna. O texto aponta 16 tokens como um padrão adotado por vLLM, TGI e TensorRT-LLM, sem tratá-lo como valor universal.

PagedAttention não altera os cálculos nem os resultados do modelo; modifica somente a disposição e a recuperação do cache KV. Ao deixar de reservar o pior caso para cada requisição, a técnica, segundo o texto e os números do trabalho original, pode permitir atender de duas a quatro vezes mais usuários por GPU.