ThriftAttention:用于长上下文FP4注意力的选择性混合精度
阅读原文· arxiv.org针对长上下文场景下FP4量化推理的质量下降问题,ThriftAttention提出了一种低比特注意力变体。该方法首先通过启发式策略快速识别出少量关键的query-key块对,对其使用FP16精度计算,其余块使用FP4计算。最后,两路计算结果通过在线softmax合并。实验表明,仅对5%的块使用FP16计算,该方法就能平均恢复FP4与FP16间89.1%的性能差距,且优势随序列长度增加而更加显著。
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.