O Seu Cache KV Não Tem um Problema de Bits. Tem um Problema de Geometria.

Em precisão de 2 bits idêntica, uma decisão sobre qual eixo quantizar alonga uma pontuação de benchmark de 2,88 para 63,53. Chaves e valores precisam de tratamento oposto — e o motivo está na equação de atenção, não no hardware.

Considere o Llama-2-13B. Agrupe seu cache de chave-valor por um tamanho de grupo de quantização de 32 em dois bits, enquanto deixa tudo o mais no lugar — mesmo modelo, mesmo orçamento de bits, mesmos tamanhos de grupo, mesmos benchmarks.

Dependendo de um único elemento de uma decisão de implementação, os resultados de precisão do CoQA resultam em 2,88 ou 63,53. A pontuação usando precisão total é 66,37.

A decisão não é sobre quantos bits totais são usados. A pergunta é simplesmente qual eixo você escolhe para agrupar ao longo da computação de cada fator de escala? Quando você decide usar o canal como sua dimensão de agrupamento (chaves) e o token como sua dimensão de agrupamento (valores), você acaba em algum lugar dentro de quatro pontos do desempenho de precisão total. Se você inverter qualquer uma dessas escolhas, você experimenta uma perda de qualidade. Se você inverter ambas as escolhas, o modelo não funciona mais.

Quatro maneiras de gastar os mesmos 2 bits no mesmo cache. Resultados da ablação KIVI no Llama-2-13B com tamanho de grupo 32.

A quantização é normalmente pensada como apenas um dial: 8 bits, 4 bits, 2 bits, com um custo de precisão suave anexado. Dentro do cache KV, não é assim. É escolher sistemas de coordenadas, e diferentes sistemas se aplicam a chaves e valores. Este artigo explica por quê. Brevemente: o erro de quantização depende do intervalo de valores dentro dos grupos; chaves e valores têm estruturas muito diferentes; e as pessoas frequentemente tropeçam porque você não pode derivar o eixo correto da distribuição de valores em tudo. Você tem que olhar como o erro muda após a atenção consumi-lo. Isso fornece um princípio geral para comprimir ativações intermediárias e um bom motivo para duvidar do erro de reconstrução como um proxy para a qualidade.

Por Que o Cache KV É Onde Isso Morde

Durante a fase de geração, um transformador armazena todos os dados de projeção de chave-valor (KV) de tokens que ele processou anteriormente em um cache para que não precise recalcular esses dados novamente. Esse cache cresce linearmente com o comprimento do contexto e o tamanho do lote. Eventualmente, isso resultará no cache crescer maior do que o modelo em si.

Esse aumento de crescimento pode ser facilmente identificado ao olhar para o consumo de memória de diferentes partes do modelo. Na análise KVQuant do LLaMA-7B, os pesos respondem por cerca de 98 por cento da memória em um comprimento de sequência de 512, com ativações em 2 por cento. Em 128K de contexto, a proporção se inverte para cerca de 16 por cento de pesos e 84 por cento de cache KV. Quando olhamos para uma análise do OPT-175B citada pelos autores do KIVI, eles encontraram resultados semelhantes. Especificamente, em um tamanho de lote de 512 com um prompt de 512 tokens, o cache KV atinge 1,2 TB — várias vezes o tamanho dos pesos do modelo.

No entanto, a capacidade é apenas metade do problema aqui. A GPU deve ler o cache KV inteiro da memória do dispositivo para cada token que gera. Isso significa que, enquanto a GPU lê o cache KV, os núcleos de processamento ficam ociosos. Portanto, reduzir o tamanho geral do cache aumenta a capacidade de processamento disponível e reduz o tempo gasto esperando transferências de dados.

O Que o Erro de Quantização É Realmente Feito

A quantização de inteiro uniforme é matematicamente direta. Para um grupo de números, você registra o menor número como um ponto zero e divide o intervalo desse grupo pelo número de níveis que podem ser representados para obter um tamanho de passo. Em seguida, arredonda cada elemento para o passo mais próximo. Dois resultados imediatos seguem. Primeiro, o erro por elemento é limitado por metade de um passo. Segundo, o tamanho do passo é o intervalo do grupo dividido por 2ᴮ − 1. Em 2 bits, você tem apenas 4 níveis para cobrir qualquer dispersão que exista dentro desse grupo. Então, um elemento que é cem vezes maior em comparação com seus vizinhos não apenas se sai mal. Ele infla o tamanho do passo para todos os outros elementos que compartilham o mesmo grupo, e todos eles se tornam mais grossos juntos. O grupo é a unidade de dano. Escolher um eixo significa decidir quais elementos sofrem juntos. Formulando a pergunta de forma diferente, não é mais “quantos bits posso me dar ao luxo de sacrificar?” mas “onde estão os valores extremos e posso isolá-los?”

Chaves: Os Valores Extremos Vivem em Canais Fixos

Os grandes modelos de linguagem contêm ativações que são incomumente grandes em comparação com a maioria das ativações. Sun e colegas catalogaram essas ativações muito grandes em diferentes famílias de modelos: no Mixtral 8x7B, a magnitude mais grande está perto de 7000, enquanto a magnitude média da característica é de cerca de 0,3 — cerca de quatro ordens de magnitude aparte. Esses são muito raros; eles permanecem fixos em dimensões que raramente mudam com a entrada e não são acidentais. Eles atuam como vieses implícitos e são o que foca a atenção em apenas alguns tokens: comportamento de sumidouros de atenção. No cache de chaves, essa estrutura é muito clara: canais específicos carregam magnitudes muito grandes consistentemente em todos os tokens de uma sequência. Agrupe ao longo dos tokens e cada grupo contém esses canais de valores extremos, então o tamanho do passo de cada grupo é definido pelos valores extremos, e todos os canais comuns pagam por isso. Agrupe ao longo dos canais e os canais de valores extremos formam seus próprios grupos. Seu intervalo interno é grande, mas autocontido; os canais comuns são deixados sozinhos. Os resultados correspondem. Média em camadas e cabeças no Llama-2-13B, o KIVI relata um erro de reconstrução de chave de 13,67 sob agrupamento por token contra 4,55 por canal, e — mais importante — erro de pontuação de atenção de 47,00 contra 9,60. A quantização de chaves por token produz cerca de cinco vezes o erro de pontuação. As pontuações concordam com métricas significativas para chaves; a quantização de canal se destaca em ambos os aspectos.

Valores: Onde a Intuição Quebra

O cache de valor não mostra um padrão de canal de valor extremo. Ele parece ser bastante plano. Sozinho, pelo argumento de intervalo, poderíamos esperar que qualquer um desses eixos produzisse uma qualidade semelhante de compressão.

Eles não o fazem. Independentemente de como o gerenciamento de chaves é implementado (os resultados de 2,80 e 2,88), a compressão por canal de valores colapsa o modelo.

E aqui está a pegadinha: se você medisse essa perda usando o erro de reconstrução bruto na tensão original para a qual cada valor foi comprimido, a quantização de valor por canal na verdade parece ligeiramente melhor, com 3,73 contra 4,57. Se você validasse a compressão da maneira óbvia, você escolheria a configuração que destrói o modelo.

Erro de quantização do cache de valor no Llama-2-13B, medido de duas maneiras. A métrica de tensor armazenado e a métrica de saída consumida discordam por mais de uma ordem de magnitude.

A resolução é que o cache de valor nunca é lido diretamente. Ele é consumido por um produto matricial: a saída de atenção é uma soma ponderada de vetores de valor em todos os tokens, com pontuações de atenção softmax como pesos. Devido a isso, o erro relevante é o introduzido durante esse processo e não dentro das próprias tensões. Medido em termos da saída de atenção, a ordem foi completamente revertida. O erro relativo relatado pelo KIVI para a saída de atenção devido à quantização de vetor de valor por token foi de 3,55 em comparação com 49,89 para a quantização por canal — mais de quatorze vezes maior para o que parecia a melhor escolha com base em como bem ele foi comprimido.

A explicação é a esparsidade da atenção, que eles mediram como 84,3 por cento. A maioria das informações contidas na saída pode ser atribuída a um pequeno número de tokens muito importantes. A quantização por token confina o erro de cada token àquele token, então os erros em tokens não importantes são multiplicados por pesos de atenção próximos a zero e efetivamente desaparecem. A quantização por canal espalha o erro de cada token por uma escala de canal compartilhada, então tokens mal representados contaminam a representação daqueles que importam. A esparsidade que torna a atenção eficiente é a mesma propriedade que torna a quantização por token segura.

A lição transferível é mais ampla do que o cache KV: medir o erro de compressão onde a tensão é consumida, não onde é armazenada. Um pressuposto implícito feito pelo erro de reconstrução é que cada componente de uma tensão tem peso igual ao contribuir para a saída final. A atenção explicitamente não o faz. Qualquer operação downstream que pondera, gera ou esparsifica sua entrada quebra esse pressuposto. Os leitores familiarizados com meu artigo anterior sobre cegueira de contradição em sistemas de recuperação reconhecerão que esses resultados são semelhantes aos falhas previamente descritas: métricas facilmente computadas que relatam algo diferente do que foi pretendido.

Embebedos Rotatórios Complicam as Chaves

Há alguns problemas ao usar Embebedos de Posição Rotatórios (RoPE). O RoPE gira pares de canais com base na posição relativa de cada token. Essa mistura dissolve parcialmente a estrutura de canal fixo que fez a quantização de chave por canal funcionar no início — um canal de valor extremo é girado para os vizinhos e os vizinhos herdam o intervalo. A resposta do KVQuant é a ordem: quantize as chaves antes de aplicar a rotação e aplique o RoPE após a desquantização. Ao lado da quantização de chave por canal, tipos de dados não uniformes e isolamento de uma pequena fração de valores extremos, isso obtém menos de 0,1 de degradação de perplexidade em 3 bits e permite servir o LLaMA-7B em até 1 milhão de tokens de contexto em um único A100-80GB.

Também é importante entender o nível de impacto do RoPE. Os autores do artigo “RotateKV” relataram um aumento de 145 por cento nos erros de quantização uma vez que o RoPE foi adicionado e notaram que canais de valores extremos diferem entre cabeças de atenção — é por isso que aplicar uma matriz de rotação compartilhada em todos os lugares é insuficiente e rotações adaptativas de cabeça fazem melhor.

O Imposto dos Sistemas e Por Que Não É um Detalhe

A quantização por token se adapta bem à decodificação. Cada token chega; você o quantiza, adiciona-o à sequência (ao longo da dimensão do token), nada mais se move.

No entanto, a quantização por canal não se encaixa. Como as estatísticas de um canal abrangem tokens que ainda não foram gerados, você não pode computar um fator de escala quando um token chega. A solução do KIVI é manter os tokens mais recentes — até 128 — em precisão total em um buffer residual e quantizar em grupos uma vez que o suficiente tenha sido acumulado.

Como acontece, o buffer residual se torna de carga, em vez de apenas uma coisa incidental. No GSM8K com o Llama-2-7B, as pontuações de precisão total são 13,50. Totalmente quantizado para 2 bits com os eixos corretos, ele pontua 5,76. Os mesmos eixos e os mesmos bits, mais o buffer residual de tokens recentemente produzidos em precisão total, pontuam 12,74. Uma janela deslizante de tokens recentemente produzidos em precisão total recuperará muito do que foi perdido devido à quantização agressiva em problemas de múltiplos passos — o que faria sentido se considerarmos quais tokens estavam sendo atendidos por uma cadeia de operações aritméticas.

Há um benefício significativo em fazer todas essas coisas corretamente — como o KIVI relata, 2,6 vezes menos uso de memória de pico para o Llama-2-7B, permitindo tamanhos de lote até 4 vezes maiores, bem como 2,35 a 3,47 vezes melhor taxa de transferência em uma tarefa de serviço real.

O Que Fazer Com Isso

1. Nunca use um quantizador para ambos. Use quantizadores diferentes para chaves (por canal) e para valores (por token). Um pipeline que aplica um único quantizador ao “cache KV” provavelmente já sacrificou a maior parte da qualidade possível ao usar um pequeno número de bits para representar cada valor.
2. Quantize as chaves antes do RoPE. Isso é uma questão de correção em vez de uma questão de preferência.
3. Armazene uma janela de precisão total de tokens recentemente gerados. Embora armazenar tal janela leve muito pouca memória em comparação com o quão grande um cache pode ser, é exatamente essa área que gera muita precisão para tarefas difíceis.
4. Não valide com base no erro de reconstrução. Sempre valide com base na saída de atenção ou no desempenho da tarefa final. A métrica de armazenamento não é apenas ruído — para os valores, ela aponta na direção errada.
5. Não valide com base em benchmarks de múltipla escolha de contexto curto. Os autores do KIVI deliberadamente evitam tarefas fechadas como o MMLU para essa avaliação, porque um único passo de decodificação lendo logítimos de saída mal exerce o cache. Qualquer avaliação que não construa um cache ao longo do tempo e, em seguida, execute a geração a partir dele nunca poderá observar os fracassos inerentes ao design do sistema.

Onde o Trabalho Está Indo

Embora ainda haja algo a ser feito com respeito à natureza geométrica do problema, muitos pesquisadores continuam a estudar maneiras pelas quais canais de valores extremos são distribuídos entre as várias cabeças do transformador e como as limitações de hardware afetam quais agrupamentos são mais baratos: o InnerQ dobra a normalização de chave por canal nas pesadas de chave e consulta durante o pré-preenchimento. Portanto, nenhuma sobrecarga adicional é incorrida no tempo de execução. Além disso, o InnerQ armazena janelas de alta precisão para tokens recentemente gerados e tokens de sumidouros de atenção. Ao fazer isso, o InnerQ elimina a oportunidade para que os valores extremos no canal de sumidouro contaminem canais vizinhos.

Outros propõem que, em vez de armazenar o cache inteiro, devemos armazenar apenas informações suficientes para ser capaz de recriar a chave e/ou valor(s) sob demanda a partir de uma representação em cache menor.

Finalmente, é importante lembrar que a precisão não é o único parâmetro que a quantização afeta. Pesquisas recentemente publicadas demonstraram degradação de alinhamento resultante da quantização de caches KV. Além disso, essa pesquisa documentou a degradação de alinhamento mesmo em ambientes de serviço de vLLM de produção que utilizam caches FP8, juntamente com um protocolo de recuperação sem treinamento que restaurou até 97 por cento do que foi perdido em termos de alinhamento. Portanto, embora uma configuração possa manter seus resultados de benchmark, não significa necessariamente que ela retém todos os outros parâmetros relevantes que você se importa.

O Princípio Geral

A ideia de quantização foi enquadrada como um “orçamento de precisão”: quantos bits posso me dar ao luxo de sacrificar? O cache KV mostra que a pergunta mais útil é estrutural. A precisão é alocada em grupos; o grupo é a unidade de dano e o eixo que você agrupa determina quais elementos compartilham seu destino. O eixo correto é aquele em que sua tensão está sendo consumida, ou seja, a maneira como você está usando sua tensão e NÃO como sua tensão aparece quando armazenada na memória. As chaves são usadas por meio de um cálculo de produto escalar contra a consulta. Um canal corrompido único envenenará todas as pontuações. Os valores são consumidos por meio de uma computação de média ponderada esparsa ao longo dos tokens. Portanto, um token corrompido único é simplesmente ponderado.

Duas tensões de dimensões idênticas e geradas por camadas consecutivas são tratadas de forma diferente. Vale a pena perguntar sobre qualquer ativação que você planeja comprimir: qual operação contrai isso e o meu agrupamento respeita isso?