Quanta memória exige o KV cache de um pedido a um LLM?
Tradução automática
Este artigo foi traduzido automaticamente a partir da versão original em inglês.
Num transformer convencional com atenção completa, a memória do KV cache aumenta com o número de tokens em cache, camadas, cabeças de chave/valor e bytes por valor. Multiplique a alocação por sequência pelo número de sequências simultâneas para estimar o volume bruto de dados. Depois reserve memória adicional para os pesos e as operações do runtime.
Esta estimativa ajuda os engenheiros de serving a dimensionar uma carga de trabalho. A atenção latente multicabeça (MLA), as caches de janela deslizante, híbridas e recorrentes precisam de estimativas específicas para cada arquitetura.
Use as cabeças KV em vez das cabeças de consulta
A cache guarda as projeções de chave e valor dos tokens anteriores para que o modelo não as volte a calcular em cada passo de geração. A atenção densa continua a ler um histórico de tokens que cresce. Para caches de atenção completa com cabeças de chave/valor separadas, como em MHA, MQA ou GQA, e sequências com o mesmo comprimento:
cache bytes = 2 × layers × KV_heads × head_dimension
× cached_tokens × concurrent_sequences × bytes_per_element
O fator dois conta as chaves e os valores. Para sequências de comprimentos diferentes, use a soma das suas contagens de tokens em cache. Leia um head_dim explícito da configuração do modelo quando estiver disponível; hidden_size / num_attention_heads é apenas uma alternativa para arquiteturas que usam essa relação.
A tabela de arquitetura do Llama 3 da Meta fornece as dimensões destes exemplos, que usam uma cache FP16/BF16 de dois bytes por valor:
| Modelo | Camadas / cabeças KV / dimensão da cabeça | Tokens por sequência | Cache bruta por sequência |
|---|---|---|---|
| Llama 3.1 8B | 32 / 8 / 128 | 8,192 | 1 GiB |
| Llama 3.1 8B | 32 / 8 / 128 | 131,072 | 16 GiB |
| Llama 3.1 70B | 80 / 8 / 128 | 4,096 | 1.25 GiB |
| Llama 3.1 70B | 80 / 8 / 128 | 131,072 | 40 GiB |
Um orçamento de cache de 40 GiB permite guardar, segundo este cálculo, no máximo 32 sequências completas de 4,096 tokens do modelo 70B. Isto conta apenas os dados da cache; escolha o limite de concorrência depois de reservar memória para o runtime.
Verifique a alocação real do servidor
Os pesos, os buffers temporários, as ativações, os blocos parcialmente preenchidos e a capacidade reservada exigem mais memória. Os servidores distribuídos podem repartir ou replicar as cabeças KV. Uma cache de um byte suportada reduz este volume bruto de dados para metade, mas os fatores de escala, o suporte da implementação e a qualidade podem alterar a capacidade utilizável.
Meça o pico de alocação e a preempção com comprimentos representativos de prompt e de saída. PagedAttention melhora a alocação de blocos; não elimina os dados de chave/valor guardados nem torna o comprimento da sequência irrelevante.
Guia de engenharia: memória do KV cache inclui um cálculo de capacidade executável.