FlashAttention 原理与应用——泽微AI/泽微一号上的极限加速

2023-10-10

🚀 性能的瓶颈：Transformer 中的 Attention 机制

自 Transformer 架构成为现代大语言模型（LLM）的基石以来，其核心机制——自注意力（Self-Attention），在带来强大建模能力的同时，也成为了计算效率的巨大瓶颈。

在序列长度为 $L$ 的情况下，Attention 机制的计算复杂度为 $O(L^2)$ ，显存复杂度为 $O(L^2)$ 。这意味着，当上下文窗口变长时（例如数万 Token），计算时间和显存消耗都会以二次方的速度暴涨，严重阻碍了 LLM 的规模扩展。

注意力计算的瓶颈并不完全在于计算量 $O(L^2)$ 本身，而在于内存访问（Memory Access）。

在 GPU 上，计算速度（FLOPS）与内存读写速度（Bandwidth）之间的差距巨大。标准的 Attention 算法需要进行多次访问显存（HBM，高带宽内存），大量的中间结果（如 $S = QK^T$ 矩阵）需要频繁地在 GPU 计算单元（SRAM，片上缓存）和显存之间进行读写。

由于内存访问速度远低于计算速度，GPU 大部分时间都花费在等待数据传输上，而不是进行矩阵乘法运算。这种受限于带宽的运算被称为 “IO-Bound”（输入/输出受限）。

FlashAttention 技术是解决这一瓶颈的革命性创新。它通过对 Attention 机制进行IO 感知型（IO-Aware） 优化，将运算转变为 “Compute-Bound”（计算受限），从而最大化 GPU 的利用率。

FlashAttention 的核心思想是减少对高带宽显存（HBM）的访问次数，最大化利用速度更快的片上缓存（SRAM）。

分块（Tiling）：将输入矩阵 $Q, K, V$ 分成小块（Tiles），每次只将一个 Tile 载入到 SRAM 中进行计算。
减少显存写入：FlashAttention 避免了将中间的 $S = QK^T$ 矩阵和归一化因子 $l$ 写入速度较慢的 HBM。相反，它利用在线 softmax 和重计算（Recomputation） 策略：
- 在线 Softmax：在分块计算过程中，实时更新 Softmax 的归一化因子，无需等待整个矩阵计算完毕。
- 重计算：虽然重计算增加了计算量，但由于计算（Compute）比内存访问（IO）快得多，通过牺牲少量计算资源，换来了巨大的 IO 性能提升。

泽微AI（或 泽微一号）平台将 FlashAttention 深度集成到我们的 LLM 训练和推理软件栈中：

极致性能提升：在 NVIDIA H100/A100 等旗舰 GPU 上，FlashAttention 能够将 Attention 机制的计算速度提升数倍，尤其是在长序列训练中效果显著。
显存效率翻倍：FlashAttention 极大地减少了中间变量的存储需求，使得 LLM 能够处理两倍于传统方法的序列长度，显著缓解了长文本带来的显存压力。
全面覆盖：FlashAttention 被应用于平台的 $haiscale$ 训练框架和 $vLLM$ 推理服务中，为用户提供端到端的加速体验。

FlashAttention 是一项革命性的技术，它通过重新设计 Attention 算法的内存访问模式，成功地将 LLM 的训练和推理从 IO-Bound 瓶颈中解放出来。

泽微AI/泽微一号 平台将 FlashAttention 作为核心优化组件，确保我们的用户在处理大模型和长上下文时，能够享受到最高的计算效率和最低的成本。