ThriftAttention: Precisión mixta selectiva para atención FP4 de contexto largo

Resumen

Los algoritmos de atención eficientes son críticos para mitigar el costo cuadrático de la atención en cargas de trabajo de contexto largo. Trabajos previos utilizan técnicas de cuantización de bloques escalados en GPUs Blackwell para mover el cómputo de atención a precisión de 4 bits y acelerar la inferencia. Sin embargo, estas técnicas resultan en una degradación significativa de la calidad en entornos de contexto largo. Mostramos que el impacto del error de cuantización en la salida es altamente no uniforme y aumenta con la importancia de cada interacción consulta-clave, concentrando el error funcionalmente relevante en un pequeño número de bloques de atención que contienen los tokens más importantes. Proponemos ThriftAttention, una variante de atención de baja precisión que ofrece una calidad de contexto largo cercana a FP16 con la eficiencia de inferencia de FP4. Este enfoque se desarrolla en dos etapas. Primero, una heurística selecciona rápidamente un pequeño número de pares de bloques consulta-clave importantes para precisión FP16. Segundo, los bloques seleccionados se computan en FP16 y los restantes en FP4, fusionándose ambas trayectorias mediante softmax en línea en una única salida. Demostramos en varios benchmarks de contexto largo y familias de modelos que, al computar solo el 5% de los bloques consulta-clave en FP16, ThriftAttention recupera en promedio el 89.1% de la brecha de rendimiento entre FP4 y FP16. Mostramos que la ventaja de ThriftAttention crece con la longitud de la secuencia, mitigando la degradación sistemática de calidad de FP4 observada en contextos más largos. El código está disponible en 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.