更快、更灵活的 FlashAttention-4：我们如何为 LLM 推理优化内核

去年 FlashAttention-4 内核源码发布时，我们深入研究了它，并以极其详尽的细节分享了我们对内核工作原理的发现。现在，你可以通过阅读这篇来自官方的文章来确认我们推断出的高层结构。

在随后的几个月里，我们对该内核做了一系列贡献，使其更适合大语言模型推理，特别是解码密集型工作负载。与预训练工作负载不同，LLM 推理工作负载通常由受内存带宽限制的"解码"或"token 生成"阶段主导（下图浅蓝色部分）。

推理工作负载通常也更加多变——批次大小和序列长度变得不均匀；键和值必须从缓存中检索（大多数情况下）。

这需要新的内核代码，而且这些代码必须快速："性能就是产品"。

在深入细节之前，先为更广泛的读者提供一些要点。

关于底层编程的高层要点

我们为将内核扩展到我们想要运行的推理工作负载所做的更改可以大致分为两类：

• 调整并行策略，即每个线程块的查询 tile 数量，以及从查询并行切换到键/值并行，以及
• 支持不规则全局内存访问，即使用 cp.async 加载来替代使用 Tensor Memory Accelerator (TMA) 的 cp.async.bulk 加载。

这两类分别由下图表示，下文将详细解释。

我们的优化之一是将"split KV"技术移植到 FA4。这可以在 KV tile 之间并行化工作（右侧）。

我们的几项优化需要处理不规则内存访问（右侧），它们使用的指令和硬件与规则访问（左侧）不同。

调整并行策略在提升现代大规模并行硬件性能方面提供了最大的杠杆。直观地说：如果你被锁定在特定的并行方法上，那么 Amdahl 定律中的串行项就是固定的。如果你能改变并行策略，就可以在算法的并行和串行组件之间移动工作。根据该定律，这通常比提高固定并行组件的速度具有更高的杠杆作用。

我们并没有选择 CUDA Templates Domain Specific Language（CuTe DSL），这是原始内核作者的选择，但它对我们来说效果很好。它通过快速的 JIT 编译支持高生产力的开发循环，且运行时成本极低或为零。它也使得表达我们的许多想法比旧工具更直接。请注意，由于它使用模板，FA4 实际上是一个内核的_家族_，如果"内核"大致意味着"可以启动到 CUDA stream 中的东西"。我们将继续称其为"内核"。

CuTe DSL 很好。但是，正如我们在上一篇文章中指出的，FA4 最好在 tile 级别从算法上理解，而不是在其实现的 warp 级别。很明显，基于 tile 的编程在人体工程学和开发速度上会更好（顺便说一句，在 agent 时代这仍然很重要）。使用基于 tile 的编程模型，程序员可以更简单地表达和操作 tile 级别的流程。这使得以更低的工程成本更改或向内核添加算法变得更加容易（第一类更改）。此外，更高级别的基于 tile 的模型使得编译器更容易实现和优化，例如，cp.async 和 TMA 加载路径（第二类），并根据大小等进行调度。

有鉴于此，我们非常期待 CUDA Tile 编程模型（区别于经典的 "CUDA SIMT" 编程模型）得到更好的支持，以构建未来的 attention 和 matmul 内核。

我们做了什么、为什么做，以及如何知道它很好

我们按 pull request 来组织我们的贡献。每个部分都以一个"品质因数"开头：用于表明该贡献提升了性能的度量。我们以性能工程师的传统格式报告这些数字：一个 ASCII 表格。

PR 2109：支持 FP8 输入（于 2026 年 4 月 17 日合并）

品质因数：相对于 bf16 基线，吞吐量提升高达 1.16 倍

训练模型通常需要更高精度的浮点数来正确累积梯度中的许多微小变化。但在推理时，我们可以使用较低的精度。将位宽减半可将内存和算术带宽需求减半，而对模型质量的影响几乎没那么大。

对于大型模型的 MLP/MoE 层尤其如此，这些层通常使用微小的、"半字节"大小的 4 位浮点数。Attention 操作，尤其是在长上下文上，涉及更多累积，因此更难量化。像 gpt-oss 这样的模型将单精度 attention 操作与 4 位 matmul 结合起来，以两全其美。

然而，像 DeepSeek-V3 和 V4 这样的关键模型家族原生（即从训练开始）就支持 8 位 attention 操作。而其他模型，如 Qwen 和 Gemma 系列，有时会部署 8 位 KV 缓存以加速推理。

因此，我们添加了对 8 位浮点数（具有四个或五个指数位，即 e4m3 或 e5m2）的支持。相对于下面讨论的其他更改，这相当直接：移动和操作的字节更少意味着推理更快！这也意味着更小的 KV 缓存，这意味着更长的上下文和/或推理期间更高的用户并发度。

值得注意的是，加速比小于你可能从位宽减半中预期的 2 倍，位宽减半将内存带宽和（有效）算术带宽的需求都减半了。确定这里的具体瓶颈需要更详细的分析。但结果与 softmax 操作中的瓶颈一致，即使 Tensor Cores 在较低精度输入上操作，softmax 操作仍然以相同的精度（在 CUDA Cores 和/或 Special Function Units 上）运行。

PR 1999 和 PR 2104：支持任意 KV 页面大小（于 2025 年 11 月 13 日合并）并优化性能（于 2026 年 1 月 15 日合并）

品质因数：对于小页面大小，吞吐量提升高达 2.40 倍

FlashAttention-4 在 tile 上操作，其大小经过调整以有效利用 Blackwell Tensor Cores。在推理的解码阶段，键和值张量的 tile 由 KV 缓存中的条目构建而成，这些条目在预填充期间填充。在原始版本的 FlashAttention-4 中，KV 缓存页面需要与 tile 大小相同。

这个限制源于内核使用了 Tensor Memory Accelerator (TMA)，这是一个用于具有 Hopper 和 Blackwell 流多处理器 (SM) 架构的 GPU 中某些规则内存访问的硬件引擎。TMA 大幅加速了大型仿射内存访问——那些看起来像"偏移量加上步长乘以形状"的访问，涉及多个步长，例如通过 CuTe Layout 进行访问时。如果页面大小足够大，这对于访问基于页面的 KV 缓存效果很好。

但是 TMA 无法在一次加载中将多个分散的块收集到一个 tile 中，并且它不会加速（甚至可能减慢）较小的加载，而较小的加载是较小页面大小的结果。

因此，我们添加了一条路径，通过 PagedKVManager 使用 cpasync，即 CuTe DSL 对 PTX cp.async 指令的封装。

在基于 TMA 的版本中，一个 warp 中的单个线程负责加载一个 tile——生产者-消费者模型中的"生产者组"是一个单线程。

在 cpasync 版本中，每个线程发出一次加载（warp 的加载由硬件合并），因此它们计算自己的 page 和页面内的 offset。这很简单但效率低下；稍后会详细介绍！

我们重新利用了原本空闲的 warp 15 来处理这个额外的工作——生产者组由两个 warp 组成。

在这个最初的 PR 中，这些较小的页面大小具有较低的算术和内存吞吐量。但在许多推理工作负载中，KV 缓存效率非常重要，因此这可能是一个很好的权衡。

首先，大的页面大小可能导致不必要的重复。如果多个请求共享一个前缀，例如 64 个 token，但在此之后不同，那么一个 page_size=128 的 attention 内核将为每个请求需要一个单独的页面，因为前缀比页面大小短。一个 page_size=16 的 attention 内核可以在请求之间共享四个页面，从而将所需的存储量按请求数量成倍减少（参见下图中左侧三个请求共享前缀"Thou shalt not"，而右侧在具有较大 page_size 的 KV 缓存中该前缀被重复了三次）。

大的页面大小会导致 KV 缓存严重的内部碎片。短序列仍然需要完整的页面——在最坏的情况下，单个 token 消耗了整个本可容纳 128 个 token 的 KV 缓存数据的页面。对于该块来说，这超过了 99% 的内部碎片。这消耗了大约 8 倍于 page_size=16 的 KV 缓存的容量，而后者"仅"有 93.75% 的内部碎片。

这对于推测性解码尤其重要。推测器在 KV 缓存中创建许多短的（约 1-16 个 token）序列，而使用大的页面大小，每个序列都会消耗更多的空间。

支持任意页面大小对于兼容性来说已经是一个胜利，但最初的实现带来了性能成本。对于 page_size=1，即最极端的情况，FA4 内核的内存受限情况下的内存吞吐量不到有效内存带宽的一半，而计算受限情况下的算术吞吐量不到有效算术带宽的三分之一。我们在后续的 PR 中修复了性能问题。

类似的问题也影响了 FlashAttention-3 内核，因此我们将该策略移植到了 FA4 的 PagedKVManager。

关键举措是将地址_生成_与地址_使用_解耦，以减少冗余计算。这是通过"转置"地址生成来完成的，如下所述。该方法也在 Zadouri 等人的这篇论文的第 4.2 节中有详细说明。

我们将每个 warp 中的 32 个线程组织成一个数组，包含四个"行"线程组，每个组有八个"列"线程：

我们最初的方法让每个线程计算它同时也负责加载的 KV 缓存行的指针。

这里的加载模式受到硬件的限制——为了获得良好的内存合并，线程应该访问连续的内存。使用按行加载，相邻线程最终会冗余计算相同的行指针。

不幸的是，这种冗余代价高昂。指针是 64 位的，而 int64 操作很昂贵（最近的数据中心 GPU 的 FLOP 和 matmul FLOP 算术带宽的增长远远超过其他操作的带宽）。当需要计算更多地址时，这种成本更高，例如在较小的页面大小下。

解决方案是提前生成所有 32 个行指针，然后循环进行加载。这引入了 warp shuffle 形式的跨线程同步，但这比地址计算更便宜。

我们使用的具体模式是转置：我们 warp 中一个"行"组中的八个线程为 1) 不同的行生成行指针，这些行 2) 在逻辑上不是连续的。相反，跨组的"列"中的线程负责计算（但不使用）连续的行指针。

这使内存吞吐量比旧方法提高了高达 2.4 倍（对于 page_size=1），达到了与我们在较大页面大小下观察到的相同或更高的吞吐量。

PR 1940：添加跨 KV 维度的并行性（于 2025 年 11 月 4 日合并）

品质因数：对于小查询长度，吞吐量提升高达 4.37 倍

推理性能通常由解码时间主导。一个"典型"的推理请求大部分时间用于根据一个或几个查询针对许多缓存的 KV 值一次或几次地生成 token。

但是原始的 FlashAttention-4 内核架构在查询维度上并行化工作，而不是在键/值维度上。对于小批次推理，这对于高交互性、延迟敏感的应用至关重要，这简直是克星。内核程序的不同可并行化实例（协作线程数组）的数量通常远低于流多处理器 (SM) 的数量，导致多达 75% 的 SM 空闲（下图中淡化的部分），从而浪费了 75% 的 GPU 峰值性能。没有这个更改，FlashAttention-4 在 B200 上通常比 FlashAttention-2 慢！

解决方案是 Flash-Decoding，也称为"split KV"，由 Tri Dao 和合作者在 FlashAttention-2 时代引入。我们在参数 num_splits 下将 split KV 移植到了 FA4。在 split KV 模式下，多个 CTA 每个查询 tile 并发工作，每个 CTA 计算序列一部分的输出，最后通过一个 reduction 步骤产生最终结果。额外的 reduction 步骤在一个单独的内核 flash_fwd_combine 中。

跨 KV 维度进行拆分确保有工作可供多个 SM 使用，理想情况下是所有 SM。

由于浮点数的非结合性，带外 reduction 会引入数值差异。在一个 split 内求和，然后再跨 split 求和，与跨扁平序列求和得到的结果不同（对 monad 爱好者来说又是一个坏消息）。在我们的 split 路径中，我们使用 32 位浮点数进行共享内存输出 tile 累积，以减少影响，但无法完全消除。

额外的 reduction 步骤及其后果意味着 split KV 并不总是胜利。因此，我们添加了一个简单的启发式方法，根据 SM 数量和序列长度来检测最佳拆分数量（通过 num_splits = 0 触发）。

PR 1993：减少小查询大小的浪费工作（于 2026 年 1 月 8 日合并）

品质因数：单 token 解码吞吐量提升高达 3.06 倍

查询并行性并非唯一反映原始 FlashAttention-4 内核面向预填充或训练（其中有许多查询 tile）的选择。它被编写为同时操作两个查询 tile，每个查询 tile 有一个专用的四个 warp 的 warpgroup 来执行 softmax 操作（总共八个 warp）。每个 tile 由 128 个查询组成，因此此设置假设至少有 256 个查询。

但是，在低延迟推理期间，许多 attention 传递中的查询远少于 256 个，即使使用推测性解码和分组查询/多查询 attention（如下所述）。查询张量只是用零填充以填满剩余部分，这导致工作浪费。特别是，如果查询少于 128 个，则第二个 tile 上的所有工作都是不必要的！

因此，我们在核心 FA4 内核中添加了另一条路径，该路径一次只操作一个查询 tile（q_stage = 1）。此优化对于解码中看到的短查询序列长度特别有用，例如 seqlen_q = 1。

每个块只操作一个查询 tile 释放了第二个 softmax warpgroup，它通常对第二个查询 tile 运行 softmax 操作。我们将其重新用于在我们在 PR 1999 中添加的非 TMA/cpasync 情况下运行额外的 KV 页面加载，如上所述。

PR 2186：通过扩展 GQA 打包加速不规则 Q::KV 头比例（于 2026 年 3 月 20 日合并）

品质因数：单 token 解码吞吐量提升 2.92 倍

解码并不意味着每个序列只运行单个查询。分组查询 attention (GQA) 是一种架构变体，它对每个序列应用多个查询向量，针对每个 KV 向量。与多查询 attention (MQA) 一样，经典的 Shazeer 作品 是其基础，GQA 提高了推理的算术强度。

有一个问题：正如我们讨论过的，FA4 按_查询_分解 attention 计算——默认情况下，GQA 组中的每个查询都被单独处理。这意味着 KV 值需要被冗余加载，抵消了预期的内存加载减少。

解决方案当然是将该组映射到一个单一的 tile 中——即"GQA 打包"，通过标志 pack_GQA。这在 FA4 中已经实现。但它只适用于某些形状。具体来说，由于此路径使用 TMA 加载，它继承了 TMA 在对齐和布局方面的限制。每个 KV 头的查询头数量需要能整除 tile 大小 (128)。我们想要运行的一些模型，如 GLM 4.7，不满足此约束。

解决方案再次是使用 cpasync 进行没有 TMA 的常规加载，但这次是针对查询 tile 而不是 KV tile。上面为 PR 2104 描述的相同的基本转置/warp shuffle 策略已经为 Hopper GPU 实现，因此我们只需要将两者连接起来。

尾声

在 Modal，我们全力支持用于推理的开源软件。我们正在为像 FA4 这样的内核、像 SGLang 这样的推理引擎以及像 SLIME 这样的训练框架做出贡献，因为我们相信我们的基础设施是将这些软件部署到生产环境作为应用程序一部分的最佳场所，无论是提供推理服务还是训练模型。

如果你想为像 FlashAttention 或 SGLang 这样的项目做出贡献——或者如果你想构建运行它们的基础设施——我们正在招聘。