FlashAttention (Metal 移植版)

本仓库将 FlashAttention 的官方实现移植到 Apple 芯片。它是一组最小化且易于维护的源文件，用于复现 FlashAttention 算法。

文档

仅支持单头注意力，以聚焦于不同注意力算法的核心瓶颈（寄存器压力、并行性）。在基础算法正确实现的前提下，添加诸如块稀疏性之类的自定义功能应相对简单。

所有内容均在运行时进行 JIT 编译。这与之前的实现形成对比，后者依赖于嵌入在 Xcode 14.2 中的可执行文件。

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

为了克服寄存器压力瓶颈，采取了许多大胆的措施。在较大的头维度（例如 256）下，没有任何矩阵块能装入寄存器。连累加器也不行。因此，进行了有意的寄存器溢出，但采用了更优化的方式。注意力算法中增加了一个第三块维度，该维度沿 D 进行分块。注意力矩阵块的宽高比被严重扭曲，以最小化寄存器溢出的带宽成本。例如，在并行化维度上为 16-32，在遍历维度上为 80-128。有一个大型参数文件，它接收 D 维度，并确定哪些操作数可以装入寄存器。然后，它分配一个块大小，以平衡许多相互竞争的瓶颈。

最终结果是在 M1 Max 上实现了一致的每秒 4400 吉指令（83% ALU 利用率），在无限序列长度和无限头维度下。前提是混合精度使用了 BF16 模拟（Metal 的 bfloat 具有符合 IEEE 标准的舍入，在没有硬件 BF16 的旧芯片上这是一个主要开销）。

Raw Data: https://docs.google.com/spreadsheets/d/1Xf4jrJ7e19I32J1IWIekGE9uMFTeZKoOpQ6hlUoh-xY/edit?usp=sharing

量化性能

在 AI 领域，性能最常以每秒十亿次浮点运算（GFLOPS）来报告。该指标反映了一种简化的性能模型，即每条指令都发生在 GEMM 中。随着硬件从早期的 FPU 发展到现代向量处理器，最常见的浮点运算被融合为单条指令。融合乘加（FMA）。当两个 100x100 矩阵相乘时，会发出 100 万条 FMA 指令。为什么我们必须将这条 FMA 视为两条独立的指令？

这个问题与注意力机制相关，因为并非所有浮点运算都是等价的。Softmax 中的指数运算在一个时钟周期内完成，前提是其他大多数指令都流向 FMA 单元。Softmax 中的一些乘法和加法，无法与附近的加法或乘法融合。我们应该将这些视为与 FMA 相同，并假装硬件只是以两倍慢的速度执行 FMA 吗？GEMM 性能模型如何解释我的着色器是否有效地利用了 ALU 硬件，尚不清楚。

我使用 gigainstructions 而非 gigaflops 来理解 shader 的性能表现。它与算法的映射更为直接。例如，一次 GEMM 是 N^3 条 FMA 指令。前向注意力执行两次矩阵乘法，即 2 * D * N^2 条 FMA 指令。反向注意力（根据 Dao-AILab/flash-attention 实现）是 5 * D * N^2 条 FMA 指令。请尝试将此表与 Flash1、Flash2 或 Flash3 论文中的 roofline 模型进行比较。

由于 FP32 原子操作的复杂性，MFA 在反向传播中采用了不同的方法。这种方法具有更高的计算成本。它将反向传播拆分为两个独立的 kernel：dQ 和 dK/dV。下拉菜单显示了伪代码。请将其与 Flash1、Flash2 或 Flash3 论文中的某个算法进行比较。

// Forward
//   for c in 0..<C {
//     load K[c]
//     S = Q * K^T
//     (m, l, P) = softmax(m, l, S * scaleFactor)
//
//     O *= correction
//     load V[c]
//     O += P * V
//   }
//   O /= l
//
//   L = m + logBaseE(l)
//
// Backward Query
//   D = dO * O
//
//   for c in 0..<C {
//     load K[c]
//     S = Q * K^T
//     P = exp(S - L)
//
//     load V[c]
//     dP = dO * V^T
//     dS = P * (dP - D) * scaleFactor
//
//     load K[c]
//     dQ += dS * K
//   }
//
// Backward Key-Value
//   for r in 0..<R {
//     load Q[r]
//     load L[r]
//     S^T = K * Q^T
//     P^T = exp(S^T - L)
//
//     load dO[r]
//     dV += P^T * dO
//
//     load dO[r]
//     load D[r]
//     dP^T = V * dO^T
//     dS^T = P^T * (dP^T - D) * scaleFactor
//
//     load Q[r]
//     dK += dS^T * Q
//   }

性能通过计算计算工作量，然后除以秒来衡量。最终结果是“每秒十亿条指令”。接下来，我们需要一个屋顶线模型。下表显示了 GINSTRS 的屋顶线，计算为 GFLOPS 的一半。ALU 利用率是（实际每秒十亿条指令）/（预期每秒十亿条指令）。例如，M1 Max 在混合精度下通常能达到 80% 的 ALU 利用率。

该模型存在局限性。在 M3 代的小头维度下，该模型会失效。不同的计算单元可能同时被利用，使得表观利用率超过 100%。在大多数情况下，基准测试提供了关于还有多少性能未被充分利用的准确模型。

var operations: Int
switch benchmarkedKernel {
case .forward:
  operations = 2 * headDimension + 5
case .backwardQuery:
  operations = 3 * headDimension + 5
case .backwardKeyValue:
  operations = 4 * headDimension + 5
}
operations *= (sequenceDimension * sequenceDimension)
operations *= dispatchCount

// Divide the work by the latency, resulting in throughput.
let instrs = Double(operations) / Double(latencySeconds)
let ginstrs = Int(instrs / 1e9)

Metal 移植版与官方 FlashAttention 仓库相比表现如何？假设我采用了“atomic dQ”算法并实现了 100% 的性能。然后，切换到实际的 MFA 仓库，发现模型训练速度慢了 4 倍。那将是官方仓库 roofline 的 25%。要得到这个百分比，请将所有三个 kernel 的平均 ALU 利用率乘以 7 / 9。针对 Apple 硬件的统计使用了更细致的模型，但这就是其要点。

为了计算 Nvidia 硬件的利用率，我使用了 FP16/BF16 ALU 的 GFLOPS。我将论文中每个图表中的最高 GFLOPS 除以 312000（A100 SXM）、989000（H100 SXM）。请注意，对于较大的 head dimensions 和寄存器密集型 kernel（backward pass），未报告基准测试。我确认他们在无限 head dimensions 下并未解决寄存器压力问题。例如，累加器始终保存在寄存器中。在撰写本文时，我尚未看到 D=256 backward gradient 以正确结果执行的具体证据。

GFLOPS

计算利用率

并排对比

尽管执行了更多的计算，Apple 硬件训练 transformers 的速度比 Nvidia 硬件完成相同工作更快。对不同 GPU 之间的尺寸差异进行归一化。仅关注 GPU 的利用效率。

也许主仓库应该尝试一种避免 FP32 原子操作并在寄存器无法容纳于 GPU 核心时故意溢出寄存器的算法。这似乎不太可能，因为他们硬编码了对可能问题尺寸的一个小子集的支持。其动机似乎是支持最常见的模型，其中 D 是 2 的幂，且小于 128。对于其他情况，用户需要依赖替代的回退实现（例如 MFA 仓库），这些实现可能使用完全不同的底层算法。

用法

设置工作流

在 macOS 上，下载 Swift 包并使用 -Xswiftc -Ounchecked 进行编译。此编译器选项对于性能敏感的 CPU 代码是必需的。不能使用 Release 模式，因为它会强制每次只要有单一更改，就从头重新编译整个代码库。在 Finder 中导航到 Git 仓库并双击 Package.swift。应该会弹出一个 Xcode 窗口。在左侧，应该有一个文件层级结构。如果你无法理清该层级结构，则说明出现了问题。

git clone https://github.com/philipturner/metal-flash-attention
swift build -Xswiftc -Ounchecked # Does it even compile?
swift test -Xswiftc -Ounchecked # Does the test suite finish in ~10 seconds?

或者，使用 SwiftUI 模板创建一个新的 Xcode 项目。用一个返回 String 的函数调用覆盖 "Hello, world!" 字符串。该函数将执行您选择的脚本，然后调用 exit(0)，以便应用程序在将任何内容渲染到屏幕之前崩溃。您将使用 Xcode 控制台中的输出作为关于您代码的反馈。此工作流兼容 macOS 和 iOS。

通过 Project > 您的项目名称 > Build Settings > Swift Compiler - Code Generation > Optimization Level 添加 -Xswiftc -Ounchecked 选项。表格的第二列列出了您的项目名称。在下拉菜单中点击 Other，并在出现的面板中输入 -Ounchecked。接下来，将此仓库添加为 Swift 包依赖项。查看 Tests/FlashAttention 下的一些测试。将这些测试中某一个的原始源代码复制到您的项目中。从上一段中的函数调用该测试。检查它在控制台上显示的内容。

要修改 Metal 代码生成（例如，添加多头或掩码支持），请将原始 Swift 代码复制到您的 Xcode 项目中。可以在单独的文件夹中使用 git clone，或者在 GitHub 上以 ZIP 形式下载原始文件。还有一种方法可以链接到 metal-flash-attention 的 fork 并将您的更改自动保存到云端，但这设置起来更困难。移除上一段中的 Swift 包依赖项。重新运行您选择的测试。它是否编译并在控制台上显示某些内容？

编辑源代码

在这些文件夹之一中定位一个多行字符串字面量：

Sources/FlashAttention/Attention/AttentionKernel
Sources/FlashAttention/GEMM/GEMMKernel

向其中一处添加随机文本。再次编译并运行该项目。应该会出现严重错误。例如，Metal 编译器可能会抛出错误。如果未发生此情况，请尝试在别处弄乱另一行代码。如果测试仍然通过，则 Xcode 未注册你的更改。

继续编写 block sparsity 或类似内容。获取关于代码是否完全可用、是否运行快速、是否在所有问题规模下都运行快速的反馈。将原始源代码集成到你的应用中，或将其转换为另一种编程语言。