Resumo Técnico: Quando a Quantização é Gratuita

Declaração do Problema

A inferência de transformadores de longo contexto é tipicamente limitada pela pegada de memória e pelos requisitos de largura de banda do cache Key-Value (KV). À medida que os comprimentos de contexto aumentam (por exemplo, 128K tokens), os tensores KV acumulados em fp16 podem exceder o tamanho dos próprios pesos do modelo, criando um gargalo onde a largura de banda de memória, e não a computação, limita a velocidade de decodificação.

A prática padrão da indústria aborda isso por meio de estratégias de "rotacionar-então-quantizar" (por exemplo, QuaRot, SpinQuant), que aplicam rotações ortogonais aleatórias para gaussianizar as ativações antes de quantizá-las para bits baixos (int4/int8). No entanto, essa abordagem é tradicionalmente enquadrada como um trade-off entre qualidade e latência: o custo computacional da rotação e da de-rotação, combinado com erros de quantização, tipicamente resulta em um aumento de 5–20% na latência de decodificação, mesmo que a memória seja economizada.

Os autores postulam que esse trade-off é invertido no Apple Silicon (arquitetura de memória unificada). Eles hipotetizam que as economias de largura de banda de um cache comprimido 3× podem superar o custo computacional do kernel de quantização, desde que o kernel seja altamente otimizado e fundido.

Metodologia

1. Transformada de Fourier Randomizada por Sinais (SRFT)

O artigo adota a SRFT (Transformada Rápida de Fourier Randomizada por Sinais) como mecanismo de rotação, definida como $SRFT(x) = \text{pack}(F \cdot \text{diag}(s) \cdot x)$, onde $F$ é a DFT unitária e $s$ é um vetor de sinais aleatórios fixo.

• Racional: Embora transformadas de Hadamard (SRHT) sejam comuns, a SRFT oferece flexibilidade para dimensões que não são potências de dois e alinha-se com operações de multiplicação de matrizes 16×16 (AMX, Tensor Cores).
• Descoberta: Os autores demonstram que SRFT e SRHT são estatisticamente indistinguíveis em termos de qualidade de perplexidade (PPL) em cinco modelos (SmolLM2, Qwen2.5, Gemma-3) e vários orçamentos de bits.

2. Kernel Metal Fundido

Para eliminar a sobrecarga de despacho, os autores implementam um único kernel de computação Metal fundido que executa todo o pipeline:

1. Carregar e Multiplicar por Sinal: Aplicar o sinal diagonal aleatório no registro.
2. FFT: Executar passagens de borboleta radix-2 de Stockham log2(d/2).
3. Empacotamento: Desempacotar a simetria hermitiana para uma representação real.
4. Escala e Quantização: Realizar redimensionamento $\lambda$ por canal (aprendido ou calibrado), redução de máximo absoluto por grupo e arredondamento para int4/int8 com sinal.
5. Empacotamento de Nibbles: Empacotar dois valores int4 em um único byte usando embaralhamentos SIMD.

Esse kernel é exposto via torch.mps.compile e integrado como uma subclasse Cache do HuggingFace (SRFTInt4Cache).

3. Estratégias Avançadas de Escala

• Por Canal + Por Grupo: Para modelos com dimensões de cabeça pequenas (por exemplo, Qwen2.5-1.5B, $d=128$), a escala por token falha catastróficamente devido a coordenadas de outliers. Os autores combinam escala por canal (uma escala por coordenada) com escala por grupo (grupos de 16) para recuperar a qualidade.
• Rotações Aprendidas: O artigo investiga calibração pós-treinamento onde uma rotação aprendível $R$ e escala $\lambda$ são adicionadas. Uma descoberta chave é que aprender $R+\lambda$ sem a base SRFT minimiza significativamente mais o MSE de calibração (84,9% vs 50,3%), mas resulta em PPL a jusante pior, sugerindo que a SRFT aleatória fixa atua como um regularizador.
• Rotações de Householder: Em $d=256$, rotações de Householder com $k=d/2$ refletores alcançam quantização efetivamente sem perdas ($\Delta PPL \approx 0.000$) com metade dos parâmetros de rotações de Cayley completas.

4. Caminho de CPU (AMX)

Um caminho de CPU complementar utiliza o coprocessador de matriz AMX da Apple via cblas sgemm. Ao agrupar DFTs de comprimento 8 como multiplicações de matrizes reais 16×16, os autores alcançam até 91% do pico de AMX de núcleo único, oferecendo uma alternativa de baixa latência para regimes de pequenos lotes onde a sobrecarga de despacho da GPU é dominante.

Resultados Principais

Desempenho (Latência)

No Apple M1 (16 GB de memória unificada), o cache KV int4 fundido supera o fp16 na latência de decodificação de ponta a ponta:

• Gemma-3 1B: Executa 3–8% mais rápido (sobrecarga de latência negativa) em prefixos de 256–4096 tokens.
• Qwen2.5-1.5B: Executa 0,7–2,6% mais rápido em contextos curtos (≤1K tokens).
• Mecanismo: A sobrecarga do kernel (~25 ns/vec) é menor que as economias de largura de banda de memória alcançadas ao transferir 3× menos dados por etapa de decodificação.

Qualidade (Perplexidade)

• SRFT vs. SRHT: Nenhuma diferença estatisticamente significativa em $\Delta PPL$ em todos os modelos testados.
• Qwen2.5-1.5B: O kernel fundido com escala por canal + por grupo reduz o $\Delta PPL$ de 4 bits de um catastrófico +7975 (linha de base por token) para +638,6 (uma redução de 12,5×), igualando a implementação de referência em Python.
• Gemma-3 1B: O $\Delta PPL$ de 4 bits melhorou de +18,7 (por token) para +3,6 com o kernel escalado fundido.
• Regime Sem Perdas: Em $d=256$, rotações de Householder aprendidas alcançam $\Delta PPL = 0.000$.

Eficiência de Memória

• Compressão: Alcança 3,0× a 3,56× de compressão de memória persistente em camadas de atenção completa.
• Impacto no Mundo Real: Para um contexto de 128K no Llama-3-70B, o cache int4 reduz a memória de ~41 GB (fp16) para ~10,9 GB.

Significado e Alegações

O artigo afirma que a quantização é "gratuita" (positiva em throughput) em hardware de memória unificada, invertendo a suposição padrão da indústria de que a quantização de bits baixos sempre incorre em uma penalidade de latência.

1. Inversão do Trade-off: A principal contribuição é demonstrar que, no Apple Silicon, as economias de largura de banda de um cache comprimido 3× superam o custo computacional do kernel de quantização. Isso torna o cache KV int4 estritamente superior ao fp16 para inferência de longo contexto sensível à latência em laptops e dispositivos de borda.
2. Design Consciente de Hardware: O trabalho destaca que a escolha da rotação (FFT vs. Hadamard) é menos sobre superioridade matemática e mais sobre alinhamento de hardware (por exemplo, a estrutura radix-8 da FFT correspondendo aos blocos 16×16 do AMX).
3. Separação entre Calibração e PPL: O artigo isola um fenômeno novel onde minimizar o MSE de calibração não garante melhor PPL a jusante, sugerindo que a rotação aleatória fixa fornece um prior de regularização necessário.
4. Viabilidade de Implantação: Ao fundir todo o pipeline em um único kernel Metal e integrá-lo ao HuggingFace, os autores provam que a quantização agressiva de 4 bits pode ser implantada sem sacrificar a qualidade ou a velocidade do modelo, desde que a arquitetura de hardware (memória unificada) e a otimização de kernel (fusão) estejam alinhadas.

Os autores concluem que, enquanto o trade-off padrão de quantização se mantém para GPUs com HBM separada e altas sobrecargas de despacho, ele é invertido para o regime de memória unificada onde operam a inferência em laptops e dispositivos de borda.