LLM2D
INT-FlashAttention:支持 INT8 量化的 FlashAttention
INT-FlashAttention: Enabling Flash Attention for INT8 Quantization
作者: Shimao Chen, Zirui Liu, Zhiying Wu, Ce Zheng, Peizhuang Cong, Zihan Jiang, Lei Su, Tong Yang
发布日期: 9/26/2024
arXiv ID: oai:arXiv.org:2409.16997v1

摘要

作为大型语言模型 (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%。