摘要
作为大型语言模型 (LLM) 的基础,自注意力模块面临着时间和内存复杂度与序列长度成二次方增长的挑战。FlashAttention 通过利用 GPU 内存层次结构来加速注意力计算并减少其内存使用。一个很有前景的研究方向是将 FlashAttention 与量化方法相结合。本文介绍了 INT-FlashAttention,这是第一个与 FlashAttention 正向工作流程兼容的 INT8 量化架构,它显著提高了 FlashAttention 在 Ampere GPU 上的推理速度。我们使用全 INT8 激活和通用矩阵乘法 (GEMM) 内核实现了 INT-FlashAttention 原型,使其成为第一个具有全 INT8 输入的注意力算子。作为一种通用的令牌级训练后量化框架,INT-FlashAttention 也兼容其他数据格式,例如 INT4 等。实验结果表明,与使用 FP16 和 FP8 数据格式的标准 FlashAttention 相比,INT-FlashAttention 的推理速度提高了 72%,量化误差减少了 82%。