Quanta memória reservar para o cache KV?

Descrever o que permanece na memória

Durante uma geração autorregressiva, as chaves e valores dos tokens já processados podem ser conservados para as etapas seguintes. Esse cache pertence às camadas de atenção. Seu volume depende, portanto, do modelo e do histórico conservado, não apenas do número de parâmetros. O contexto de uma conversa também inclui as instruções, as mensagens anteriores e os documentos adicionados pela aplicação.

Sua primeira decisão é operacional: quantas sequências precisarão permanecer ativas, até qual comprimento? Uma fila de vinte requisições das quais duas são executadas simultaneamente não representa necessariamente vinte caches residentes. Registre a admissão real das requisições e as eventuais sequências adicionais criadas pela geração.

Prepare uma ficha com a revisão do modelo, as camadas de atenção, as cabeças KV, a dimensão das chaves e valores, seu dtype e a estratégia de cache. Conte os tokens após o tokenizer e o template de conversa. Um limite em caracteres não descreve essa alocação.

Fontes técnicas: Hugging Face — funcionamento e forma dos caches por camada

Usar as cabeças KV, sobretudo com GQA

O número de cabeças de consulta Q e o das cabeças KV podem diferir. Na atenção multi-cabeça clássica, eles coincidem. Com MQA, uma única cabeça KV é compartilhada; GQA agrupa várias cabeças Q em torno de uma cabeça KV. Leia num_key_value_heads na configuração quando esse campo existe e verifique seu significado para a arquitetura.

Por exemplo, quarenta cabeças Q e oito cabeças KV formam cinco cabeças Q por grupo KV. A fórmula do cache armazenado usa oito, não quarenta. Essa proporção não descreve toda a memória da atenção: operações ou conversões podem produzir temporários.

Não modifique simplesmente esse número para reduzir a necessidade de memória de um modelo já treinado. O esquema de atenção faz parte da arquitetura dele. Dois modelos com números de cabeças diferentes não se tornam variantes equivalentes por um cálculo de capacidade; a qualidade deles deve ser avaliada separadamente.

Fontes técnicas: Hugging Face — campos do LlamaConfig e distinção entre MHA, MQA e GQA · PyTorch — dimensões Q/K/V e restrições da atenção GQA

Estabelecer a fórmula e suas unidades

No caso uniforme, L é o número de camadas, Hkv o número de cabeças KV, D a dimensão delas, T o número de posições conservadas por sequência, B o número de sequências e q o número de bytes por valor. O fator dois conta K e V. Essa aproximação pressupõe chaves e valores com a mesma dimensão e o mesmo formato, sem compressão nem compartilhamento de prefixo.

Converta apenas o resultado final: um GiB equivale a 1.073.741.824 bytes; um GB decimal equivale a 1.000.000.000 bytes. Mantenha os bytes na sua ficha para evitar que um arredondamento ou uma mudança de unidade mascare uma diferença.

Para comprimentos diferentes sem padding, substitua B × T pela soma das posições realmente armazenadas. Para camadas heterogêneas, faça uma soma camada por camada. Um armazenamento denso com padding, uma alocação por blocos ou uma reserva estática exigem contar os espaços alocados, que podem exceder os tokens úteis.

Fontes técnicas: Hugging Face — dimensões dos tensores do cache · NIST — unidades decimais e prefixos binários

Exemplo resolvido: seis sequências, sem medição de GPU

Tomemos uma arquitetura fictícia de quarenta camadas, oito cabeças KV e uma dimensão de 128. Suponhamos um cache uniforme com dois bytes por valor. Cada sequência recebe no máximo 3.072 tokens de entrada e uma reserva para 1.024 tokens adicionais, ou seja, um limite de 4.096 posições. São hipóteses pedagógicas, não a configuração atestada de um modelo.

O custo calculado por posição e por sequência é 2 × 40 × 8 × 128 × 2 = 163.840 bytes. Uma sequência de 4.096 posições representa então 671.088.640 bytes, ou seja, 0,625 GiB. Seis sequências dão 4.026.531.840 bytes, ou seja, 3,75 GiB só para o cache.

Dobrar o comprimento conservado dobra esse item nessa fórmula. Substituir oito cabeças KV por quarenta o multiplica por cinco, mantendo inalteradas todas as outras hipóteses. Essas proporções não anunciam nenhuma aceleração, perda de qualidade ou compatibilidade de um modelo real.

Adaptar o orçamento à estratégia de cache

Um cache dinâmico cresce com as posições conservadas. Um cache estático reserva uma capacidade máxima: dimensione essa reserva, não apenas a requisição curta observada na inicialização. Para uma atenção com janela deslizante, algumas camadas podem limitar seu histórico; as camadas com atenção completa exigem um cálculo à parte.

A quantização do cache e sua transferência para a CPU são outras estratégias, dependentes do modelo e do software. Elas modificam as restrições de armazenamento, de transferência ou de cálculo. Quantizar os pesos não prova que o cache tem o mesmo formato.

Anote a classe de cache e seus parâmetros explícitos. Verifique também a liberação dos espaços após o fim ou o cancelamento de uma requisição. Para uma carga que mistura sequências curtas e longas, uma reserva máxima uniforme pode pesar mais do que a mera soma dos conteúdos úteis.

Fontes técnicas: Hugging Face — caches dinâmicos, estáticos, quantizados e transferidos

Verificar uma carga representativa por etapas

Construa três casos: entrada habitual, entrada longa esperada e número máximo de requisições simultâneas permitido. Fixe modelo, tokenizer, gabarito, limite de geração e regra de fim. Varie uma dimensão de cada vez; uma resposta encurtada ou um documento truncado altera o trabalho realizado.

Meça separadamente o carregamento, o processamento inicial da entrada e a geração. Sincronize o device em torno das fases cronometradas, conserve os níveis de memória iniciais e os picos. O método de memória explica por que allocated e reserved não se somam e por que a diferença entre seus máximos não isola o cache.

Seu controle deve resultar em um limite testado e em saídas aceitáveis: IDs das requisições concluídas, erros, comprimento produzido e critério de qualidade. Uma execução que se mantém com uma entrada curta não valida a concorrência máxima. Um erro de memória antes do fim não constitui uma medida da necessidade de uma execução completa.

Fontes técnicas: PyTorch — sincronização do trabalho no device escolhido · PyTorch — escopo dos contadores de memória

Descartar os erros que distorcem a escolha

Não transforme o cache calculado em capacidade total da GPU. Acrescente uma análise dos pesos, das ativações temporárias, das saídas mantidas e do software. Também não divida esse volume automaticamente pelo número de placas: o posicionamento das camadas ou das cabeças deve ser realmente configurado e verificado em cada device.

O notebook IteraGPU Lab é um exercício de instrumentação em uma pequena rede densa sem atenção. Ele ajuda a ler fases de memória; seu parâmetro context não valida esse cálculo de KV. Use a ficha de carga e a pasta de inferência para preparar o teste do seu próprio modelo.

Compare as capacidades das ofertas somente depois de distinguir o volume aritmético, o máximo realmente observado e os casos ainda não testados. O plano de locação organiza sua janela de trabalho; ele não garante nenhum comprimento de contexto nem cadência de geração.

• Confundir cabeças Q e cabeças KV: retomar a configuração da arquitetura.
• Orçar apenas o prompt: incluir a geração e a reserva efetiva.
• Ler "pesos 4 bits" como "cache 4 bits": registrar os dois formatos.
• Esquecer a concorrência: contar as sequências realmente residentes.
• Prometer memória compartilhada entre placas: verificar o posicionamento real.

Perguntas práticas