FlashAttention 深度解析：大模型长序列的算力解放

Transformer 的长序列处理一直被算力瓶颈卡着脖子——不是算不动，是显存带宽不够用。FlashAttention 的思路很直接：把数据搬运算搬到 SRAM 里分批做完，减少对 HBM 的依赖。没有魔法，但足够聪明。

一、为什么 Attention 这么慢？

要理解 FlashAttention，得先搞清楚 GPU 的三层存储结构：

层级	容量	带宽	速度特征
寄存器	~64KB	~50 TB/s	最快，直连计算单元
SRAM（共享内存）	~192KB	~20 TB/s	次快，每个SM独立
HBM（显存）	~80GB	~2 TB/s	慢一个数量级

← 表格可左右滑动查看

标准 Attention 的计算过程看起来简单——O = softmax(QKᵀ) · V——但在 GPU 上的实际执行是这样的：

① Q, K, V 从 HBM 加载到 SRAM
② 计算 S = Q·Kᵀ，写回 HBM
③ 读回 S 做 softmax → P，写回 HBM
④ 读回 P，乘 V → O，写回 HBM

一共 6 次 HBM 访问。 每一次都在等显存带宽。

虽然 HBM 有 80GB 的大容量（可以放下长序列的 Attention 矩阵），但 2 TB/s 的带宽和 SRAM 的 20 TB/s 差了一个数量级。Attention 的计算本身很快（矩阵乘法是 GPU 的强项），但数据搬运占了绝大部分时间。

简单估算一下：对于 64K 长度的序列，标准 Attention 需要中间矩阵 S 的大小是 64K × 64K = 4B 个元素，用 FP16 存储就是 8GB。每轮计算要在 8GB 数据和 HBM 之间来回搬运 6 次 → 48GB 的数据传输。按 2TB/s 算，光搬数据就占 24ms，而计算本身可能只需要几毫秒。

IO 复杂度：O(N²) 的数据搬运，这就是长序列时的瓶颈所在。

二、FlashAttention 的核心：Tiling

FlashAttention 的改进思路很直接：不要让大矩阵来回穿梭 HBM，而是把计算切分成小块，全部在 SRAM 里完成后再合并。

这个技术叫 Tiling（分块计算）。

核心流程：

分块： 把 Q、K、V 按 SRAM 容量切分成小块（block），每块约 64-128 个 token
加载： 一次只加载一个 Q 块和对应的 K、V 块到 SRAM
局部计算： 在 SRAM 里算 S = Q_block · K_blockᵀ（小矩阵乘法）
在线 softmax： 在同一块 SRAM 里完成 softmax——这是关键创新点
累加： 把结果累加到输出 O 的对应位置
重复： 所有块都算完，得到完整的 Attention 输出

在线 softmax 算法

传统 softmax 需要知道整个序列的最大值和分母求和，这就要求全量数据在同一个地方。FlashAttention 用了一个技巧：

对每个 block 计算局部 softmax
用统计方法（rescaling）在累加时修正之前 block 的贡献
最终结果在数学上等价于全局 softmax，但 不需要一次加载所有数据

这个算法细节是 FlashAttention 真正的技术贡献——它让分块+精确的 Attention 计算变得可能，无需近似或妥协。

三、性能对比

指标	标准 Attention	FlashAttention	提升
HBM 访问次数	6 次	2 次	3× 减少
速度（4K 序列）	基线	1.5-2×	—
速度（64K 序列）	基线	2.5-4×	长序列优势更大
精确度	标准	严格等价	🎯 无损

← 表格可左右滑动查看

为什么长序列提升更大？因为序列越长，中间矩阵 S 越大，IO 负担越重，Tiling 节省的效果越显著。

注意： FlashAttention 不是近似方法。它和标准 Attention 在数学上完全等价——每 bit 都一样。

数据来源：FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Tri Dao et al., 2022)

四、后续演进

FlashAttention 自 2022 年提出后，已经有多个迭代版本：

FlashAttention-2（2023）：减少了 non-matmul 操作，优化了线程块调度，速度快了约 2×
FlashAttention-3（2024）：针对 Hopper 架构（H100）的 FP8 优化，进一步利用 WGMMA 指令
Block-Sparse FlashAttention：支持稀疏 Attention 的加速版本

如今，PyTorch 的 torch.nn.functional.scaled_dot_product_attention 底层已经集成了 FlashAttention（有 GPU 和 cuDNN 支持的情况下自动启用）。也就是说，大多数大模型已经在用了，可能只是你不知道。

五、总结

瓶颈不是算力，是带宽 — Attention 的计算/IO 比极低，数据搬运才是大头
Tiling + 在线 softmax = FlashAttention 的核心技术创新
无损加速 — 在数学上严格等价的前提下，做到 2-4× 速度提升
应用广泛 — 已是现代 Transformer 框架的标准组件

锐Owl，不是更快，是更聪明。

参考文献：

[1] Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv preprint arXiv:2205.14135.
[2] Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv preprint arXiv:2307.08691.
[3] Shah, J., et al. (2024). FlashAttention-3: Fast and Accurate Attention with Hardware-Awareness. arXiv preprint arXiv:2407.08614.

FlashAttention 深度解析：大模型长序列的算力解放

本期要点

关键实测数据

FlashAttention 深度解析：大模型长序列的算力解放

一、为什么 Attention 这么慢？

二、FlashAttention 的核心：Tiling

在线 softmax 算法

三、性能对比

四、后续演进

五、总结