Cache KV

O que é um cache KV?

Em um modelo de linguagem Transformer autorregressivo (um LLM), cada etapa de decodificação prevê o próximo token dado todos os tokens anteriores. As camadas de autoatenção (self-attention) do modelo calculam três projeções para cada token:

  • Q (consultas)
  • K (chaves)
  • V (valores)

Durante a geração, as consultas do token atual “atendem” às chaves/valores de todos os tokens anteriores. Um cache KV armazena as chaves e os valores calculados para tokens anteriores para que etapas futuras de decodificação possam reutilizá-los em vez de recalculá-los repetidamente.

Este é um dos mecanismos de engenharia mais importantes no serving de LLMs porque afeta fortemente:

  • Latência (tempo para gerar cada novo token)
  • Vazão (tokens/segundo ou requisições/segundo que um servidor consegue produzir)
  • Memória (especialmente VRAM de GPU), o que por sua vez limita a concorrência

O cache KV é usado em praticamente todas as pilhas de inferência de LLM em produção (por exemplo, vLLM, TGI, TensorRT-LLM) e normalmente vem habilitado por padrão.

Contexto relacionado: o cache é uma consequência direta do cálculo de atenção descrito em Arquitetura Transformer.

Prefill vs decode: onde o cache KV fica

A inferência de LLM costuma ser dividida em duas fases:

Prefill (também conhecido como processamento do prompt)

Dado um prompt de entrada de comprimento (P), o modelo processa os tokens do prompt em paralelo (com mascaramento causal). Durante essa fase, ele calcula K/V para cada camada e para cada token do prompt e grava isso no cache KV.

  • O custo é dominado por processar um prompt longo
  • Essa fase é altamente paralela em GPUs, mas ainda assim cara para contextos longos

Decode (também conhecido como geração token a token)

Em seguida, o modelo gera tokens um por vez. No passo (t), ele:

  1. calcula o estado oculto do novo token,
  2. calcula K/V apenas para o novo token (por camada),
  3. anexa isso ao cache,
  4. executa a atenção usando K/V em cache para todos os tokens anteriores.
  • O custo é dominado por atenção repetida contra um cache crescente
  • Essa fase é menos paralela e mais sensível à latência

O cache KV acelera principalmente a fase de decodificação.

Por que o cache KV reduz computação (e latência)

Sem um cache KV, cada passo de decodificação teria que recalcular K/V para todos os tokens anteriores, mesmo que esses tokens nunca mudem. Isso é trabalho desperdiçado.

Uma visão simplificada para uma única camada:

  • No passo de decodificação (t), há (t) tokens de contexto (prompt + gerados até então).
  • Calcular K/V para todos os tokens é trabalho (O(t)).
  • Fazer isso em cada passo de 1 até (T) torna o trabalho total (O(T^2)) ao longo da geração.

Com um cache KV:

  • No passo (t), você calcula K/V apenas para o novo token: projeção incremental de K/V (O(1)) por passo.
  • A atenção ainda precisa ler o K/V em cache para (t) tokens para computar a saída, então a própria atenção cresce com o comprimento do contexto, mas você evita a projeção repetida de K/V sobre todo o prefixo.

Na prática, isso produz uma grande redução na latência por token durante a decodificação, especialmente para contextos mais longos.

Um modelo mental prático

  • Sem cache: “reler e recodificar todo o histórico toda vez”
  • Cache KV: “lembrar o histórico codificado e codificar apenas o novo token”

Memória do cache KV: o que é armazenado e qual é o tamanho?

O principal trade-off é que o cache KV consome memória significativa — frequentemente o custo dominante de VRAM durante serving em escala.

O que é armazenado?

Para cada camada Transformer, o cache armazena:

  • Chaves: forma aproximadamente ([B, H_{kv}, L, D])
  • Valores: forma aproximadamente ([B, H_{kv}, L, D])

Onde:

  • (B) = número de sequências ativas (tamanho do lote / requisições concorrentes)
  • (H_{kv}) = número de cabeças de chave/valor (pode diferir das cabeças de consulta sob MQA/GQA)
  • (L) = comprimento atual da sequência (prompt + gerados até então)
  • (D) = dimensão da cabeça (por exemplo, 128)

Isso é armazenado para todas as camadas.

Fórmula aproximada de memória

Uma estimativa comum “de guardanapo”:

[ \text{bytes KV} \approx B \cdot L \cdot N_{layers} \cdot 2 \cdot H_{kv} \cdot D \cdot \text{bytes_per_element} ]

  • O fator 2 é para K e V.
  • bytes_per_element é tipicamente 2 para FP16/BF16, 1 para INT8, etc.

Exemplo resolvido (modelo típico da classe 7B)

Assuma (semelhante a formas no estilo LLaMA-7B):

  • (N_{layers} = 32)
  • (H_{kv} = 32) (atenção multi-cabeças padrão, não MQA/GQA)
  • (D = 128)
  • dtype = FP16 (2 bytes)

Por token, por camada:

  • elementos K+V = (2 \cdot 32 \cdot 128 = 8192) elementos
  • bytes = (8192 \cdot 2 \approx 16{,}384) bytes ≈ 16 KB

Ao longo de 32 camadas:

  • (16 \text{ KB} \cdot 32 \approx 512 \text{ KB por token})

Para uma única sequência de comprimento 2048:

  • (512 \text{ KB} \cdot 2048 \approx 1{,}048{,}576 \text{ KB} \approx 1 \text{ GB})

Assim, uma única conversa longa com 2k tokens pode consumir na ordem de ~1 GB de VRAM apenas para cache KV (pesos do modelo e ativações são adicionais).

É por isso que serving com contexto longo frequentemente é limitado por memória, e por que técnicas como GQA/MQA, quantização de KV e paginação são tão importantes.

Calculadora rápida (Python)

def kv_cache_bytes(
    batch_size: int,
    seq_len: int,
    n_layers: int,
    n_kv_heads: int,
    head_dim: int,
    bytes_per_element: int = 2,  # fp16/bf16
) -> int:
    # 2 for K and V
    return batch_size * seq_len * n_layers * 2 * n_kv_heads * head_dim * bytes_per_element

def fmt_gb(n_bytes: int) -> str:
    return f"{n_bytes / (1024**3):.2f} GiB"

print(fmt_gb(kv_cache_bytes(
    batch_size=1, seq_len=2048, n_layers=32, n_kv_heads=32, head_dim=128, bytes_per_element=2
)))

Como o cache KV impacta a vazão (e por que a memória vira o gargalo)

No serving de LLMs, normalmente você quer maximizar a vazão executando múltiplas requisições concorrentemente e/ou usando batching. O cache KV complica isso porque:

  • Cada requisição ativa tem seu próprio cache crescente
  • Prompts mais longos e gerações mais longas aumentam o tamanho do cache
  • Mais requisições concorrentes multiplicam o uso de memória

Como resultado, o sistema frequentemente fica limitado por memória de KV antes de ficar limitado por computação.

Latência vs vazão: o trade-off central

  • Mais concorrência (lote maior, mais sequências ativas) melhora a utilização da GPU e a vazão.
  • Mas mais concorrência consome mais memória de cache KV, o que pode forçar:
    • tamanho máximo de lote menor,
    • contexto máximo menor,
    • OOMs mais frequentes,
    • ou técnicas caras como paginação/offloading.

Isso interage diretamente com políticas de batching discutidas em Batching e Agendamento de Inferência.

Se você usar beam search com (k) beams, você efetivamente replica o estado de decodificação, incluindo o cache KV, para cada beam. A memória pode crescer aproximadamente proporcional ao número de beams (dependente da implementação, mas frequentemente próximo disso).

Por exemplo, beam size 4 pode tornar o cache KV ~4× maior do que a decodificação gulosa.

Padrões de implementação: como o cache KV é representado em código

“past_key_values” em APIs de modelos

Muitas bibliotecas expõem o cache KV explicitamente. No Hugging Face Transformers, o cache frequentemente aparece como past_key_values, e use_cache=True diz ao modelo para retornar caches atualizados.

Conceitualmente:

# Pseudocode illustrating the idea
past = None
for step in range(num_new_tokens):
    outputs = model(input_ids=current_token_ids, past_key_values=past, use_cache=True)
    logits = outputs.logits
    past = outputs.past_key_values  # contains K/V for every layer
    current_token_ids = sample_next_token(logits)

No serving real, normalmente você não chama o modelo em um loop Python (por desempenho), mas o mecanismo subjacente é semelhante.

Gerenciamento de KV no lado do servidor

Mecanismos de serving em produção tratam o cache KV como um sistema de memória de primeira classe:

  • alocar blocos/pedaços para cada requisição,
  • anexar tokens conforme eles são gerados,
  • liberar blocos quando uma requisição termina,
  • às vezes compactar ou paginar blocos para reduzir fragmentação.

É aqui que projetos como “atenção paginada” (block-based KV caching) importam: eles reduzem fragmentação e permitem melhor empacotamento de muitas sequências com comprimentos diferentes.

Técnicas para reduzir a memória do cache KV (e melhorar a vazão)

Como o cache KV frequentemente domina o uso de VRAM, muitas otimizações o visam diretamente.

1) Atenção Multi-Query (MQA) e Atenção Grouped-Query (GQA)

Na atenção multi-cabeças padrão, cada cabeça de consulta tem suas próprias cabeças de K/V. MQA/GQA reduz o número de cabeças de K/V:

  • Consultas: muitas cabeças
  • Chaves/Valores: menos cabeças compartilhadas entre cabeças de consulta

Isso reduz a memória do cache KV aproximadamente proporcional à redução em (H_{kv}). Por exemplo, passar de 32 cabeças KV para 8 cabeças KV pode cortar o tamanho do cache KV em ~4×.

Trade-off: pode afetar levemente a qualidade do modelo, a menos que o modelo tenha sido treinado com essa variante de atenção.

2) Quantização do cache KV (INT8/FP8 e além)

Armazenar K/V em FP16 é caro. Muitas pilhas suportam quantizar o cache para reduzir memória e largura de banda:

  • KV em INT8: ~2× redução vs FP16
  • KV em FP8 (onde suportado): objetivos semelhantes, às vezes com melhores trade-offs de acurácia/vazão

Trade-off: a quantização pode degradar levemente a qualidade ou estabilidade da geração dependendo do método e do modelo.

3) Janela deslizante / atenção local (truncar histórico em cache)

Alguns modelos ou modos de serving usam uma janela de atenção fixa, mantendo apenas os (W) tokens mais recentes no cache. A memória vira (O(W)) em vez de (O(L)).

Trade-off: o modelo não consegue mais atender a tokens mais antigos do que a janela, o que pode prejudicar tarefas de contexto longo, a menos que o modelo seja projetado para isso.

4) Paginação e offloading de KV

Quando a memória de GPU acaba, você pode:

  • paginar blocos de KV entre memória da GPU e memória da CPU, ou
  • manter parte do KV na CPU e buscar conforme necessário.

Isso pode aumentar a concorrência, mas frequentemente aumenta a latência por causa de transferências PCIe/NVLink e de um agendamento mais complexo. Normalmente é usado quando você precisa servir contextos longos ou muitas sessões concorrentes.

5) Prefill em chunks / chunking de prompt

Prompts longos podem causar grandes picos temporários de ativação e alta latência de prefill. Alguns mecanismos “quebram” o prefill em chunks para:

  • reduzir memória de pico,
  • intercalar trabalho de prefill entre requisições,
  • melhorar a latência de cauda.

Isso não reduz o tamanho final de KV (você ainda precisa de K/V para todo o prompt), mas pode tornar prompts longos mais manejáveis operacionalmente.

6) Cache de prefixo (cache compartilhado de prompt)

Em muitas aplicações, uma grande parte do prompt se repete:

  • prompts de sistema,
  • instruções de ferramentas,
  • templates longos,
  • contexto recuperado que se repete entre usuários (às vezes).

Se múltiplas requisições compartilham um prefixo idêntico, você pode armazenar o KV para esse prefixo uma vez e reutilizá-lo. Isso pode reduzir dramaticamente tanto a computação de prefill quanto o tempo até o primeiro token.

Trade-offs e ressalvas:

  • Requer identidade exata do prefixo (ou gerenciamento cuidadoso)
  • Deve garantir isolamento correto (sem vazamento entre usuários)
  • Funciona melhor para prompts de sistema e templates estáticos

Detalhes de latência: tempo até o primeiro token vs tempo por token de saída

O cache KV impacta diferentes métricas de latência de maneiras distintas:

  • TTFT (time to first token) é dominado pelo prefill para prompts longos.

    • O cache KV está sendo construído aqui, não reutilizado muito.
    • Otimizações: kernels de prefill mais rápidos, chunking, cache de prefixo.

  • TPOT (time per output token) é dominado pelo decode.

    • O cache KV é fortemente reutilizado aqui.
    • Otimizações: atenção eficiente sobre K/V em cache, KV quantizado, batching melhor.

Um padrão comum em produção:

  • otimizar TTFT para UX de chat interativo,
  • otimizar TPOT e vazão para cargas de trabalho de geração em lote.

Detalhes de vazão: batching, “batching contínuo” e crescimento do cache

Diferentemente da inferência clássica de deep learning em que entradas têm tamanho fixo, o serving de LLMs tem sequências de comprimento variável que crescem ao longo do tempo. Isso torna o batching mais difícil:

  • requisições chegam em momentos diferentes,
  • prompts têm comprimentos diferentes,
  • cada requisição gera um número diferente de tokens.

Mecanismos modernos usam batching contínuo (continuous batching, também chamado de “in-flight batching”): eles mesclam muitas sequências ativas em um lote de decodificação a cada passo. O cache KV viabiliza isso, mas também o torna intensivo em memória porque você mantém muitas sequências vivas simultaneamente.

Esta é uma razão pela qual o layout do cache KV importa: empacotamento eficiente de memória e baixa fragmentação podem se traduzir diretamente em tamanhos de lote sustentáveis maiores e maior vazão.

Para mais sobre o lado de agendamento, veja Batching e Agendamento de Inferência.

Interação com decodificação especulativa

A Decodificação Especulativa tenta reduzir a latência fazendo um modelo pequeno “rascunho” propor múltiplos tokens e então verificando-os com o modelo grande.

O cache KV ainda importa:

  • O modelo grande precisa manter um cache KV correto para os tokens aceitos.
  • Tokens rascunho rejeitados podem exigir rollback ou gerenciamento cuidadoso de atualizações do cache.
  • Implementações frequentemente usam estratégias para evitar regravações caras do cache (por exemplo, anexar e então descartar blocos, ou aplicar atualizações em etapas).

A decodificação especulativa pode reduzir computação por token aceito, mas a memória do cache KV continua sendo uma restrição central à concorrência.

Preocupações operacionais em serving de produção

Modo de falha comum: OOM por sequências “descontroladas”

Como o cache KV cresce a cada token gerado, uma requisição que gera muito mais do que o esperado pode explodir a memória.

Mitigações:

  • impor max_new_tokens e limites máximos de contexto,
  • aplicar controle de admissão (não iniciar requisições que você não consegue terminar),
  • implementar cancelamento/limpeza robustos para que requisições abortadas liberem KV rapidamente.

Fragmentação e overhead de alocação

Se KV for armazenado como tensores contíguos por requisição, comprimentos variáveis podem causar fragmentação e operações caras de realocação/cópia.

Mitigações:

  • alocadores baseados em blocos (KV paginado),
  • estratégias de pré-alocação,
  • políticas de compactação.

Métricas que valem a pena acompanhar

Painéis práticos de serving frequentemente incluem:

  • sequências ativas
  • total de tokens em cache (prompt + gerados)
  • memória de cache KV usada / livre
  • tamanho do lote de decodificação ao longo do tempo
  • TTFT e TPOT (p50/p95/p99)
  • taxas de despejo/paginação (se aplicável)

Quando você poderia desabilitar o cache KV?

Desabilitar o cache KV é incomum para decodificação autorregressiva, mas há casos em que ele não é necessário ou não é benéfico:

  • Pontuação de sequência completa (computar logits para uma entrada fixa sem geração passo a passo): você pode executar um único forward pass e não reutilizar cache.
  • Treinamento (teacher forcing): frameworks computam sequências completas e fazem backprop; cache KV para decodificação não é a principal preocupação (embora existam truques de memória relacionados, como activation checkpointing — assunto diferente).
  • Gerações muito curtas: o benefício existe, mas pode ser menor; a maioria das pilhas ainda mantém ligado porque o overhead é baixo em relação aos ganhos.

Resumo: o trade-off central de engenharia

O cache KV é central para geração rápida em LLMs:

  • Ele reduz a computação na decodificação ao reutilizar chaves/valores de atenção do passado.
  • Ele melhora a latência por token e torna a geração em tempo real prática.
  • Ele consome memória substancial que escala com:
    • comprimento do contexto,
    • tokens gerados,
    • número de camadas,
    • número de cabeças KV,
    • número de sequências concorrentes,
    • estratégia de decodificação (por exemplo, beam search).

Em produção, o melhor desempenho de serving vem de tratar o cache KV como um recurso restrito a ser gerenciado — via variantes de atenção (GQA/MQA), quantização, paginação, cache de prefixo e batching/agendamento inteligentes — em vez de como um detalhe incidental de implementação.