Resumo Técnico: Pesquise as Escalas de Ponto Flutuante em Blocos!

Declaração do Problema

A quantização é uma técnica padrão para acelerar a inferência de modelos generativos, mas manter a precisão em larguras de bits ultra-baixas continua sendo um desafio. Embora os formatos de Ponto Flutuante em Blocos (BFP) ofereçam um equilíbrio entre alcance dinâmico, eficiência de memória e taxa de transferência, os algoritmos BFP padrão dependem de uma escala fixa determinada pela magnitude absoluta máxima dentro de um bloco. Os autores observam que essa heurística de "escala máxima" é frequentemente subótima, levando a erros de quantização desnecessários. Essa questão é particularmente relevante com o advento da arquitetura Blackwell da NVIDIA, que introduz suporte de hardware para formatos BFP de microescala como o NVFP4. Diferentemente do BFP tradicional, o NVFP4 inclui bits de mantissa dentro de sua escala de bloco de 8 bits (formato E4M3), fornecendo uma granularidade mais fina para a seleção de escala que a escala máxima padrão não consegue explorar.

Metodologia

ScaleSearch

A contribuição central é o ScaleSearch, um algoritmo que aproveita os bits de mantissa em formatos de microescala para realizar uma busca de alta granularidade pela escala de bloco ótima.

• Mecanismo: Em vez de selecionar rigidamente a escala com base no valor máximo do bloco ($s = \text{round}(\|x\|_\infty / 6.0)$), o ScaleSearch explora uma vizinhança de escalas ao redor desse valor padrão. Ele trata a escala como um deslocamento inteiro $f$ aplicado à representação de bits da escala padrão: $s(f) = \text{reinterpret}(\text{reinterpret}(s, \text{int8}) + f, \text{fp8}_{UE4M3})$.
• Estratégia de Busca: O algoritmo busca exaustivamente um pequeno intervalo de deslocamentos (por exemplo, $f \in [-2, 6]$) para encontrar a escala que minimiza o Erro Quadrático Médio (MSE) para a distribuição do bloco dada.
• Racional: A análise empírica em dados Gaussianos sintéticos e ativações reais de redes neurais (por exemplo, estados Chave do Llama 3.1) revela uma distribuição bimodal de deslocamentos ótimos. Enquanto a escala padrão (deslocamento 0) mapeia a maior magnitude para o valor FP4 6, uma parte significativa dos blocos é melhor representada mapeando a maior magnitude para o valor FP4 4, exigindo um fator de escala aproximadamente 1,5 vezes maior. O ScaleSearch identifica esses mapeamentos ótimos.

ScaleSearchAttention

Para estender esses benefícios ao mecanismo de atenção e ao cache KV, os autores propõem o ScaleSearchAttention. Este é um algoritmo de atenção FP4 de ponta a ponta projetado para modelagem de linguagem causal que evita sobrecargas de desquantização.

• Execução Nativa NVFP4: Realiza multiplicações de matrizes QKT e PV diretamente nos Tensor Cores usando precisão NVFP4 com acumulação FP32.
• Processamento de Incoerência (IP) e Redução de Magnitude: Para mitigar valores atípicos e reduzir o erro de quantização nos tensores Q e K, o método emprega transformadas de Hadamard (seguindo o QuIP#) e um par adicional de transformações lineares ($Q' = QR^{-T}, K' = KR$) que preserva as pontuações de atenção enquanto reduz a magnitude quadrática média dos tensores projetados.
• Cache de Precisão Mista Consciente de Sink de Atenção: Reconhecendo que as pontuações de atenção se concentram nos tokens iniciais e recentes (Sink de Atenção), o método armazena os primeiros e últimos blocos do cache KV em precisão total, enquanto os blocos intermediários são armazenados em NVFP4 comprimido. Isso garante alta precisão para tokens críticos sem escalar o uso de memória com o comprimento da sequência.

Principais Contribuições

1. Algoritmo ScaleSearch: Uma técnica de busca de escala para quantização de Ponto Flutuante em Blocos que minimiza o erro de quantização por bloco explorando escalas representáveis vizinhas. Ela aproveita especificamente a resolução de mantissa do formato de escala NVFP4 E4M3.
2. ScaleSearchAttention: Um mecanismo de atenção FP4 co-otimizado que permite inferência de ponta a ponta com degradação de precisão próxima de zero. Integra ScaleSearch, processamento de incoerência e cache KV de precisão mista para realizar operações QKT e PV diretamente em NVFP4 nos Tensor Cores.
3. Validação Empírica: Experimentos extensivos demonstrando ganhos consistentes em perplexidade e precisão de benchmarks sobre bases FP4 anteriores (como SageAttention3 e Naive-FP4) em modelos de linguagem e modelos de difusão em grande escala.

Resultados Experimentais

Quantização Pós-Treinamento (PTQ)

• Modelos de Linguagem: No modelo Qwen3-8B, o ScaleSearch melhorou o desempenho em até 15 pontos no benchmark MATH-500 comparado à quantização NVFP4 padrão. Também fechou significativamente a lacuna nos benchmarks GPQA e MMLU.
• Modelos de Difusão: Quando aplicado aos modelos Mochi e CogVideoX usando SageAttention3, o ScaleSearch melhorou as métricas VQA-a, VQA-t e FScore, enquanto igualou as pontuações de alinhamento de precisão total (CLIPSIM, CLIP-T).

Atenção e Perplexidade

• Perplexidade (Wikitext-2): O ScaleSearchAttention reduziu a perplexidade em até 22% para o Llama 3.1 70B comparado à quantização FP4 padrão. Para o Llama 3.1 70B, alcançou uma PPL de 2,6348, uma melhoria significativa sobre o SageAttention3 (3,3899) e o Naive-FP4 (3,4000).
• GPQA Diamond: No modelo Llama 3.1 8B Instruct, o ScaleSearchAttention alcançou 32,32% de precisão, superando o SageAttention3 (26,26%) e igualando a linha de base de precisão total (31,81%).

Eficiência e Sobrecarga

• Sobrecarga de Quantização: O ScaleSearch introduz uma sobrecarga mínima de 1,74x durante a fase de quantização de FP32 para NVFP4 (0,0449 ms vs. 0,0258 ms da linha de base) para um intervalo de busca restrito.
• Taxa de Transferência: Em kernels de atenção, o ScaleSearch alcança 98,3% da taxa de transferência da linha de base (SageAttention3) para atenção não causal em comprimentos de sequência de 32K, demonstrando impacto de tempo de execução negligenciável.
• Latência: A latência de ponta a ponta para geração de texto para vídeo aumentou apenas marginalmente (por exemplo, de 353,40s para 364,68s para o Mochi), confirmando a praticidade do método para cargas de trabalho do mundo real.

Significado

O artigo afirma que o ScaleSearch aborda uma subotimalidade fundamental nas estratégias de escala BFP padrão ao utilizar os bits de mantissa anteriormente subutilizados em formatos de microescala. Ao buscar sistematicamente a escala ótima em vez de confiar em um máximo heurístico, o método reduz significativamente o erro de quantização. A introdução do ScaleSearchAttention demonstra que esses ganhos podem ser estendidos ao mecanismo de atenção intensivo em memória e computação, permitindo inferência de alta taxa de transferência e baixa precisão com perda de precisão próxima de zero. Os autores posicionam este trabalho como um passo rumo ao aproveitamento total das capacidades emergentes de hardware de 4 bits (como o NVIDIA Blackwell) para IA generativa eficiente.