antonio leandro

ia generativa

Efficient Memory Management for Large Language Model Serving with PagedAttention

paper · núcleo · Kwon, Woosuk, Li, Zhuohan, Zhuang, Siyuan, Sheng, Ying, Zheng, Lianmin, Yu, Cody Hao, Gonzalez, Joseph E., Zhang, Hao · · ~57 min de leitura do original

a tese

o gargalo de servir llm não é flop, é fragmentação de memória: paginar o kv cache como um sistema operacional pagina ram rende 2 a 4 vezes mais throughput sem encostar no modelo

o que fica

  1. Nos sistemas anteriores, apenas 20,4% a 38,2% da memória reservada para KV cache guardava estado de token de verdade; o resto era reserva parada, fragmentação interna e fragmentação externa.
  2. No OPT-13B, o KV cache de um único token ocupa 800 KB, e uma sequência de 2.048 tokens chega a 1,6 GB — é a memória, não o poder de cálculo, que define o tamanho do batch.
  3. O kernel de PagedAttention é 20% a 26% mais lento que o do FasterTransformer e mesmo assim vence de 2 a 4 vezes no fim a fim, porque o ganho vem de caber mais requisições no batch.
  4. Beam search economiza de 37,6% a 55,2% da memória com blocos compartilhados, enquanto parallel sampling economiza bem menos, de 6,1% a 9,8%: quanto mais as sequências divergem tarde, mais o compartilhamento paga.
  5. O bloco padrão do vLLM tem 16 tokens porque bloco pequeno desperdiça o paralelismo da GPU e bloco grande traz a fragmentação interna de volta.
  6. Recomputar o KV cache despejado nunca custou mais que 20% da latência de trazê-lo de volta da RAM da CPU, já que prompt e tokens gerados voltam num único passo de prefill.

o problema

Servir um llm é um problema de memória disfarçado de problema de compute. A geração é autorregressiva: o modelo emite um token por vez, e cada passo precisa dos vetores key e value de todos os tokens anteriores. Guardar esses vetores — o KV cache — evita recomputar o passado a cada iteração, mas cobra caro. No OPT-13B, um único token ocupa 800 KB de KV cache: 2 vetores × 5.120 de hidden state × 40 camadas × 2 bytes em FP16. Uma sequência de 2.048 tokens chega a 1,6 GB. Numa A100 de 40 GB, os pesos ficam com cerca de 65% da memória e sobram uns 30% para o cache de todas as requisições ao mesmo tempo. Quem decide o throughput é o tamanho do batch, e quem decide o tamanho do batch é essa fatia de memória.

Os sistemas da época — FasterTransformer, Orca — guardavam o KV cache de cada requisição num tensor contíguo, porque é assim que os frameworks de deep learning querem tensores. Como o comprimento da saída é desconhecido de antemão, eles pré-alocavam o pior caso: 2.048 posições, mesmo quando a resposta teria 25 tokens. Daí saem três desperdícios: reserva ocupada durante toda a vida da requisição, fragmentação interna do que nunca será preenchido e fragmentação externa, já que cada bloco pré-alocado tem tamanho diferente. Na medição dos autores, só 20,4% a 38,2% da memória reservada guardava estado real. Pior: com cada requisição presa na sua área contígua, duas amostras do mesmo prompt não conseguiam dividir nem um byte do prefixo idêntico.

a ideia

O problema já tinha nome em 1962, em outro contexto: fragmentação de memória. A resposta do sistema operacional foi memória virtual com paginação — quebrar a memória em páginas de tamanho fixo e deixar que páginas lógicas contíguas apontem para páginas físicas espalhadas. O trabalho leva essa solução para dentro do kernel de attention. Blocos são páginas, tokens são bytes, requisições são processos. A analogia se sustenta até o fim: tem tabela de páginas, tem contagem de referências, tem copy-on-write, tem swap para a “disk” (aqui, a RAM da CPU).

como funciona

O KV cache de uma sequência é cortado em KV blocks de tamanho fixo — 16 tokens por padrão. Um bloco físico é uma fatia de um pedaço contíguo de DRAM da GPU; um bloco lógico é o que a requisição enxerga, preenchido da esquerda para a direita. Uma block table por requisição mapeia lógico para físico e registra quantas posições estão preenchidas. Novo bloco físico só é alocado quando o anterior enche, então o desperdício de uma requisição fica limitado a um bloco.

A conta de attention é reescrita bloco a bloco: o kernel pega a query do token atual, multiplica pelas keys de um bloco, obtém aquele pedaço do vetor de scores e acumula a saída contra os values do mesmo bloco. Nada exige que os blocos sejam vizinhos na memória física. O prefill continua usando attention convencional; a PagedAttention entra na fase de decodificação.

O compartilhamento vem de graça a partir daí. Em parallel sampling, os blocos do prompt são mapeados nas duas sequências com refcount 2; quando uma delas escreve no último bloco compartilhado, o sistema copia aquele bloco e só ele. Em beam search, candidatos compartilham blocos intermediários e o padrão muda a cada passo, como uma árvore de processos. Em prefixo compartilhado — system prompt, exemplos few-shot —, o provedor pré-computa os blocos e cada requisição só mapeia. Tudo isso se expressa com três operações: fork, append, free.

Quando a memória acaba, o escalonador é FCFS e o despejo é tudo-ou-nada por sequência: ou saem todos os blocos daquela sequência, ou nenhum. Para trazer de volta, duas opções — swap para a RAM da CPU ou recomputar. No modo distribuído, existe um único KV cache manager central; os workers recebem a block table junto com os tokens de entrada a cada iteração e nunca precisam sincronizar decisão de memória entre si.

o que isso custou

Indireção não é grátis. Consultar a block table, tratar branches extras e comprimentos variáveis deixa o kernel de attention 20% a 26% mais lento que a implementação do FasterTransformer. O argumento dos autores é que isso só afeta o operador de attention, e que o batch maior paga a diferença com folga — o que os números fim a fim sustentam, mas é um trade-off explícito, não um almoço grátis.

O ganho também não é universal. Em OPT-175B com sequências curtas do Alpaca, sobra tanta memória para KV cache que o sistema vira compute-bound e a vantagem sobre o Orca encolhe. Os próprios autores dizem que a técnica não deve ser aplicada a qualquer carga de GPU: em treino, os shapes são estáticos e a alocação pode ser resolvida antes; em inferência de redes que não são llm, o gargalo costuma ser cálculo, e aí a indireção só atrapalha.

A política de preempção é grosseira de propósito. Depois de despejar uma sequência, o sistema para de aceitar requisições novas até que as preemptadas terminem. E o preço de implementação é real: kernels CUDA próprios para escrita em bloco, leitura fundida com attention e cópia de blocos em lote, sem os quais a ideia perde o pé.

onde isso aparece hoje

O vocabulário deste paper virou o vocabulário do serving: bloco de KV cache, prefix caching, copy-on-write entre sequências. O código saiu público junto com o artigo, e o vLLM deixou de ser um protótipo de paper para virar um servidor de inferência de uso corrente.

O caso do prefixo compartilhado é o que mais migrou para produto: reaproveitar o estado de um system prompt longo entre requisições é exatamente o que aparece hoje como Prompt caching nas apis. O prefill continua sendo território de FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, que ataca o mesmo gargalo de memória por outro lado — movimentação entre níveis de memória, não fragmentação de alocação. E há um terceiro caminho, complementar: encolher o próprio cache na origem, como faz GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints ao compartilhar keys e values entre attention heads. Os três convivem, porque atacam camadas diferentes do mesmo fato: desde o Attention Is All You Need, o custo de servir um transformer é o custo de carregar estado.

lido na íntegra por pipeline de llm, revisado por antonio leandro antes de publicar ·