FlashAttention 如何用 IO 感知分块加速 Transformer 注意力计算

一个在线大模型服务把上下文窗口从较短对话扩展到长文档问答后，经常会出现一种看似矛盾的现象：GPU 的理论计算能力很高，模型参数也已经放进显存，但请求的首个 Token 延迟仍随着输入长度迅速上升。工程团队继续增加算力并不一定能按比例改善延迟，因为 Attention 不只是在做乘法，它还要不断搬运 Q、K、V 和中间结果。序列越长，中间数据越大，内存系统承担的压力越明显。

FlashAttention 解决的正是这个执行层问题。它没有把标准 Attention 换成近似算法，也没有从根本上改变其二次复杂度，而是重新安排计算顺序：将大矩阵拆成能放进片上高速存储的小块，在块内完成尽可能多的计算，并用可增量更新的 softmax 统计量避免把完整 Attention Score 和概率矩阵写回 HBM。对长上下文训练和长 Prompt 的 prefill 阶段，这种 IO-aware 设计能够显著减少数据搬运；但在逐 Token decode、短序列或 Attention 本来就不是瓶颈的场景中，收益会明显不同。

标准 Attention 为什么会被中间矩阵拖慢

标准缩放点积 Attention 可以写成 softmax(QKᵀ)V。从数学上看，过程并不复杂：先让查询矩阵 Q 与键矩阵 K 相乘得到分数，再对每一行执行 softmax，最后用得到的概率对 V 做加权求和。问题在于，一个直接实现往往会把这些阶段拆成多个独立 kernel，每个阶段都需要从显存读取输入，并把中间结果重新写回显存，供下一个 kernel 使用。

假设输入序列长度为 8000。单看某个 Attention head，QKᵀ 就对应一个 8000×8000 的分数矩阵。序列长度继续翻倍时，这个矩阵的元素数量会增长到原来的四倍。即使 GPU 能快速完成矩阵乘法，把如此大的中间结果写入 HBM，再读回来执行 softmax，随后再次读写用于与 V 相乘，也会产生可观的数据移动成本。

这里需要区分计算复杂度和数据移动。FlashAttention 并没有让 QKᵀ 少做一个数量级的乘加，因此它不是通过把 O(N²) 变成 O(N) 来加速标准 Attention。它改变的是哪些数据必须落到 HBM、哪些数据可以在片上保存、一个数据块被加载后能完成多少工作。对于受内存带宽限制的执行阶段，少一次大规模 HBM 往返就可能比减少少量算术指令更重要。

在实际服务中，这也是为什么只看 GPU FLOPS 很容易误判。一个 kernel 可能没有把计算单元跑满，不是因为缺少待计算的乘法，而是因为执行单元正在等待数据。判断 FlashAttention 是否值得启用，应观察 Attention kernel 的时间占比、显存带宽压力、序列长度分布和整体请求阶段，而不是只看模型参数量。

IO-aware 的关键：让数据块在片上完成更多工作

GPU 存储存在明显的层级差异。HBM 容量较大，适合保存模型权重、激活和 KV Cache，但访问成本高于寄存器、共享内存等片上资源。片上存储速度快，却无法容纳完整的长序列 Attention 矩阵。FlashAttention 的基本策略不是试图把整个矩阵塞进片上存储，而是采用 tiling，也就是分块。

计算时，kernel 取一小块 Q，再依次加载对应的 K、V 块。当前块产生的局部分数会立即参与 softmax 统计和输出累积。只要这部分中间状态仍能保留在寄存器或共享内存中，就不必把完整分数矩阵写回 HBM。处理完一个 K、V 块后，kernel 更新当前输出和归一化所需的统计量，再继续处理下一块。

这与“先生成整个 Attention Score，再交给下一个阶段”有本质区别。传统分阶段实现的中间矩阵是跨 kernel 的物化结果；FlashAttention 则把多个逻辑阶段融合到一个更紧凑的数据流中。数学结果仍对应标准 Attention，但执行路径不再要求完整保存那个二次增长的中间矩阵。

flowchart TD
  A[加载一块 Q] --> B[加载一块 K 和 V]
  B --> C[计算局部分数]
  C --> D[更新行最大值与归一化统计]
  D --> E[累积当前输出]
  E --> F{还有 K/V 块吗}
  F -->|有| B
  F -->|没有| G[写回最终 Attention 输出]
  G --> H[进入后续 Transformer 层]

对长 Prompt 的 prefill 来说，这个流程尤其重要。此时大量 Token 会同时进入 Attention 计算，Q 与 K 的二维交互规模大，中间矩阵的物化成本突出。FlashAttention 通过块级复用把更多计算留在一次 kernel 生命周期中，减少 HBM 往返。这里的“Flash”不应理解为闪存技术；真正有工程意义的关键词是 IO-aware，即显式考虑内存层级和数据移动成本来设计算法实现。

不保存完整矩阵，softmax 为什么还能算对

分块带来的一个直接问题是 softmax。对某一行分数 x，softmax 需要知道整行所有元素，因为分母包含所有指数项。如果只拿到当前块，直接对每个块分别做 softmax，再把结果拼起来，得到的并不是原始 Attention 的结果。

FlashAttention 能够分块处理，是因为 softmax 可以用增量方式维护。处理第一批分数时，kernel 记录当前行的最大值、指数和以及对应的输出累积；处理下一批分数时，如果发现更大的行最大值，就按照新的尺度重新缩放之前的累计量，再把新块贡献合并进去。这样只需要保存每一行少量统计状态，而不需要保存所有分数。

这种 online softmax 是 FlashAttention 保持“精确 Attention”语义的关键。这里的“精确”指算法目标仍然是标准 softmax Attention，而不是通过稀疏化、低秩近似或截断窗口改变要计算的注意力关系。实际 GPU 浮点计算仍可能因为运算顺序不同出现细小数值差异，因此更准确的工程表述是：它保持标准 Attention 的数学定义，而不是承诺不同 kernel 的每一位浮点结果完全一致。

从数据生命周期看，Q、K、V 从 HBM 被分块加载，局部分数只在片上短暂存在；行最大值和归一化因子随着块推进不断更新；输出向量也以累积形式逐步形成。直到一整行相关块处理完毕，最终输出才需要写回 HBM。原本最大的两个中间对象——完整分数矩阵和完整 softmax 概率矩阵——不再成为必须跨 kernel 保存的全局中间结果。

这也解释了 FlashAttention 为什么同时改善速度和训练时的激活内存压力。减少中间矩阵物化不仅意味着少搬数据，也意味着无需长期为这些对象保留同等规模的显存空间。反向传播仍需要正确重建所需信息，但实现可以通过保存较小的统计量并重新计算部分中间值，在计算与内存之间做更合适的交换。

FlashAttention-2 优化的是并行工作分配

第一代 FlashAttention 解决了“数据怎么少搬”的核心问题，但一个高效 GPU kernel 还要回答另一个问题：工作如何分给不同线程块和 warp，才能让更多执行单元持续忙碌。FlashAttention-2 的重点就是重新组织并行方式和工作划分，减少不必要的同步与非矩阵乘运算开销，并提高不同序列长度和 head 配置下的硬件利用率。

用户提供的 FlashAttention-2 资料中报告，在 A100 上相对第一版 Attention kernel 还能获得进一步加速，并给出了达到较高理论峰值利用率的实验结果。这类数字应理解为论文特定硬件、形状和实现条件下的 kernel 结果，而不是“任何模型端到端都固定快两倍”。模型整体还包含线性层、MLP、通信、采样、KV Cache 访问等开销，Attention kernel 加速多少，最终只能按它在整个请求中的占比折算。

从工程视角看，FlashAttention-2 的意义在于说明：IO 优化完成后，新的瓶颈可能转移到并行度、线程间通信和指令组成。优化 GPU kernel 往往是逐层推进的。第一步消除巨大的 HBM 中间结果，下一步再提升 occupancy 和 workload partitioning；硬件架构变化后，还可能继续调整数据搬运指令、矩阵乘单元和片上存储之间的配合。

因此，部署时不应把“FlashAttention”当成一个抽象开关后就停止验证。不同框架可能根据 GPU、数据类型、head dimension、mask 形式和序列形状选择不同 Attention backend。某些输入会走优化 kernel，另一些输入可能回退到其他实现。生产环境需要确认实际执行路径，而不是只确认配置文件里出现了某个名称。

训练、prefill 和逐 Token decode 的收益并不相同

“FlashAttention 能加速推理”这句话过于宽泛。大模型在线推理至少可以拆成 prefill 和 decode 两个阶段。prefill 处理完整 Prompt，需要让输入 Token 之间执行大规模 Attention，计算形态更接近训练中的长序列前向过程；decode 每轮只新增一个 Token，新的 Query 去读取已有 KV Cache，计算和访存形态已经不同。

在长上下文 prefill 中，避免物化大规模 Attention 中间矩阵通常具有直接价值。输入越长，传统实现的中间数据压力越明显，FlashAttention 的 IO 优势越容易体现。训练场景也类似，因为需要同时处理大量 Token，并考虑前向和反向的激活内存。这里既可能获得 kernel 时间收益，也可能因为降低内存占用而允许使用更长序列或更合适的 batch。

逐 Token decode 则不能简单套用同一结论。每轮只有少量新 Query，但必须访问此前累计的 K、V。此时 KV Cache 读取、批处理调度、模型权重访问等因素可能成为主要瓶颈。高性能推理引擎仍会使用专门优化的 Attention kernel，但“把完整 N×N 中间矩阵消掉”不再是唯一主线，因为 decode 本身通常不会按训练阶段的方式物化同样形态的大矩阵。

场景	主要 Attention 形态	FlashAttention 类优化的典型价值	需要同时关注的瓶颈
长序列训练	大量 Query 与 Key 形成二维交互	减少中间矩阵读写和激活内存压力	反向计算、通信、整体 batch 规划
长 Prompt prefill	一次处理较多输入 Token	降低 HBM 流量，缩短首 Token 前的 Attention 阶段	权重计算、并行通信、Prompt 长度分布
逐 Token decode	少量新 Query 读取历史 KV Cache	依赖具体 kernel 和批处理形态，收益不等同于 prefill	KV Cache 带宽、调度、并发和采样
短序列小模型	Attention 数据规模较小	优化收益可能被 kernel 启动和其他层开销稀释	MLP、框架开销、请求调度

这张表也是评估“为什么换了 FlashAttention，端到端吞吐却没有翻倍”的起点。若线上流量以短回答和短 Prompt 为主，Attention 占总时间不高，那么优化它只影响请求的一小部分。若业务是长文档问答，TTFT 被长 Prompt prefill 主导，则同一个优化可能更有价值。

与近似 Attention 的区别决定了适用边界

FlashAttention 常被和长上下文算法放在一起讨论，但它与稀疏 Attention、滑动窗口 Attention、低秩近似等路线解决的问题不同。后者通常通过减少需要参与计算的 Token 对，改变计算图或注意力连接范围，从算法层降低工作量；FlashAttention 的核心目标则是在仍计算标准 Attention 的前提下，让执行过程更符合 GPU 内存层级。

这一区别直接影响工程决策。需要保持现有模型语义、权重和 Attention 结构时，替换为兼容的 FlashAttention kernel 通常比改模型结构更容易控制风险。若上下文继续增长到标准二次 Attention 本身已经无法接受，即使每次 HBM 访问都非常高效，O(N²) 的算术工作仍然存在，此时可能需要窗口化、稀疏化或其他结构性方案。

方案	是否改变标准 Attention 计算范围	主要优化对象	主要收益	主要代价或边界
直接物化的标准 Attention	否	无专门 IO 优化	实现直观、兼容性高	长序列中间矩阵读写和内存压力大
FlashAttention	否	HBM IO 与片上数据复用	减少中间矩阵物化，提升长序列执行效率	依赖高质量底层 kernel 与硬件适配
FlashAttention-2	否	IO 加并行工作划分	在前述基础上提高并行效率	实际收益仍依赖形状、硬件和端到端瓶颈
稀疏或窗口 Attention	通常会	参与 Attention 的 Token 对数量	可降低超长序列的算术与存储压力	改变注意力范围，需要模型和效果验证

因此，“FlashAttention 让长上下文变得无限便宜”是不成立的。它减少的是不必要的数据移动和中间存储，不会消除标准 Attention 随序列长度平方增长的核心计算量。它能把原本低效的执行推向更合理的硬件利用率，但不能绕过算法规模本身。

在推理引擎中落地时要看真实执行路径

用户原稿提到了 vLLM、TensorRT-LLM、Hugging Face Transformers 和 PyTorch 生态。对使用者来说，真正重要的不是这些框架是否“支持 FlashAttention”这个静态结论，而是当前模型、GPU、数据类型和输入形状是否实际命中了目标 kernel。框架升级、模型结构变化或某个不兼容参数，都可能让执行路径发生改变。

生产环境中可以把验证分成三层。第一层是功能正确性：同一批固定输入在切换 Attention backend 前后，输出应保持可接受的一致性，不能因为 mask、因果关系或 padding 处理错误造成语义变化。第二层是 kernel 级性能：用 profiler 确认 Attention 阶段耗时、HBM 流量和 GPU 利用情况确实改善。第三层是服务级指标：观察 TTFT、每 Token 延迟、吞吐、显存峰值和并发能力，确认局部优化最终转化为用户可见收益。

还要特别注意 fallback。一个服务可能在常见 head dimension 上使用优化 kernel，但遇到特殊模型配置后回退；也可能在训练路径使用一种实现，在推理路径使用另一种实现。若只做一次短序列 benchmark，很容易得出错误结论。更可靠的测试应该覆盖真实的 Prompt 长度分布、batch 大小、并发度和目标硬件。

对于长上下文服务，还可以把 prefill 和 decode 分开观测。如果启用 FlashAttention 后 TTFT 明显改善而 TPOT 变化不大，这并不矛盾，反而可能说明优化命中了长 Prompt Attention，而逐 Token decode 的瓶颈仍在 KV Cache 带宽或调度。如果两个指标都没有变化，则应继续确认是否实际使用了目标 kernel，以及 Attention 是否本来就不是当前服务的主瓶颈。

哪些情况下收益会退化

第一类退化来自序列太短。分块、融合和专用 kernel 本身也有调度成本。当 Attention 中间矩阵很小，传统实现的数据移动尚未成为主要问题时，优化空间有限，端到端差异可能被其他层覆盖。

第二类来自硬件和 kernel 不匹配。FlashAttention 依赖针对 GPU 执行模型设计的底层实现，不同硬件架构具有不同的片上资源、矩阵计算能力和数据搬运机制。即使算法思路相同，最佳分块大小、warp 分工和流水方式也可能不同。用户提供的 Hopper CUDA kernel 案例资料本身就体现了这种硬件相关性：高性能实现需要结合具体架构重新组织 kernel，而不是写一次 Python 逻辑就自动获得同样性能。

第三类来自瓶颈迁移。Attention 变快后，MLP、跨卡通信、KV Cache、采样或请求调度可能占据更高比例。此时继续优化 Attention 的边际收益会下降。性能工程不能停在“单 kernel 快了多少”，而要重新做端到端 profile，找到新的最长路径。

第四类来自算法规模本身。FlashAttention 不消除二次计算复杂度。上下文长度持续增长时，即使中间矩阵不落 HBM，实际乘加次数仍快速增加。到了这一边界，仅靠 IO-aware 实现无法解决全部问题，需要结合模型结构、上下文策略或近似 Attention 做更上层的设计选择。

如何判断自己的系统是否值得启用

最实用的判断方法是先定位瓶颈，再选择优化。对于长文档问答服务，可以挑选接近生产分布的短、中、长三组 Prompt，在保持模型、硬件、batch 和并发条件一致的情况下，对比不同 Attention backend。除了端到端延迟，还要拆分 prefill 与 decode，并记录显存峰值和吞吐变化。

如果长 Prompt 的 prefill 时间占比高，Attention kernel 在 profiler 中明显消耗时间，同时 HBM 数据移动压力突出，那么 FlashAttention 类实现通常具有较强的适配理由。若业务主要是极短输入、极短输出，或者性能已经被跨卡通信和模型权重读取主导，则应先解决真正的限制因素。

最终应把 FlashAttention 理解为一种系统级执行优化：数学上的 Attention 仍然存在，Q、K、V 仍要参与同样的核心关系计算，但中间数据不再按最直观、最昂贵的方式在 GPU 内存层级之间来回搬运。它展示了大模型性能优化中一个重要事实——决定速度的不只是“算多少”，还包括“数据在哪里、什么时候移动、移动后能复用多少次”。当这些问题被纳入算法实现本身，GPU 才更有机会把理论计算能力转化为真实的训练和推理效率。