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 上实现了稳定的每秒 44000 亿次指令执行速度（ALU 利用率为 83%），并且支持无限序列长度和无限头维度。前提是混合精度计算中使用 BF16 模拟运算（Metal 的 bfloat 实现遵循 IEEE 规范的舍入规则，这在不支持硬件 BF16 的旧款芯片上会造成较大的性能开销）。

原始数据：https://docs.google.com/spreadsheets/d/1Xf4jrJ7e19I32J1IWIekGE9uMFTeZKoOpQ6hlUoh-xY/edit?usp=sharing

量化性能

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

这个问题与注意力机制密切相关，因为并非所有的浮点运算都是一样的。在 softmax 过程中的指数运算可以在一个时钟周期内完成，前提是大多数其他指令都被送往 FMA 单元。然而，softmax 过程中的一些乘法和加法操作无法与附近的加法或乘法操作融合。我们是否应该把这些操作也当作 FMA 来对待，并假装硬件只是以两倍的速度执行 FMA 呢？目前尚不清楚 GEMM 性能模型如何解释我的着色器是否有效地利用了 ALU 硬件。

因此，我改用每秒十亿次指令（GINSTRS）来衡量着色器的性能表现，因为它更直接地对应于算法本身。例如，一次 GEMM 计算需要 N^3 条 FMA 指令。前向注意力机制执行两次矩阵乘法，即 2 * D * N^2 条 FMA 指令。而后向注意力机制（由 Dao-AILab/flash-attention 实现）则需要 5 * D * N^2 条 FMA 指令。不妨将这张表与 Flash1、Flash2 或 Flash3 论文中的屋顶模型进行对比。

由于 FP32 原子操作的复杂性，MFA 在后向传播中采用了不同的方法，这导致更高的计算开销。它将后向传播拆分为两个独立的内核：dQ 和 dK/dV。下拉菜单中展示了伪代码，可以将其与 Flash1、Flash2 或 Flash3 论文中的一种算法进行比较。

// 前向
//   for c in 0..<C {
//     加载 K[c]
//     S = Q * K^T
//     (m, l, P) = softmax(m, l, S * 缩放因子)
//
//     O *= 校正系数
//     加载 V[c]
//     O += P * V
//   }
//   O /= l
//
//   L = m + logBaseE(l)
//
// 后向查询
//   D = dO * O
//
//   for c in 0..<C {
//     加载 K[c]
//     S = Q * K^T
//     P = exp(S - L)
//
//     加载 V[c]
//     dP = dO * V^T
//     dS = P * (dP - D) * 缩放因子
//
//     加载 K[c]
//     dQ += dS * K
//   }
//
// 后向键值
//   for r in 0..<R {
//     加载 Q[r]
//     加载 L[r]
//     S^T = K * Q^T
//     P^T = exp(S^T - L)
//
//     加载 dO[r]
//     dV += P^T * dO
//
//     加载 dO[r]
//     加载 D[r]
//     dP^T = V * dO^T
//     dS^T = P^T * (dP^T - D) * 缩放因子
//
//     加载 Q[r]
//     dK += dS^T * Q
//   }

性能是通过计算总的计算工作量，再除以时间（秒）来衡量的。最终结果就是“每秒十亿次指令”。接下来，我们需要一个屋顶模型。下表显示了基于 GINSTRS 的屋顶线，其数值为 GFLOPS 的一半。ALU 利用率等于实际每秒十亿次指令数除以理论上的每秒十亿次指令数。例如，M1 Max 在混合精度下通常能达到 80% 的 ALU 利用率。

不过，这种模型也有局限性。在 M3 芯片上，当头维度较小时，该模型就会失效。不同的计算单元可能会同时被利用，从而导致 apparent 利用率超过 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

// 将总工作量除以延迟时间，得到吞吐量。
let instrs = Double(operations) / Double(latencySeconds)
let ginstrs = Int(instrs / 1e9)

Metal 版本与官方 FlashAttention 仓库相比，性能如何呢？假设我采用了“原子 dQ”算法并达到了 100% 的性能，随后切换到实际的 MFA 仓库，却发现模型训练速度慢了四倍。这意味着我只发挥了官方仓库屋顶线的 25%。要计算这个百分比，只需将三个内核的平均 ALU 利用率乘以 7 / 9 即可。虽然针对 Apple 硬件的统计使用了更为精细的模型，但大致思路就是这样。

为了计算 Nvidia 硬件的利用率，我使用了 FP16/BF16 ALU 的 GFLOPS 数据。我将论文中每张图给出的最高 GFLOPS 数值分别除以 A100 SXM 的 312000 和 H100 SXM 的 989000。需要注意的是，对于较大的头维度和寄存器密集型的内核（如后向传播），并未报告任何基准测试结果。我确认他们并未解决无限头维度下的寄存器压力问题。例如，累加器始终保存在寄存器中。截至撰写本文时，我尚未看到 D=256 的后向梯度能够正确执行的具体证据。

GFLOPS

计算利用率

并排对比

尽管发出的计算量更多，但苹果硬件在训练 Transformer 模型时的速度比执行相同任务的英伟达硬件更快。这一比较已根据不同 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 # 是否能成功编译？
swift test -Xswiftc -Ounchecked # 测试套件是否能在约 10 秒内完成？

或者，可以使用 SwiftUI 模板创建一个新的 Xcode 项目，将默认的 "Hello, world!" 字符串替换为调用一个返回 String 的函数。该函数将执行你选择的脚本，然后调用 exit(0)，使应用在渲染任何内容之前崩溃。你可以通过 Xcode 控制台输出来获取代码反馈。此工作流程同时适用于 macOS 和 iOS。

通过 项目 > 你的项目名称 > 构建设置 > Swift 编译器 - 代码生成 > 优化级别 添加 -Xswiftc -Ounchecked 选项。表格的第二列显示你的项目名称。点击下拉菜单中的 其他,并在出现的面板中输入 -Ounchecked。接下来，将此仓库添加为 Swift 包依赖项。浏览一下 Tests/FlashAttention 下的一些测试，将其中一个测试的原始源代码复制到你的项目中，并从上一段提到的函数中调用该测试，检查控制台输出的内容。

要修改 Metal 代码生成（例如添加多头或掩码支持），可将原始 Swift 代码复制到你的 Xcode 项目中。你可以选择在单独的文件夹中使用 git clone，或直接从 GitHub 下载 ZIP 格式的原始文件。此外，还可以链接到你自己的 metal-flash-attention 分支，并自动将更改保存到云端，但这种方式设置起来较为复杂。移除上一段中的 Swift 包依赖项，然后重新运行你选择的测试。它能否成功编译并在控制台显示内容？

编辑源代码

在以下任一文件夹中找到一个多行字符串字面量：

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

向其中任意一个添加随机文本，然后再次编译并运行项目。此时应该会出现严重错误，例如 Metal 编译器抛出异常。如果未发生这种情况，请尝试修改其他地方的代码行。若测试仍能通过，则说明 Xcode 并未识别你的更改。

接下来可以继续开发 块稀疏性 或其他功能，以验证代码是否有效、运行速度如何以及在各种问题规模下是否都能保持高效。最后，将原始源代码集成到你的应用中，或将其转换为其他编程语言。

metal-flash-attention 快速上手指南

metal-flash-attention 是官方 FlashAttention 算法在 Apple Silicon (M1/M2/M3/M4 等) 芯片上的 Metal 移植版本。它专注于核心性能瓶颈优化，支持即时编译 (JIT)，并在大维度头 (Head Dimension) 下通过独特的寄存器溢出策略实现了极高的 ALU 利用率。

环境准备

• 操作系统: macOS (推荐最新版本)
• 硬件: Apple Silicon 芯片 (M1, M2, M3, M4 系列)
• 开发工具:
  • Xcode (需包含命令行工具)
  • Swift 5.9 或更高版本
• 依赖: 无额外第三方依赖，仅需标准 Swift 包管理器。

安装步骤

1. 克隆仓库 打开终端，执行以下命令下载源码：

   git clone https://github.com/philipturner/metal-flash-attention
   cd metal-flash-attention
   

2. 编译项目 为了获得最佳性能，必须使用 -Ounchecked 优化标志进行编译。该标志针对性能敏感的 CPU 代码进行了优化，但会跳过部分安全检查（仅限受信任的代码库）。

   swift build -Xswiftc -Ounchecked
   

3. 运行测试（可选） 验证安装是否成功及核心算法是否正确：

   swift test -Xswiftc -Ounchecked
   

4. Xcode 集成（可选） 如果你偏好使用 Xcode 进行开发或调试：

   • 在 Finder 中进入项目目录。
   • 双击 Package.swift 文件。
   • Xcode 将自动打开并解析项目结构。

基本使用

该项目主要作为底层内核库供其他 Swift 项目调用。以下是一个最简单的集成示例，展示如何在你的 Swift 项目中引入并使用它。

1. 添加依赖

在你的项目根目录下的 Package.swift 文件中，添加 metal-flash-attention 作为依赖：

// Package.swift
dependencies: [
    .package(url: "https://github.com/philipturner/metal-flash-attention.git", branch: "main")
],
targets: [
    .executableTarget(
        name: "YourApp",
        dependencies: ["metal-flash-attention"]
    )
]

2. 代码示例

在你的 Swift 源代码中导入模块并调用注意力机制。由于该库采用 JIT 编译，运行时会根据当前的序列长度和头维度动态生成优化的 Metal 着色器。

import MetalFlashAttention
import Metal

// 初始化 Metal 设备
let device = MTLCreateSystemDefaultDevice()!
let commandQueue = device.makeCommandQueue()!

// 配置参数 (示例)
let headDimension = 128
let sequenceLength = 2048
let batchSize = 1

// 创建输入张量 (此处仅为伪代码示意，实际需分配 Metal Buffer)
// let queryBuffer = ...
// let keyBuffer = ...
// let valueBuffer = ...

// 执行前向传播 (Forward Pass)
// 注意：具体 API 调用请参考库中暴露的最新接口，核心逻辑如下：
/*
let attention = FlashAttention(device: device)
let outputBuffer = attention.forward(
    query: queryBuffer,
    key: keyBuffer,
    value: valueBuffer,
    headDim: headDimension,
    seqLen: sequenceLength
)
*/

print("FlashAttention kernel compiled and ready on Apple Silicon.")

提示: 该库针对无限序列长度和无限头维度进行了优化，特别是在 D=256 等大维度场景下，通过特殊的分块策略避免了寄存器压力瓶颈，性能表现优于传统实现。