注意力机制是驱动大型语言模型（LLMs）的变压器架构中的基本组成部分。然而，随着LLMs处理的输入序列不断增长，注意力的计算成本逐渐成为一个显著的瓶颈。
为了解决这一问题，Colfax Research、Meta、Nvidia、乔治亚州理工学院、普林斯顿大学和Together AI的合作团队推出了FlashAttention-3。这项前沿技术在Nvidia Hopper GPU（H100和H800）上大幅加速了注意力计算。
理解LLMs中的注意力计算挑战
变压器模型的注意力机制能够有效评估输入序列中各个令牌之间的关系。尽管效果显著，但这种机制的计算量巨大。随着输入序列长度的增加，注意力计算的成本呈平方增长，这会导致LLMs扩展能力的显著瓶颈。
现代GPU主要针对矩阵乘法操作进行了优化，而其他操作（如指数运算）则速度较慢，进一步加剧了性能问题。注意力计算结合了矩阵乘法和其他复杂函数（如用于标准化权重的softmax函数），这使得计算量更高，限制了性能。有效调度工作负载以避免操作冲突并充分利用内存资源至关重要。
利用FlashAttention增强硬件资源利用率
FlashAttention于2022年推出，通过减少GPU上的高带宽内存（HBM）与静态随机存取存储器（SRAM）之间的内存传输，改善了注意力计算的效率。通过将注意力权重分成更小的“块”进行处理，FlashAttention提高了效率，使得LLMs能够将上下文窗口从数千个扩展到数百万个令牌。
然而，随着硬件性能的提升，对进一步优化的需求也随之增加。2023年推出的FlashAttention-2优化了GPU资源的使用，在Nvidia A100 GPU上达到了70%的最大性能，但在H100上仅发挥了35%的能力。
FlashAttention-3的创新
FlashAttention-3充分利用了新的Nvidia Hopper GPU功能提升性能，包括矩阵乘法的吞吐量改善和更快速的数据传输。这使得低精度操作的效率得到了提升。
FlashAttention-3的关键创新包括：
1. 优化调度：作业组织旨在最大化计算与数据移动的重叠，减少GPU空闲时间。
2. 无缝操作交错：通过结合矩阵乘法与softmax操作，FlashAttention-3减少潜在的瓶颈。
3. 增强量化模型性能：优化的操作设计确保即使在使用低位表示以减小模型规模的情况下也能实现更快速、更准确的计算，解决了量化通常带来的精度折衷问题。
研究表明，FlashAttention-3可以利用H100 GPU最大性能的75%，为训练和部署LLMs提供了1.5到2倍的速度提升。
FlashAttention-3的优势
FlashAttention-3加快的注意力计算对LLMs的开发和应用产生深远影响：
- 加速训练：提高的效率可以显著缩短训练时间，使研究人员能够探索更大的模型和数据集。
- 扩展上下文窗口：通过高效处理更长的序列，FlashAttention-3为LLM应用开辟了新领域，如长篇文档理解和多示例上下文学习。
- 成本效益：更高的GPU利用率可以减少LLM操作所需的加速器数量，从而降低生产成本。
FlashAttention-3已在宽松许可下开源，并计划集成到流行的深度学习库，如PyTorch和Hugging Face Transformers。这将帮助研究人员和开发者利用FlashAttention-3的进步。Together AI团队在博客中指出：“设计利用硬件特性的算法可以显著提高效率，解锁新的模型能力。”他们期待进一步优化LLM推理，并将其技术应用于各种硬件架构。
排行榜
