Resumo Técnico: OSCAR - Rotação Offline Consciente da Covariância Espectral para Quantização de Cache KV de 2 bits

1. Declaração do Problema

A inferência de contexto longo em Modelos de Linguagem Grandes (LLMs) está cada vez mais limitada pela pegada de memória do cache de Chave-Valor (KV). À medida que o comprimento do contexto, o tamanho do lote e a profundidade do modelo aumentam, o cache consome memória significativa da GPU, limitando a vazão e a escalabilidade. Embora comprimir o cache KV para precisão ultra-baixa (por exemplo, INT2) ofereça uma redução de memória de 8×, enfrenta um desafio crítico de precisão: as ativações KV contêm outliers severos por canal.

Sob quantização de baixa precisão, esses outliers dominam a escala de quantização, forçando a maioria dos valores "normais" a poucos níveis de quantização efetivos e degradando severamente a qualidade da atenção. Métodos baseados em rotação existentes (por exemplo, transformadas de Hadamard) tentam suavizar esses outliers redistribuindo a energia entre os canais. No entanto, esses métodos são cegos aos dados; eles suavizam a distribuição bruta do cache, mas falham em alinhar-se com as direções específicas que os mecanismos de atenção a jusante realmente consomem. Em INT2, onde existem apenas quatro níveis de quantização, esse desalinhamento faz com que métodos de rotação ingênuos colapsem, resultando em precisão próxima de zero em tarefas de raciocínio.

2. Metodologia: OSCAR

Os autores propõem o OSCAR (Rotação Offline Consciente da Covariância Espectral), um método que estima estruturas de covariância conscientes da atenção offline para derivar rotações fixas e limites de corte. Diferentemente de trabalhos anteriores que minimizam o erro de reconstrução euclidiana bruto ($\|K - \hat{K}\|_F$), o OSCAR minimiza a distorção nas pontuações e saídas de atenção a jusante.

2.1 Motivação Teórica

O artigo argumenta que o alvo de rotação ideal deve ser derivado das estatísticas da própria computação de atenção:

• Para Chaves ($K$): A distorção nos logits de atenção é controlada pela covariância da consulta $Q^\top Q$, e não apenas pela covariância da chave $K^\top K$.
• Para Valores ($V$): A distorção na saída de atenção é ponderada pelas pontuações de atenção $S$, levando a uma covariância alvo de $V^\top S^\top S V$.

2.2 Design do Algoritmo

O OSCAR opera em duas fases:

A. Fase de Calibração Offline:

1. Estimação de Covariância: Usando um pequeno conjunto de dados de calibração, o método computa covariâncias alvo empíricas:
   • $C_Q = \frac{1}{N} \sum q_n^\top q_n$ (Alvo de chave consciente da consulta).
   • $C_S = \frac{1}{N} V^\top S^\top S V$ (Alvo de valor consciente da pontuação).
2. Derivação de Rotação: O método realiza decomposição espectral nessas matrizes para encontrar as bases de autovetores ótimas ($U_Q$ e $U_S$).
3. Construção de Rotação Composta: As matrizes de rotação finais ($R_K, R_V$) são construídas como um produto de três componentes:
   • $U_Q/U_S$: A base de autovetores consciente da atenção que alinha os canais com direções de importância.
   • $H_{Had}$: Uma transformada de Hadamard que equaliza a métrica de importância entre os canais e espalha a energia dos outliers.
   • $P_{br}$: Uma permutação de reversão de bits que intercala canais de alta e baixa variância para garantir faixas dinâmicas equilibradas dentro dos grupos de quantização.
   • Fórmula: $R_K = U_Q H_{Had} P_{br}$ e $R_V = U_S H_{Had} P_{br}$.
4. Limites de Corte: Limites de corte por camada são estimados a partir das ativações de calibração rotacionadas para controlar outliers.

B. Fase de Serviço Online:

• Layout de Precisão Mista: O sistema mantém um cache KV paginado onde os tokens "sumidouro" (iniciais) e "recentes" são mantidos em precisão completa BF16. O contexto médio histórico é armazenado em INT2 após a aplicação da rotação fixa do OSCAR e do corte.
• Integração de Kernel: Os autores desenvolveram um kernel fundido personalizado em Triton para o framework de serviço SGLang. Este kernel trata a descompactação, desquantização, rotação inversa e computação de atenção para o histórico INT2, garantindo compatibilidade com atenção paginada e cache de prefixo sem introduzir tráfego de memória extra.

3. Contribuições Principais

1. Identificação do Alvo Faltante: O artigo identifica que rotações genéricas espalham outliers, mas falham em INT2 porque não se alinham com a atenção a jusante. O alvo de rotação deve ser induzido pelas estatísticas de atenção ($Q^\top Q$ e $V^\top S^\top S V$), e não pela reconstrução bruta do cache.
2. Framework OSCAR: Um sistema de quantização de cache KV em INT2 totalmente implantável que usa calibração offline para derivar rotações conscientes da atenção. A análise teórica mostra que esses alvos são ótimos sob uma suposição de erro substituto congelado.
3. Sistema Pronto para Produção: Um kernel de atenção INT2 personalizado integrado ao SGLang que suporta layouts de cache KV paginados e cache de prefixo. Ele alcança integração perfeita com frameworks de serviço modernos (SGLang, vLLM) enquanto mantém alta vazão.

4. Resultados Experimentais

Os autores avaliaram o OSCAR em modelos de raciocínio (Qwen3-4B, Qwen3-8B, Qwen3-32B) e um modelo de fronteira (GLM-4.7, 358 bilhões de parâmetros) em cinco benchmarks (AIME25, GPQA, HumanEval, LiveCodeBench, MATH500) e recuperação de contexto longo (RULER-NIAH até 128k).

• Preservação de Precisão:

  • Em Qwen3-4B-Thinking e Qwen3-8B, o OSCAR reduz a lacuna de precisão em relação ao BF16 para 3,78 e 1,42 pontos, respectivamente. Em contraste, a rotação ingênua INT2 colapsa para precisão quase zero.
  • Em Qwen3-32B e GLM-4.7, o OSCAR permanece efetivamente no mesmo nível do BF16 (dentro de 0,02–0,27 pontos), enquanto outras bases INT2 sofrem degradação significativa.
  • Na recuperação de contexto longo (RULER-NIAH), o OSCAR mantém desempenho robusto até 128k tokens, enquanto métodos INT2 ingênuos e apenas Hadamard colapsam em comprimentos muito menores (por exemplo, 16k).

• Desempenho do Sistema:

  • Memória: Reduz a pegada de memória do cache KV em aproximadamente 8× em comparação com BF16.
  • Vazão: Melhora a vazão em até 7× em tamanhos de lote grandes sob o mesmo orçamento de memória.
  • Latência: Acelera a decodificação com tamanho de lote 1 em até 3× sobre BF16 devido à redução da sobrecarga de largura de banda de memória.
  • Escalabilidade: Demonstra a capacidade de lidar com até 28 solicitações simultâneas com entradas de 100k tokens em uma única H100, um regime onde as bases BF16 e INT4 esgotam a memória.

5. Significado e Alegações

O artigo afirma que o OSCAR conecta com sucesso a lacuna entre a promessa teórica do cache KV de ultra-baixa precisão (INT2) e a implantabilidade prática. Ao deslocar o alvo de rotação da reconstrução de dados bruta para a covariância consciente da atenção, o OSCAR recupera a precisão perdida pela quantização ingênua, mantendo ao mesmo tempo os benefícios massivos de memória e vazão do armazenamento de 2 bits.

Os autores enfatizam que sua abordagem não é meramente uma melhoria algorítmica, mas uma solução em nível de sistema. Ao fornecer um kernel personalizado que se integra à atenção paginada e ao cache de prefixo, eles demonstram que a quantização INT2 pode ser usada em frameworks de serviço de produção (como SGLang) sem sacrificar os recursos de eficiência (como reutilização de prefixo) nos quais o serviço moderno de LLM depende. O trabalho sugere que, para raciocínio de contexto longo, a chave para a quantização bem-sucedida reside em alinhar a estratégia de compressão com as demandas estatísticas específicas do mecanismo de atenção.