ThriftAttention：长上下文 FP4 注意力的选择性混合精度

高效的注意力算法对于减轻长上下文工作负载中注意力的二次成本至关重要。之前的工作利用 Blackwell GPU 上的块级量化技术将注意力计算转移到 4 位精度，以加速推理。然而，这些技术会导致长上下文环境中的质量显着下降。我们表明，量化误差的输出影响是高度不均匀的，并且随着每个查询键交互的重要性而增加，将功能相关的误差集中在包含最重要标记的少量注意块中。我们提出了 ThriftAttention，这是一种低位注意力变体，能够以 FP4 推理效率提供接近 FP16 的长上下文质量。这种方法分两个阶段进行。首先，启发式方法快速选择少量重要的查询密钥块对以实现 FP16 精度。其次，所选块在 FP16 中计算，其余块在 FP4 中计算，两条路径通过在线 softmax 合并为单个输出。我们在长上下文基准测试和模型系列中证明，通过仅计算 FP16 中 5% 的查询键块，ThriftAttention 平均可以弥补 FP4 到 FP16 性能差距的 89.1%。我们证明 ThriftAttention 的优势随着序列长度的增加而增长，减轻了在较长上下文中观察到的系统 FP4 质量下降。该代码可在 https://github.com/joesharratt1229/ThriftAttention 获取。