PagedAttention

PagedAttention é um algoritmo de atenção e um gestor de memória para servir modelos de linguagem de grande porte (LLMs) que armazena a cache KV de cada pedido em pequenos blocos de tamanho fixo, colocados em qualquer local da memória da GPU, e acompanha-os através de uma tabela de blocos por pedido, à semelhança de como um sistema operativo mapeia páginas de memória virtual. Woosuk Kwon, Zhuohan Li e colegas da UC Berkeley e de universidades parceiras apresentaram-na juntamente com o motor de serving vLLM no artigo Efficient Memory Management for Large Language Model Serving with PagedAttention, apresentado no SOSP 2023. A eliminação da fragmentação causada por alocações de cache contíguas permite que uma GPU mantenha muito mais sequências em simultâneo, e o artigo refere um débito 2–4 vezes superior ao do FasterTransformer e do Orca com a mesma latência.

O problema de memória que resolve#

Durante a geração de texto, um transformador mantém os vetores de chaves e valores de cada token anterior para não ter de os recalcular. Essa cache aumenta uma entrada por token e por camada, e uma única sequência do modelo LLaMA de 13 mil milhões de parâmetros pode ocupar até 1,7 GB. O comprimento da saída é desconhecido quando chega um pedido, pelo que os primeiros sistemas de serving reservavam para cada pedido uma região contígua dimensionada para o comprimento máximo possível. A maior parte dessa região ficava vazia, e as lacunas entre regiões eram demasiado pequenas para reutilizar. A equipa vLLM mediu que estes sistemas desperdiçavam 60–80% da memória da cache KV, o que limitava o tamanho do lote muito antes de esgotarem a capacidade de cálculo.

Como funciona a paginação#

A PagedAttention divide a cache em blocos, cada um com as chaves e os valores de um número fixo de tokens — 16 por predefinição no vLLM. Um pedido recebe um novo bloco apenas quando o último fica cheio, a partir de qualquer bloco livre, e a sua tabela de blocos regista a ordem. O kernel do mecanismo de atenção segue essa tabela para reunir as chaves e os valores bloco a bloco, pelo que a cache de uma sequência não precisa de ser contígua. O desperdício limita-se à parte não preenchida do último bloco de cada sequência: menos de 4% nas medições do vLLM.

Este exemplo de código PyTorch armazena as chaves de um pedido em blocos dispersos de um conjunto partilhado e lê-as através da tabela de blocos:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Apresenta a tabela de blocos [41, 7, 19] e keys intact: True e, em seguida, 8 posições não utilizadas para o pedido paginado, em comparação com 2.008 para uma região reservada para 2.048 tokens.

Partilha de blocos entre sequências#

As tabelas de blocos também permitem que várias sequências apontem para os mesmos blocos físicos. A amostragem paralela, a pesquisa em feixe e os pedidos que começam com o mesmo prompt de sistema reutilizam uma única cópia do prefixo partilhado, com uma contagem de referências em cada bloco. Quando uma sequência precisa de escrever num bloco ainda utilizado por outras, o motor copia primeiro esse bloco, num passo de cópia na escrita emprestado dos sistemas operativos.

Os autores do vLLM referem que esta partilha reduz até 55% o uso de memória da amostragem paralela e da pesquisa em feixe, permitindo até 2,2 vezes mais débito. A aplicação de hash a blocos completos estende a ideia a vários pedidos, que é a forma como o vLLM implementa a cache de prompts.

Técnicas relacionadas#

A PagedAttention determina onde ficam as chaves e os valores em cache e combina-se com métodos que alteram a quantidade de dados a armazenar ou a velocidade de leitura:

• Flash Attention: Divide o cálculo da atenção em blocos, para que a matriz completa de pontuações nunca seja escrita na memória principal da GPU. Os motores atuais executam kernels do tipo FlashAttention que leem blocos paginados.
• Agrupamento contínuo: Aceita um novo pedido sempre que outro termina, e a paginação devolve imediatamente os blocos da sequência concluída ao conjunto livre, libertando espaço para o pedido seguinte.
• Atenção de consultas agrupadas: Partilha cabeças de chaves e valores entre cabeças de consulta, reduzindo o conteúdo de cada bloco.
• Quantização da cache FP8: Armazena chaves e valores com precisão de 8 bits, reduzindo aproximadamente para metade a memória da cache em comparação com formatos de 16 bits.

A interface LLM da Ultralytics encaminha os argumentos dos pedidos sem alterações; assim, uma chamada Chat Completions com n=3, enviada através de base_url para um servidor vLLM local, devolve três amostras que partilham os blocos em cache do prompt, e apenas as continuações divergentes ocupam memória nova.