首次实现支持INT8量化的FlashAttention，显著加速GPU推理。
INT-FlashAttention: Enabling Flash Attention for INT8 Quantization
- 将INT8量化融入FlashAttention计算流程，保持原架构兼容性。
- 在Ampere GPU上实现72%加速，量化误差降低82%。
- 适用于INT4/INT8等格式，适合大模型高效部署场景。
作为大语言模型的核心组件，自注意力模块随序列长度呈二次方增长的时间与内存复杂度。FlashAttention通过利用GPU内存层级结构加速计算并降低内存占用。将FlashAttention与量化方法结合是极具前景的研究方向。本文提出INT-FlashAttention，首个兼容FlashAttention前向流程的INT8量化架构，在Ampere GPU上显著提升推理速度。我们实现了全INT8激活与通用矩阵乘法（GEMM）核的原型系统，使该注意力算子成为首个支持全INT8输入的算子。作为通用的令牌级后训练量化框架，INT-FlashAttention还兼容如INT4等其他数据格式。实验结果表明，相比使用FP16和FP8数据格式的标准FlashAttention，INT-FlashAttention实现72%的推理加速与82%的量化误差降低。
原文摘要 · Abstract (English)
As the foundation of large language models (LLMs), self-attention module faces the challenge of quadratic time and memory complexity with respect to sequence length. FlashAttention accelerates attention computation and reduces its memory usage by leveraging the GPU memory hierarchy. A promising research direction is to integrate FlashAttention with quantization methods. This paper introduces INT-FlashAttention, the first INT8 quantization architecture compatible with the forward workflow of FlashAttention, which significantly improves the inference speed of FlashAttention on Ampere GPUs. We implement our INT-FlashAttention prototype with fully INT8 activations and general matrix-multiplication (GEMM) kernels, making it the first attention operator with fully INT8 input. As a general token-level post-training quantization framework, INT-FlashAttention is also compatible with other data formats like INT4, etc. Experimental results show INT-FlashAttention achieves 72% faster inference speed and 82% smaller quantization error compared to standard FlashAttention with FP16 and FP8 data format.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。