Resumo Técnico: RDKV - Alocação de Bits de Distorção-Taxa para Evicção e Quantização Conjunta do Cache KV

1. Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) enfrentam um gargalo crítico durante a inferência de contexto longo: a largura de banda e a capacidade de memória. O tamanho do cache de Chave-Valor (KV) cresce linearmente com o comprimento da sequência, exigindo a re-leitura de dezenas de gigabytes de dados da memória de alta largura de banda (HBM) fora do chip para a memória dentro do chip a cada etapa de decodificação. Isso torna a inferência limitada pela memória em vez de limitada pelo cálculo.

As estratégias de compressão existentes geralmente se enquadram em duas categorias isoladas:

1. Evicção: Pontuação de tokens ou canais e descarte completo dos de menor classificação (manter ou descartar binário). Isso não deixa nenhum nível de precisão para tokens de importância moderada.
2. Quantização: Retenção de todos os tokens, mas reduzindo sua largura de bits. Isso retém até mesmo tokens negligenciáveis, desperdiçando memória.

Embora híbridos recentes tentem combinar essas abordagens, eles frequentemente dependem de pontuações de roteamento heurísticas ou proporções pré-alocadas determinadas por busca offline, falhando em tratar a evicção e a quantização como um problema de otimização unificado.

2. Metodologia: RDKV

O artigo propõe o RDKV (Compressão de cache KV de Distorção-Taxa), um framework que formula a compressão do cache KV como um problema de distorção-taxa. Nesta visão, a evicção e a quantização não são ações separadas, mas dois extremos de uma única curva de alocação de bits.

2.1 Formulação Teórica

O RDKV deriva um peso de distorção para cada unidade de cache (token no cache V, canal no cache K) com base na perturbação induzida por sua remoção ou quantização no cálculo de atenção.

• Cache V (Tokens): O peso de distorção $w_t$ para um token é a distância de variação total da distribuição de atenção após a evicção. É calculado como a soma das pontuações de atenção recebidas pelo token em todas as posições de consulta: $w_t = \sum_{\tau} a_{\tau,t}$.
• Cache K (Canais): O peso de distorção $w_c$ para um canal é a norma espectral do desvio de logit causado por zerar esse canal. É calculado como o produto das normas das colunas de consulta e chave correspondentes: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

Esses pesos servem como coeficientes de custo em uma função objetivo de distorção-taxa:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{sob} \quad \sum_{u} b_u \leq B$$
onde $b_u$ é a largura de bits atribuída à unidade $u$, e $B$ é o orçamento total de bits.

2.2 Pipeline de Alocação de Bits

A alocação é resolvida via enchimento reverso de água (relaxamento Lagrangiano) adaptado para larguras de bits discretas suportadas por hardware $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. Estágio 1 (Pesagem): Calcule os pesos dos tokens ($w_t$) e os pesos dos canais ($w_c$) uma vez após o estágio de preenchimento usando o cache completo não comprimido.
2. Estágio 2 (Alocação de Tokens): Resolva um Problema da Mochila de Múltipla Escolha (MCKP) para o cache V. Os tokens são atribuídos larguras de bits de $\{0, 2, 4, 8, 16\}$. Uma largura de bits de 0 corresponde à evicção.
3. Estágio 3 (Alocação de Canais): Resolva o MCKP para o cache K, mas apenas para os tokens retidos no Estágio 2. Os canais são atribuídos larguras de bits independentemente da largura de bits V do token, permitindo que um único token tenha componentes K e V de precisão mista.

2.3 Decodificação Compactada Eficiente (TriZone)

Para realizar economias reais de memória e latência, o RDKV introduz um layout de cache TriZone que funde a dequantização no kernel de atenção, evitando a necessidade de descompactar dados para FP16 antes do processamento.

• Zona A: Fluxo de cache antigo compactado contendo linhas K e linhas V retidas com larguras de bits $\{2, 4, 8\}$. Estas são agrupadas por largura de bits para permitir caminhos de dequantização uniformes.
• Zona B: Linhas V em FP16 para tokens atribuídos precisão de 16 bits.
• Zona C: Linhas K e V em FP16 para tokens recém-gerados.

O kernel de atenção calcula logits e saídas combinando essas três zonas, com a dequantização fundida diretamente na multiplicação de matrizes, eliminando sobrecarga extra de largura de banda de memória.

3. Contribuições Principais

1. Formulação de Distorção-Taxa: O artigo apresenta a compressão do cache KV como um problema unificado de distorção-taxa, provando que a solução ótima mistura intrinsecamente unidades de taxa zero (evicção) e unidades de taxa finita (quantização). Isso elimina a necessidade de proporções heurísticas ou buscas offline.
2. Alocação Discreta via Enchimento Reverso de Água: Os autores instanciam a alocação como um problema de mochila discreto solucionável via relaxamento Lagrangiano, atribuindo larguras de bits ótimas de $\{0, 2, 4, 8, 16\}$ com base em pesos de distorção derivados da dinâmica de atenção.
3. Layout Compactado TriZone: Uma implementação consciente de hardware que funde a dequantização no kernel de atenção, convertendo as economias teóricas de largura de bits em reduções reais no pico de memória e na latência de decodificação.

4. Resultados Experimentais

Os experimentos foram conduzidos em LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B e Qwen2.5-72B através de LongBench, RULER e InfiniteBench.

• Precisão: O RDKV supera a melhor linha de base avaliada (Ada-SnapKV) em 9,1% em média. No LongBench, ele recupera 97,81% da precisão do cache completo com apenas 2,48% de retenção de cache.
• Desempenho de Contexto Longo:
  • RULER: Em contexto de 128K, o RDKV alcança 66,95% de precisão comparado a 58,43% para a linha de base mais forte (AdaKV), uma diferença de 8,5 pontos.
  • InfiniteBench: O RDKV ocupa o primeiro lugar entre métodos de compressão com uma pontuação média de 39,46, superando o AdaKV em 1,40 pontos.
  • Agulha no Palheiro: O RDKV mantém precisão de recuperação quase uniforme (média de 0,99) em todas as profundidades, enquanto linhas de base apenas de evicção falham significativamente em profundidades intermediárias.
• Eficiência:
  • Latência: Com um contexto de 128K, o RDKV alcança um aceleração de decodificação de 4,5× comparado à decodificação FlashAttention-2 com cache completo.
  • Memória: Ele alcança uma redução de 1,9× no pico de memória em contexto de 128K e executa com sucesso contexto de 256K em uma única GPU A100 de 64GB, onde o FullKV esgota a memória (OOM).

5. Significado e Alegações

O artigo afirma que tratar evicção e quantização conjuntamente dentro de um único framework de otimização é superior a abordagens em etapas ou baseadas em proporções. Ao derivar pontuações de alocação diretamente da distorção induzida no cálculo de atenção, o RDKV evita a "armadilha binária" de métodos apenas de evicção (que descartam tokens moderados) e o "desperdício" de métodos apenas de quantização (que retêm tokens negligenciáveis).

Os autores enfatizam que seu método:

• Requer nenhum re-treinamento do modelo.
• É agnóstico à arquitetura, performando consistentemente entre diferentes tamanhos e famílias de modelos.
• Fornece um trade-off Pareto-ótimo entre latência e precisão, superando estritamente as linhas de base em ambas as métricas em comprimentos de contexto longos.
• Demonstra que o "ótimo contínuo" da teoria de distorção-taxa, quando discretizado para hardware, gera naturalmente uma estratégia de bits mistos que supera híbridos atuais do estado da arte.

O artigo conclui que essa abordagem conjunta melhora a precisão sem sacrificar a eficiência, permitindo inferência de contexto longo em hardware com recursos limitados.