Conceitos
KV cache paginado e batching contínuo para LLMs
Intuição
Pense num restaurante que, para cada grupo de clientes, reserva de antemão uma mesa enorme, grande o bastante para o pior caso de quantos pratos o grupo pode pedir. Como ninguém sabe quantos pratos virão, a maior parte da mesa fica vazia e o restaurante recusa novos grupos mesmo estando quase vazio. A solução é parecida com a de um sistema operacional: em vez de uma mesa gigante por grupo, distribuir pequenas bandejas de tamanho igual, só quando o grupo realmente precisa de mais espaço. Assim cabem muito mais grupos ao mesmo tempo.
Como funciona
Um LLM autorregressivo gera um token por vez e guarda, para cada token já visto, os vetores de chave e valor da atenção (o KV cache) para não recalculá-los; na GPU, os pesos ficam fixos, então o espaço restante para o KV cache decide quantas requisições cabem no batch. Os sistemas anteriores citados no paper (FasterTransformer, Orca) guardam o KV cache de cada requisição em um bloco contíguo e reservam o comprimento máximo possível, o que gera reserva ociosa, fragmentação interna e fragmentação externa; nas medições dos autores com suas reimplementações do Orca (Max, Pow2 e Oracle), só 20,4% a 38,2% da memória do KV cache continha estados de tokens de fato. O PagedAttention, inspirado na memória virtual paginada, divide o KV cache de cada sequência em blocos de tamanho fixo que não precisam ser contíguos; uma block table por requisição mapeia blocos lógicos em blocos físicos, e um bloco novo só é alocado quando os anteriores estão cheios, de modo que o desperdício fica restrito ao último bloco. Como o mapeamento passa por uma tabela, vários blocos lógicos podem apontar para o mesmo bloco físico, com contador de referências: o KV cache de um mesmo prompt é guardado uma só vez em amostragem paralela, e, quando uma sequência precisa escrever em um bloco compartilhado, o vLLM faz copy-on-write no nível do bloco (aloca outro, copia e reduz o contador do original). Somado ao batching em nível de iteração (técnica anterior, como no Orca), em que requisições concluídas saem e novas entram após cada iteração, o vLLM, quando faltam blocos físicos, usa FCFS e preempta as requisições mais recentes primeiro, de forma tudo ou nada, e recupera os blocos por swapping para a CPU ou por recomputação do KV cache.
Na prática
No SOSP 2023, o paper reporta que o vLLM, com 96,3% da memória do KV cache contendo estados de tokens, 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) no ShareGPT com amostragem básica, com latências semelhantes (o Orca é uma reimplementação dos autores). Os ganhos crescem quando há compartilhamento: no beam search a economia de memória no Alpaca foi de 37,6% a 55,2%, e com prefixo compartilhado de 5 exemplos o throughput foi 3,58x o do Orca (Oracle). O custo medido é que os kernels de atenção ficaram 20% a 26% mais lentos que os do FasterTransformer, e o desempenho ponta a ponta ainda foi superior.
Use em um desenho
Pratique na escola de System Design os casos que dependem deste conceito.