Quanta memória o modelo vai consumir: como GQA, janelas e MLA mudam o KV-cache

Uma situação conhecida: você colocou um modelo pronto para rodar, os pesos couberam com folga na memória da GPU, mas o servidor já cai com out-of-memory na décima requisição simultânea. O culpado é o KV-cache, não o tamanho do modelo. Seu volume depende da arquitetura da camada de atenção, e essa escolha já não está nas suas mãos: ela foi feita durante o treinamento.

O checkpoint define a arquitetura

A configuração e os pesos determinam o número de cabeças de consulta e de cabeças key-value, a dimensão de cada cabeça, o padrão de atenção local e as dimensões da representação latente. Trocar MHA por GQA depois do treinamento significa alterar o modelo e, muito provavelmente, será preciso fazer treinamento adicional para recuperar a qualidade. Trocar atenção completa por uma janela deslizante significa descartar parte do contexto que o modelo se acostumou a enxergar.

Na prática, portanto, abra a configuração (no Hugging Face, com AutoConfig) e confira o número de camadas, num_attention_heads, num_key_value_heads, head_dim e os campos de sliding window. Uma ressalva: o campo da janela pode existir na configuração, mas estar desativado por uma flag ao lado. Verifique tanto o valor quanto a flag. Alguns modelos alternam camadas locais e globais, e o código da implementação é sempre a referência definitiva.

Como fazer uma conta rápida do cache

A memória persistente permanece alocada entre os passos de decodificação e determina quantas requisições você consegue atender ao mesmo tempo. Buffers temporários, como os blocos de scores de atenção, podem ser liberados após a camada. Para um cache padrão, a fórmula é: 2 × bytes por elemento × camadas × tokens × número de cabeças KV × dimensão da cabeça. O fator 2 corresponde às chaves mais os valores. Para um batch, multiplique pelo tamanho dele.

Considere um modelo hipotético: 32 camadas, 32 cabeças, head_dim 128, fp16 e contexto de 8192 tokens. Com MHA clássico, uma requisição ocupa 4 GiB. Se o modelo usa GQA com 8 cabeças KV, são 1 GiB. Com MQA e uma única cabeça KV, apenas 128 MiB. A diferença entre “cabem 4 requisições” e “cabem 16” surge sem alterar o número de parâmetros.

Em arquiteturas mistas, fazemos a conta camada por camada. Suponha que, nesse mesmo modelo GQA de 32 camadas, 24 sejam locais, com janela de 4096, e 8 usem atenção completa, com contexto de 32 768 tokens. Usar atenção completa em todas as camadas consumiria 4 GiB; a arquitetura mista, cerca de 1,4 GiB. Mas isso não impõe um limite ao cache total: cada novo token ainda acrescenta memória nas camadas com atenção completa. Ele cresce mais devagar, mas não para de crescer.

As variantes de atenção

• MHA: cada cabeça de consulta tem seu próprio par de chave e valor. É o cache mais caro e também o que exige mais leitura de memória na decodificação.
• MQA e GQA: as cabeças de consulta compartilham cabeças KV. O cache e o número de parâmetros das projeções K/V diminuem exatamente na mesma proporção que o número de cabeças KV.
• Janela deslizante: a camada mantém apenas os últimos W tokens, e o cache pode ficar em um buffer circular de tamanho fixo. As posições são calculadas pelo índice lógico do token, não pela posição no buffer.
• MLA: armazena um vetor latente comprimido em vez das chaves e dos valores completos. O cache é menor, mas a decodificação exige mais cálculos. Isso costuma compensar porque o gargalo da decodificação é a largura de banda da memória, não a aritmética.

Não é possível encaixar MLA na fórmula de MHA por meio de num_key_value_heads: os tensores armazenados em cache são outros. Calcule com base no que a implementação realmente armazena e use kernels especializados.

O que o runtime pode mudar

Vale separar quatro questões: quais tokens atendem a quais outros (atenção completa ou em janela), quantas representações KV foram aprendidas (MHA/GQA/MQA), o que fica em cache (tensores KV ou representação latente) e como isso é executado. O checkpoint determina as três primeiras; o motor de inferência, a quarta. FlashAttention calcula a atenção exata em blocos e não grava matrizes enormes de scores na memória: os buffers temporários e o tráfego diminuem, mas o cache lógico e o número de cabeças KV permanecem iguais.

Paged attention armazena o cache em blocos que não precisam ser contíguos e melhora a alocação e o compartilhamento, mas, por si só, não reduz o tamanho do cache. O alocador precisa respeitar a regra de descarte do modelo: não pode descartar entradas antigas de uma camada com atenção completa só porque as camadas vizinhas usam janelas.

Conclusão prática

Antes de escolher um checkpoint, leia a configuração e estime o cache para o contexto e o batch que pretende usar. Se dois candidatos tiverem qualidade próxima, escolha aquele cujas cabeças KV ou janelas consumam menos memória: isso se traduz diretamente no número de requisições simultâneas por placa. E não se esqueça do padding: em batches densos, você pode consumir memória com posições vazias; uma implementação ragged evita isso.

A limitação do método é simples: o cálculo estima a memória persistente, mas não inclui pesos, ativações nem o overhead do alocador. Os números acima são ilustrativos; para um modelo real, use os valores da configuração dele.