arXiv:2605.23081cs.LG2026-05被引 2

只用5%的高精度计算,就让4位精度注意力接近16位精度效果。

ThriftAttention: Selective Mixed Precision for Long-Context FP4 Attention

论文配图:ThriftAttention: Selective Mixed Precision for Long-Context FP4 Attention
图 1 · 摘自论文原文
  • 选关键查询-键块用FP16,其余用FP4,动态混合精度
  • 仅5%块用高精度,平均恢复89.1%的精度损失
  • 适合长序列任务,越长越有优势,适合部署优化

高效注意力算法对缓解长上下文场景下的二次复杂度至关重要。先前工作在Blackwell GPU上使用分块量化技术将注意力计算降至4位精度以加速推理,但导致长上下文质量显著下降。我们发现量化误差的影响高度非均匀,且随查询-键交互重要性增加而加剧,功能相关误差集中于少数包含最重要标记的注意力块中。为此提出ThriftAttention,一种低比特注意力变体,在保持FP4推理效率的同时实现接近FP16的长上下文性能。该方法分两阶段:首先通过启发式快速筛选出少量重要查询-键块对进行FP16计算;其次,将选定块的FP16结果与其余块的FP4结果通过在线Softmax合并为单一输出。在多个长上下文基准和模型族上验证,仅需对5%的查询-键块采用FP16,ThriftAttention即可平均恢复89.1%的FP4到FP16性能差距。其优势随序列长度增长,有效缓解了长上下文下普遍存在的系统性精度下降问题。代码已开源:https://github.com/joesharratt1229/ThriftAttention。

原文摘要 · Abstract (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.

注意力机制混合精度长序列

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。