ThriftAttention : Précision mixte sélective pour l'attention FP4 en contexte long

Résumé

Les algorithmes d'attention efficaces sont essentiels pour atténuer le coût quadratique de l'attention dans les charges de travail à contexte long. Des travaux antérieurs utilisent des techniques de quantification à échelle de blocs sur les GPU Blackwell pour déplacer le calcul de l'attention vers une précision 4 bits afin d'accélérer l'inférence. Cependant, ces techniques entraînent une dégradation significative de la qualité dans les contextes longs. Nous montrons que l'impact de l'erreur de quantification est hautement non uniforme et augmente avec l'importance de chaque interaction requête-clé, concentrant l'erreur fonctionnellement pertinente dans un petit nombre de blocs d'attention contenant les jetons les plus importants. Nous proposons ThriftAttention, une variante d'attention à faible précision qui offre une qualité de contexte long proche du FP16 avec une efficacité d'inférence FP4. Cette approche se déroule en deux étapes. Premièrement, une heuristique sélectionne rapidement un petit nombre de paires de blocs requête-clé importants pour la précision FP16. Deuxièmement, les blocs sélectionnés sont calculés en FP16 et les blocs restants en FP4, les deux chemins étant fusionnés via un softmax en ligne en une seule sortie. Nous démontrons sur des bancs d'essai à contexte long et des familles de modèles qu'en calculant seulement 5 % des paires de blocs requête-clé en FP16, ThriftAttention comble en moyenne 89,1 % de l'écart de performance entre FP4 et FP16. Nous montrons que l'avantage de ThriftAttention augmente avec la longueur de la séquence, atténuant la dégradation systématique de la qualité FP4 observée dans les contextes plus longs. Le code est disponible à l'adresse 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.