Paged Attention em Large Language Models LLMs

Ao executar LLMs em escala, a limitação real é a memória da GPU, e não a capacidade de computação, principalmente porque cada requisição exige um cache KV para armazenar dados em nível de token. Em configurações tradicionais, um grande bloco de memória fixo é reservado por requisição com base no comprimento máximo da sequência, o que leva a um desperdício significativo de espaço e limita a concorrência. O Paged Attention melhora isso dividindo o cache KV em blocos menores e flexíveis que são alocados apenas quando necessários, de forma semelhante ao funcionamento da memória virtual. Ele também permite que várias requisições com o mesmo prompt inicial compartilhem memória e a dupliquem apenas quando suas saídas começam a ser diferentes. Essa abordagem melhora muito a eficiência da memória, permitindo uma taxa de transferência significativamente maior com pouquíssima sobrecarga. Neste artigo, simulamos o alocador de cache KV ingênuo, construímos uma implementação funcional de Paged Attention com uma tabela de blocos e compartilhamento de prefixo Copy-on-Write, e medimos a lacuna de utilização em tamanhos de lote de 10 a 200 requisições simultâneas. Importando as dependências Copiar Código Copiado Usar um navegador diferente import math import random import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as mpatches from collections import defaultdict random.seed(42) np.random.seed(42) Configurando as Constantes Antes de simular qualquer coisa, precisamos saber quanto de memória da GPU um único token realmente custa. Isso depende inteiramente da arquitetura do modelo. Usamos uma configuração estilo GPT — 32 camadas, 32 heads de atenção, 128 dimensões por head, armazenados em fp16. O fator de 2 na frente corresponde às projeções Key (Chave) e Value (Valor) (não há cache Q — as queries são recalculadas a cada etapa). Multiplicando esses valores, obtemos 524.288 bytes, ou 512 KB, por token. Esta é a unidade fundamental sobre a qual tudo o mais é construído — tamanhos de pré-alocação, contagens de páginas e memória desperdiçada são todos dimensionados diretamente a partir desse número. Copiar Código Copiado Usar um navegador diferente NUM_LAYERS = 32 NUM_HEADS = 32 HEAD_DIM = 128 BYTES_FP16 = 2 PAGE_SIZE = 16 # tokens por página (padrão vLLM) MAX_SEQ_LEN = 2048 KV_BYTES_PER_TOKEN = 2 * NUM_LAYERS * NUM_HEADS * HEAD_DIM * BYTES_FP16 KV_MB_PER_TOKEN = KV_BYTES_PER_TOKEN / 1024 / 1024 Cache KV Ingênuo A abordagem ingênua é simples: quando uma requisição chega, um bloco contíguo de memória da GPU é alocado com o tamanho máximo da sequência — 2048 tokens neste caso. Isso acontece porque o comprimento da resposta é desconhecido antecipadamente, então o pior cenário é reservado. AVG_RESPONSE é definido como 500, o que é uma média realista.