Compressão e Poda Dinâmica de KV Cache: StreamingLLM, SnapKV, PyramidKV e DuoAttention

A inferência de LLMs em contextos longos esbarra no teto físico de VRAM.

O armazenamento de tensores de Chave e Valor (KV Cache) cresce linearmente com a janela.

Sem compressão, um lote de 16 requisições em 128k tokens consome centenas de gigabytes. Essa carga satura os barramentos HBM3 e derruba o throughput dos aceleradores.

Quatro abordagens resolvem esse gargalo: attention sinks, retenção de heavy-hitters, poda guiada no prefill e decomposição funcional de cabeças de atenção.

O gargalo de memória do KV Cache em contextos longos

O consumo de memória do KV Cache decorre da necessidade de reter tensores passados para evitar o recálculo quadrático da autoatenção. Cada token processado gera tensores que permanecem na VRAM durante toda a geração.

A equação exata do dimensionamento em bytes por sequência é:

Memoria_KV = 2 × n_layers × n_heads_kv × d_head × seq_len × bytes_por_elemento × batch_size

Onde n_layers representa as camadas e n_heads_kv as cabeças KV. A variável d_head é a dimensão interna e bytes_por_elemento a precisão numérica.

Em arquiteturas com GQA (EMNLP 2023), a proporção n_heads_kv = n_heads_q ÷ 8 atenua o problema. Ainda assim, em janelas de 128k a 1M tokens o consumo continua crítico.

A tabela abaixo detalha o consumo de VRAM por fluxo individual em precisão FP16:

O gargalo se agrava na decodificação autorregressiva (decode). Cada token emitido lê todo o KV Cache da HBM para executar uma operação matriz-vetor (GEMV).

A intensidade aritmética no decode cai para 1 a 2 FLOPs por byte lido. Isso satura os canais de memória enquanto as unidades de processamento permanecem ociosas.

Attention Sinks e o StreamingLLM para fluxos contínuos

O StreamingLLM viabiliza janelas de contexto infinitas ao fixar os primeiros tokens da sequência como âncoras numéricas estáveis. Os modelos baseados em Transformer concentram scores desproporcionais nos primeiros tokens (t_0 a t_3), independentemente do conteúdo semântico.

O estudo de Xiao et al. (ICLR 2024) analisou a Softmax. O operador exige soma unitária das probabilidades calculadas.

Quando tokens intermediários não têm afinidade com a consulta, o modelo descarrega o residual numérico nos tokens inaugurais.

Se uma janela deslizante descarta esses tokens inaugurais, o denominador da Softmax colapsa. A perplexidade salta instantaneamente de 12 para valores superiores a 10^3.

Estrutura StreamingLLM: [ t_0  t_1  t_2  t_3 ] + [ Janela Móvel: t_{N-W} ... t_N ]
                        \___________________/   \_____________________________/
                           Attention Sinks              Contexto Recente
                           (4 tokens fixos)             (W tokens móveis)

O StreamingLLM divide o buffer de KV Cache em duas partições contíguas:

• Tokens de Âncora (Attention Sinks): Preserva fixos os 4 primeiros tokens (t_0 a t_3) na VRAM.
• Janela Móvel (Rolling Window): Retém os últimos W tokens recentes em um buffer circular.

A complexidade de memória passa de O(N) linear para O(W + 4) = O(1) constante. O modelo sustenta mais de 4 milhões de tokens contínuos sem colapso de perplexidade.

Heavy-Hitter Oracle (H2O) e a esparsidade de cauda pesada

O algoritmo H2O reduz a memória do KV Cache descartando tokens com base no acúmulo histórico de pontuações de atenção. A autoatenção em LLMs exibe distribuição de cauda pesada (heavy-tailed sparsity), onde menos de 20% dos tokens concentram mais de 80% do peso de ativação.

No paper de Zhang et al. (NeurIPS 2023), o H2O rastreia o valor do token j. Ele soma a atenção recebida a cada passo t:

s_j = ∑_{t} α_{t,j}

Quando o cache alcança o limite alocado (K_{budget}), o gerenciador executa a poda seletiva:

1. Partição Recente: Reserva cota garantida para os tokens locais mais novos (L_{local}).
2. Partição Heavy-Hitters: Ordena os tokens históricos pelo score s_j e retém apenas os K_{HH} maiores.
3. Poda e Desalocação: Remove os tensores dos tokens descartados da memória física.

Em tarefas como GSM8K e CNN/DailyMail, o H2O descarta até 80% do cache. A perda de acurácia fica abaixo de 0,5%.

Em tarefas de busca pontual (Needle In A Haystack), o método pode descartar fatos cruciais se eles não acumularem atenção prévia.

SnapKV e a seleção de contexto no prefill

O SnapKV comprime o KV Cache no final do prefill antes da geração de tokens, aproveitando a consistência dos mapas de atenção de cada cabeça. Cada cabeça de atenção concentra sua ativação em posições estruturadas previsíveis durante o processamento do prompt.

O estudo do SnapKV (2024) analisou o prefill. Uma pequena janela final basta para identificar os tokens críticos do texto.

O pipeline de compressão do SnapKV opera em quatro etapas:

Passo 1: Janela de Observação (L_obs) --> Extrai matriz de atenção A_{h, L_obs, N}
Passo 2: Agregação por Média          --> Calcula score bruto s_i = (1 ÷ L_obs) ∑ A_{h, t, i}
Passo 3: Pooling Convolucional 1D     --> Suaviza scores S = Conv1D(s, kernel_size = 5)
Passo 4: Poda Top-K por Cabeça        --> Mantém as posições com maior valor

O pooling convolucional 1D preserva a coerência contextual. Tokens vizinhos a palavras de alta atenção contêm valor sintático indispensável.

Ao aplicar o filtro antes da seleção top-k, o SnapKV retém blocos semânticos íntegros, evitando a perda de entidades e termos técnicos.

O algoritmo retém apenas 16% do KV Cache em sequências de 32k tokens mantendo a acurácia de recuperação acima de 98%.

Alocação piramidal com o PyramidKV

O PyramidKV melhora a retenção de memória ao distribuir orçamentos de KV Cache assimétricos entre as camadas da rede. O fluxo de informação em Transformers segue um funil de abstração ao longo da profundidade.

Conforme demonstrado por Cai et al. (EMNLP 2024), camadas iniciais exigem contexto detalhado. Camadas superiores processam conceitos condensados e suportam podas mais agressivas.

O PyramidKV aplica uma função linear decrescente para determinar o orçamento B_l na camada l:

B_l = B_{base} + ΔB × ((N_L - 1 - l) ÷ (N_L - 1))

Onde N_L é o total de camadas, B_{base} o orçamento mínimo no topo e ΔB o gradiente de alocação.

A restrição orçamentária total obedece à meta global de capacidade:

∑_{l=0}^{N_L - 1} B_l = N_L × B_{alvo}

A distribuição piramidal aumenta a precisão em raciocínios longos sem introduzir sobrecarga computacional na inferência.

Decomposição funcional com o DuoAttention

O DuoAttention elimina desperdícios ao separar cabeças de recuperação global de cabeças puramente locais. Apenas uma fração das cabeças em um modelo de contexto longo precisa de acesso a todo o histórico.

O estudo de Xiao et al. (NeurIPS 2024) comprovou essa especialização funcional:

• Retrieval Heads (Recuperação): Representam de 20% a 30% do total. Localizam informações dispersas e exigem retenção de KV Cache completo (O(N)).
• Streaming Heads (Locais): Representam de 70% a 80% do total. Operam com buffer constante de O(1) (64 tokens locais e 4 attention sinks).

A classificação das cabeças utiliza otimização com regularização L1 antes do deploy:

min_{γ} L_{val}(γ) + λ × ‖γ‖_1 sujeito a γ_{l,h} ∈ [0, 1]

Onde γ_{l,h} = 1 define uma cabeça de recuperação e γ_{l,h} = 0 uma cabeça local. A calibração demora menos de 15 minutos em uma GPU.

Com kernels de FlashDecoding esparso, o DuoAttention alcança 100% de acurácia em Needle In A Haystack em 128k tokens, economizando até 70% de VRAM.

Implementação prática em vLLM e SGLang

A aplicação eficiente da poda de KV Cache em produção exige integração com o gerenciamento de páginas de memória. O descarte de tokens isolados quebra o alinhamento de memória de 128 bytes da GPU.

Motores como vLLM e SGLang adotam despejo por bloco físico (Block-Level Eviction).

No PagedAttention, os tensores são alocados em blocos de 16 ou 32 tokens. A engine soma a importância dos tokens para avaliar o bloco inteiro:

Score(Bloco_k) = ∑_{t ∈ Bloco_k} s_t

Ao esgotar a VRAM, o agendador desaloca os blocos com menor score na tabela de páginas (Page Table). Os blocos retornam ao pool livre sem cópias de memória.

No SGLang, a árvore RadixAttention retém prefixos comuns ponderados pela importância das cabeças de recuperação.

Para acelerar a leitura dos vetores mantidos, os runtimes usam o FlashDecoding. Ele paraleliza o cálculo ao longo do tempo.

Comparativo de desempenho e benchmarks

O equilíbrio entre redução de memória, latência de resposta e fidelidade determina o método ideal para cada aplicação. Avaliações no Benchmark RULER (arXiv:2404.06654) e no LongBench comprovam esses comportamentos.

A tabela compara os métodos no Llama-3-8B com contexto de 128k tokens em nó NVIDIA H100:

O StreamingLLM maximiza o throughput em tarefas de fluxo contínuo. Para recuperação factual em documentos longos, o DuoAttention e o PyramidKV entregam a maior fidelidade com economia substantiva de hardware.

Diretrizes de arquitetura para clusters de produção

A seleção do mecanismo de compressão deve alinhar o padrão de tráfego aos requisitos de recuperação factual do sistema.

Para agentes autônomos e atendimentos contínuos, utilize o StreamingLLM com 4 attention sinks e buffer de 4.096 tokens. A complexidade O(1) elimina o risco de estouro de memória em sessões longas.

Para análise documental e fluxos RAG com janelas de 64k a 256k tokens, adote o DuoAttention ou o PyramidKV. Essa configuração preserva fatos pontuais e quadruplica a densidade de requisições por GPU.

A compressão estruturada de tensores viabiliza o atendimento escalável de contextos massivos com custos operacionais controlados.

Perguntas Frequentes sobre Compressão de KV Cache

Por que o KV Cache satura a memória em contextos longos?

O KV Cache armazena tensores para cada camada e token processado. Em contextos longos sob alto tráfego, esse volume supera rapidamente o tamanho dos pesos do próprio modelo.

O que são Attention Sinks e como impedem o colapso do modelo?

São os primeiros 4 tokens da sequência que absorvem a massa residual de probabilidade da Softmax. Fixá-los na memória estabiliza o denominador do cálculo em buffers circulares.

Qual é a diferença prática entre o SnapKV e o PyramidKV?

O SnapKV poda tensores após o prefill de forma homogênea nas camadas. O PyramidKV aloca mais memória nas camadas iniciais e reduz o espaço nas camadas superiores.

O que são as Retrieval Heads no algoritmo DuoAttention?

São cabeças de atenção (20% a 30% do total) que recuperam fatos dispersos no histórico. O DuoAttention mantém cache completo nessas cabeças e poda as demais para buffers locais.

Como a paginação por blocos evita a perda de desempenho na GPU?

Ao agrupar tensores em blocos físicos de 16 ou 32 posições (como no PagedAttention), o motor descarta páginas inteiras, mantendo os acessos de memória alinhados e sem fragmentação.