ThriftAttention: Выборочная смешанная точность для длинноконтекстного FP4-внимания

Аннотация

Эффективные алгоритмы внимания критически важны для снижения квадратичной стоимости внимания в задачах с длинным контекстом. Предшествующие работы применяют методы блочно-масштабированной квантизации на графических процессорах Blackwell для перевода вычислений внимания в 4-битную точность с целью ускорения инференса. Однако такие методы приводят к значительному ухудшению качества в условиях длинного контекста. Мы показываем, что влияние ошибки квантизации на выходные данные крайне неравномерно и возрастает с увеличением важности взаимодействия «запрос-ключ», концентрируя функционально значимые ошибки в небольшом числе блоков внимания, содержащих наиболее важные токены. Мы предлагаем ThriftAttention — вариант низкобитного внимания, обеспечивающий качество, близкое к FP16, при эффективности инференса на уровне FP4. Этот подход состоит из двух этапов. Сначала эвристика быстро выбирает небольшое количество важных пар блоков «запрос-ключ» для использования точности FP16. Затем выбранные блоки вычисляются в FP16, а остальные — в FP4, причём оба пути объединяются через онлайн-softmax в единый выходной результат. На примере эталонов длинного контекста и различных семейств моделей мы демонстрируем, что, вычисляя лишь 5% пар блоков «запрос-ключ» в FP16, ThriftAttention в среднем восстанавливает 89,1% разрыва в производительности между FP4 и FP16. Мы показываем, что преимущество ThriftAttention возрастает с увеличением длины последовательности, компенсируя систематическое ухудшение качества FP4, наблюдаемое на более длинных контекстах. Код доступен по адресу 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.