Awesome System Papers Wiki

Flash-Attention

10分钟阅读

Flash-Attention

FlashAttention 是一族 I/O-aware 的精确注意力 kernel。它用分块、online softmax 和算子融合，避免把完整的 L×L attention score matrix 写到 GPU HBM；输出语义仍是 dense attention，并没有通过丢弃 token 来减少计算 pairs。

核心思想

朴素实现把 QKᵀ、softmax 概率和 PV 分成多个 kernel，中间矩阵反复写入和读取 HBM。FlashAttention 把 Q/K/V 分成能放进 SRAM/SMEM 的 tiles，在片上逐块计算 score，并为每一行增量维护最大值、归一化因子和输出累加值。最后只把必要结果写回 HBM。

backward 通常不保存完整概率矩阵，而是根据 Q/K/V 和少量统计量重新计算。这用额外 FLOPs 换显存与 HBM I/O。因为 GPU 上矩阵乘很快、HBM 往返较贵，这个交换常有利。

“FlashAttention”既指 2022 年的算法，也常被用作后续 kernel family 的统称。不同版本共享 I/O-aware 思想，但不是同一份实现：

为什么重要

FlashAttention 的重要性不只是“一个更快 kernel”。它证明 attention 的性能上限要同时看计算和内存层次，也让更长 sequence 能在不保存二次大小中间状态的情况下训练。之后的编译器、DSL、低精度 attention、稀疏 attention 和 mega-kernel 都必须与最新强基线比较，而不是只胜过朴素 PyTorch 实现。

OSDI 2026 的 Twill-OSDI26 又把手写优化变成一个可求解问题：software pipelining 和 warp specialization 必须联合考虑，否则理论上 initiation interval 很小的 schedule 可能因寄存器、同步和跨 warp 通信无法执行。

关键观察 / 隐含假设

  • 减少 HBM I/O 可以在不近似模型的情况下加速。 FlashAttention-NeurIPS22 的核心收益来自不物化二次大小中间矩阵，而不是少算 attention pairs。它在 A100 和所测模型上成立，具体倍数不能直接外推到 decode、稀疏或新 GPU。
  • 算法相同，工作划分仍会限制硬件利用率。 FlashAttention-2-ICLR24 在 FA1 已经 I/O-aware 的基础上，通过 sequence 维并行和 split-Q 减少 shared-memory 通信。说明“复杂度最优”并不代表实现已接近峰值。
  • 新 GPU 的异步原语会改变整个流水。 FlashAttention-3-NeurIPS24 用 TMA/WGMMA overlap GEMM 与 softmax；FP8 更快，但必须处理量化和 outlier。结果主要来自 H100，不能把 FP8 throughput 当作无条件精度等价。
  • 算力增长快于其他单元时，softmax 也会成为瓶颈。 FlashAttention-4-MLSys26 发现 Blackwell 的 SMEM traffic 和 exponential 可超过 MMA 时间，于是用 TMEM、FMA 多项式和 conditional rescaling。部分指数近似和低精度路径需要单独检查误差。
  • 最小理论周期可能不可实现。 Twill-OSDI26 联合整数规划与 SMT，在固定 H100/B200 FP16 non-causal shape 上距 FA3/FA4 约 1%–2%；但实际 ptxas register spill 仍迫使作者收紧模型，高性能 CUDA 也不是自动 lowering 的最终输出。
  • 同一 kernel 在生产中可能因一个 tile 配置退化。 StriaTrace-OSDI26 的真实案例发现 build 漏掉特定 FlashAttention tile，回退版本在只剩 30% SM 时明显变慢。这说明 binary coverage、运行配置和可观测性是部署性能的一部分。
  • FlashAttention 的重计算也可成为可靠性信号。 AEGIS-OSDI26 利用 softmax 行和不变量，以及 forward/backward 的确定性重计算来检测 SDC。它只覆盖满足相应不变量或重放条件的错误，不能保证全部 FlashAttention 计算正确。
  • kernel 故障具有输入和实现亲和性。 SDCHunter-OSDI26 的生产缺陷 GPU 样本中包含只在 FlashAttention 等特定 kernel、dtype 和数值范围触发的 SDC；通用 GEMM stress test 和 ECC 可能看不见。
  • 把 KV 放到慢层需要重写数据流。 DirectKV-OSDI26 以 FA3 为基础，让 CPU-resident KV 在片上尽量复用，并融合 projection。它改变了内存层次和 kernel 边界，不是把现成 FA3 直接指向 host pointer 就得到同样结果。

设计空间与取舍

  • tile 大小与循环顺序：大 tile 增加复用，也占更多 SMEM/register；小 tile 并行度高，额外同步和加载更多。
  • 保存中间量或 backward 重算：重算减少 HBM 和显存，增加计算，也要求 forward/backward 的数值和 RNG 状态能正确对应。
  • 同 warp 协作或 warp specialization：专门 warp 可重叠 load、MMA、softmax 和 store，角色间同步与寄存器分配更难。
  • BF16/FP16、FP8 或整数路径：低精度提高吞吐，scale、指数近似和累加误差必须按任务验证。IntAttention-MLSys26 走全整数 softmax，是另一种硬件/数值设计点。
  • causal、non-causal、ragged、GQA 和 paged KV：mask、head sharing 和非连续地址改变 tile 合法性与负载均衡；一个固定 benchmark 不能代表所有 serving shape。
  • dense exact 或 sparse/approximate 组合：FlashAttention 本身不减少 pairs。与 sparse selection 组合后，索引和不规则 gather 可能破坏 dense coalescing，需要新 kernel。
  • 手写库、DSL 或自动搜索：手写能使用最新 ISA；Twill-OSDI26、Flashlight-MLSys26、WAVE-MLSys26、HipKittens-MLSys26 等提高生成或迁移能力，但各自仍有 shape、backend 和编译边界。
  • 单算子或 mega-kernel：独立 FA 易复用和调试；MPK-OSDI26 将 attention 与其他算子放进 persistent kernel，能跨边界重叠，却扩大编译、资源隔离和线上公平性问题。

引用本概念的论文

算法家族与自动优化

系统组合与替代数据流

作为组件或基线的引用

已知局限 / 开放问题

  • 自动跨 GPU 代际选择 tile、warp、pipeline、register 和数值路径仍未闭合；编译模型和实际 ptxas 行为会不一致。
  • ragged batch、paged KV、动态 mask、sparse attention 和 continuous batching 下，需要同时报告 kernel 与端到端 SLO，而不只报规则 dense shape 的 TFLOPs。
  • 低精度 online softmax、指数近似和 backward 重算需要更系统的误差界与任务质量验证。
  • 高度融合 kernel 的测试空间很大。应覆盖 dtype、head dimension、mask、极端 logits、别名、非确定性和 SDC，而不是只做正常输入性能测试。
  • 部署成本包括编译时间、binary 体积、tile coverage、fallback 和 driver/CUDA 兼容；这些成本常被 microbenchmark 忽略。

关系图谱

反向链接