FlashAttention（Metal 移植版）

警告：代码尚未完成。目前仅为"最小可行产品"。即完整复现 Flash2 论文，具有合理的性能、大幅节省内存且无 bug。熟悉该领域的人可以通过查看代码获得有意义的见解。剩余的性能问题将在 2024 年 7-8 月解决。

更新：基本的系统级设计缺陷已被纠正。通过有意的寄存器溢出，在无限大的头部维度下，FP32 前向传播已达到超过 50% 的 ALU 利用率。溢出开销通过矩形注意力矩阵块得到了分摊。

本仓库将FlashAttention的官方实现移植到 Apple 芯片上。它是一组最小化、可维护的源文件，重现了 FlashAttention 算法。

源代码树包含了一个定制版的统一 GEMM 内核，这是一个独立脚本，用于在矩阵乘法中达到峰值性能。GEMM 内核与 FlashAttention 内核是不同的。修改后的 GEMM 内核有几个用途，比如测试朴素注意力算法。与 GEMM 相关的代码及其维护不在 metal-flash-attention 的范围内。

重要信息

支持 macOS 和 iOS。可以在 iPad 上的 Swift Playground 中编译。

所有内容都在运行时即时编译。这与之前的实现不同，后者依赖于嵌入 Xcode 14.2 中的可执行文件。

所有计算和存储都使用完整的 32 位精度。

广泛使用异步复制，主要是为了简化代码设计。即使在 M3 上，有时也能帮助提高性能。

仅支持单头注意力，以专注于不同注意力算法的核心瓶颈（算术密度、并行性）。

对 FlashAttention 的修改

反向传播使用更少的内存。官方实现为原子操作和部分和分配了临时空间。Apple 硬件缺乏原生 FP32 原子操作（metal::atomic<float> 是模拟的）。在尝试绕过硬件支持不足的过程中，FlashAttention-2 反向内核中的带宽和并行化瓶颈被揭示出来。设计了一种替代的反向传播方法，计算成本更高（7 个 GEMM 而不是 5 个 GEMM）。它在注意力矩阵的行和列维度上都达到了 100% 的并行化效率。最重要的是，它更容易编码和维护。

反向传播的替代方案考虑

待办：更新此文档，或将其删除。由于时间限制，我不会探索使用 5 个 GEMM 的反向注意力。

一个有趣的发现是最佳反向内核。如果序列长度足够小，注意力矩阵（dS）应该在 RAM 中明确具体化。批次中的每个推理或头部（通常总共 8 个）可以顺序计算。任何时候只需在内存中保留一个 $O(n^2)$ 注意力矩阵。这个注意力矩阵将消耗与 Flash2 反向中 dQ 累积的部分和大致相同的内存。

最小计算成本、最大并行的反向传播是：

• 在芯片上累积 dV，同时将 dS 写入内存
  • 寄存器压力可能会阻止 dK 的融合累积
• 在第二次传递中，使用 GEMM 生成 dK 和 dQ
  • 每个 GEMM 沿不同方向遍历注意力矩阵，在该方向上不需要处理单元之间的同步

如果头部大小（D）大于或等于 32*，这应该比原子累积 dQ 的内核使用更少的 HBM 带宽。它需要 5 个 GEMM 的计算成本，就像 Flash2 一样。这种变体可以将内存限制在 $O($ 处理器数量 $)$ 而不是 $O(n^2)$。要做到这一点而不降低性能，需要了解特定机器的参数（GPU 核心数量、缓存大小等）。在生产环境中运行该算法需要大量针对特定硬件的调优。

*在 Apple 芯片上，由于寄存器压力限制，最佳块大小为 32x32。

初步数据支持明确具体化 $O(n^2)$ 矩阵可以提高性能的预测。数据质量受到寄存器压力瓶颈和将注意力矩阵压缩为 BF16 时性能问题的限制。

待办事项列表

这里列出了优先事项。按照最大到最小的问题排序。

剩余的寄存器压力优化：

• 混合精度。
  • 研究在溢出到 RAM 时频繁截断累加器的影响。
• 块大小、精度等选择启发式，在广泛的目标问题范围内提供稳健的性能。

可移植性：

• 测试注意力矩阵不是正方形的问题。
• 支持一些关键特性：
  • 安全处理未对齐矩阵边缘的密集掩码（"偏置"）。
  • 用于处理多头注意力的前导维度。

文档：

• 提供构建准确基准测试的示例代码。
  • 评估推测性优化或其他代码更改对性能的影响。
  • 解释如何计算理论性能上限。
  • 由于时间限制，可能无法实现。
• 解释移除稀疏块掩码的原因，注意其最近被添加到 PyTorch 中。
  • 解释如何实现稀疏块掩码：两级稀疏性层次结构。将注意力掩码分解为块对齐、零填充的块，通过指针使用一级间接访问。
  • PyTorch 对 FlexAttention 的初始实现将问题大小限制为能被 128 整除的情况。MFA 设计时就考虑到了对奇数和奇怪块大小的一流支持。块大小不能整除序列长度也没关系。
  • 解释 128x128 + 1024x1024 不是一个可行的块大小对。小块大小应等于原生 GEMM 块大小（例如 32x64、16x128）。原生 GEMM 块大小能整除 960x960，后者有更多质因数。
    • 960 可被 64 整除。
    • 960 不能被 128 整除。相反，在所有块大小选择启发式中，用 120 替换 128。120 仍然可以被 8 整除（SIMD 组矩阵正方形边长）。
  • 仓库所有者鼓励人们尝试修改此代码，以克服注意力的二次计算成本缩放。由于时间限制，作者本人无法尝试这一点。