ThriftAttention: Precisão Mista Seletiva para Atenção FP4 de Contexto Longo

Resumo

Algoritmos de atenção eficientes são cruciais para mitigar o custo quadrático da atenção em cargas de trabalho de contexto longo. Trabalhos anteriores utilizam técnicas de quantização em escala de bloco em GPUs Blackwell para mover o cálculo da atenção para precisão de 4 bits, acelerando a inferência. No entanto, essas técnicas resultam em degradação significativa da qualidade em contextos longos. Mostramos que o impacto do erro de quantização na saída é altamente não uniforme e aumenta com a importância de cada interação consulta-chave, concentrando erros funcionalmente relevantes em um pequeno número de blocos de atenção que contêm os tokens mais importantes. Propomos o ThriftAttention, uma variante de atenção de baixa precisão que oferece qualidade de contexto longo próxima ao FP16 com eficiência de inferência FP4. Essa abordagem ocorre em duas etapas. Primeiro, uma heurística seleciona rapidamente um pequeno número de pares de blocos consulta-chave importantes para precisão FP16. Segundo, os blocos selecionados são calculados em FP16 e os blocos restantes em FP4, com ambos os caminhos mesclados via softmax online em uma única saída. Demonstramos, em benchmarks de contexto longo e famílias de modelos, que, ao calcular apenas 5% dos blocos consulta-chave em FP16, o ThriftAttention recupera, em média, 89,1% da lacuna de desempenho entre FP4 e FP16. Mostramos que a vantagem do ThriftAttention cresce com o comprimento da sequência, mitigando a degradação sistemática de qualidade FP4 observada em contextos mais longos. O código está disponível em https://github.com/joesharratt1229/ThriftAttention.

English

Efficient attention algorithms are critical to mitigate the quadratic cost of attention in long-context workloads. Prior work utilises block-scaled quantisation techniques on Blackwell GPUs to move attention computation to 4-bit precision to accelerate inference. However, these techniques result in significant quality degradation in long-context settings. We show that the output impact of quantisation error is highly non-uniform and increases with the importance of each query-key interaction, concentrating functionally relevant error in a small number of attention blocks that contain the most important tokens. We propose ThriftAttention, a low-bit attention variant that delivers near-FP16 long-context quality at FP4 inference efficiency. This approach proceeds in two stages. First, a heuristic rapidly selects a small number of important query-key block pairs for FP16 precision. Second, the selected blocks are computed in FP16 and the remaining blocks in FP4, with both paths merged via online softmax into a single output. We demonstrate across long-context benchmarks and model families that by computing only 5% of query-key blocks in FP16, ThriftAttention recovers on average 89.1% of the FP4-to-FP16 performance gap. We show ThriftAttention's advantage grows with sequence length, mitigating the systematic FP4 quality degradation observed at longer contexts. The code is available at https://github.com/joesharratt1229/ThriftAttention.