FlashAttention 与算子融合

标准注意力要生成 $s \times s$ 的注意力矩阵再乘回——128k 序列的矩阵是 160 亿个元素，64 GB 起步。FlashAttention 的洞察：注意力矩阵从不需要完整存在。用分块计算 + online softmax，把显存从平方级压到线性，速度还快 2–4 倍。它是过去十年最重要的内核，没有之一。本章讲清楚它为什么快、怎么做到的、以及工程上的注意点。

标准 Attention 浪费在哪

标准实现的流程：

S = QKᵀ          ← 写出 s×s 矩阵到 HBM（平方级显存！）
S = softmax(S)   ← 再读一遍 s×s，写出 s×s
O = SV           ← 再读一遍 s×s

三次 HBM 往返，每次都是 $B \times h \times s \times s$ 大小。以 batch 4、16 头、8k 序列（BF16）：

s^2 \text{ 矩阵} = 4 \times 16 \times 8192^2 \times 2\ \text{B} = 8.6\ \text{GB} \times 3 \text{次往返} \approx 26\ \text{GB 流量}

而注意力的计算量只有 $4 \times B h s^2 d$ FLOPs——算术强度被 s² 矩阵的读写拖垮（Roofline 的访存受限侧）。128k 序列更是直接 OOM： $s^2$ 矩阵 64 GB × 3 份，HBM 根本装不下。

FlashAttention 的两个关键

一、分块计算（tiling）：不生成完整 $s^2$ 矩阵，把 Q/K/V 分块，一块一块算注意力——中间结果只存在共享内存/寄存器里，从不落到 HBM。

二、online softmax：分块算 softmax 有个数学问题——softmax 的分母需要全局 max 和全局和。标准算法必须先看全所有数据。online softmax 的解法：维护一个运行时的 max 和 sum，每来一个新块就增量修正：

m_{\text{new}} = \max(m_{\text{old}},\ m_{\text{块}}) \qquad \ell_{\text{new}} = \ell_{\text{old}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + \ell_{\text{块}} \cdot e^{m_{\text{块}} - m_{\text{new}}}

输出的部分和按同样方式增量缩放——数学上与一次性 softmax 完全等价，但每块算完就丢。

两个关键合起来：

	标准 Attention	FlashAttention
显存	$O(s^2)$ （s² 矩阵 ×3 份）	$O(s)$ （只存 O 与少量统计量）
HBM 流量	~3 × s²	O(s²·d/M)（M 是 SRAM 大小，远小于 s²）
速度	访存受限	计算受限，快 2–4 倍
8k 序列流量	~26 GB	~1 GB
128k 序列	64 GB ×3 ✗	~16 GB ✅

它同时解决了"装不下"（显存）与"跑不快"（带宽）——这是它成为标配的原因。

代际演进

版本	关键改进	收益
FA1（2022）	online softmax + 分块	显存平方→线性，快 2–4 倍
FA2	循环重排（K/V 外层循环减少 rescale）、并行度提升	快 ~2 倍，达到理论 FLOPs 的 50–70%
FA3	Hopper 特性：TMA 异步搬运、wgmma、warp 专用化（softmax 与 GEMM 重叠）	再快 1.5–2 倍，75–85% 理论算力
通信感知扩展	与上下文并行配合的环形 K/V 传递	长序列多卡

每代的主题一致：减少 HBM 往返、提高 Tensor Core 占空比——与GPU 架构一章的硬件主线完全对应。

算子融合：更大的图景

FlashAttention 是"算子融合"的极致案例，但融合不只是注意力：

融合	合并的算子	收益
FlashAttention	QKᵀ + softmax + SV + mask	消 s² 往返
fused LayerNorm	LN + 残差 + bias	消 2 次往返
fused optimizer	AdamW 的逐元素链	训练 step 的隐形大户
fused MLP	bias + GELU + 下游 tile	消激活往返

共同规律：逐元素与 reduction 类算子（访存受限，见 Roofline 的表）是融合的目标——它们算术强度 ~1，每多一次往返就是纯浪费。融合后的算术强度翻倍，时间减半。

工程上获得融合的三种方式（成本递增）：

torch.compile 自动融合（免费，见上一章）
现成的融合内核（FlashAttention、apex fused optimizer、框架内置）
Triton/手写（新算子，见Triton）

工程实践

默认启用：现代框架（PyTorch SDPA、transformers 库）已内置 FlashAttention 路径，attn_implementation="flash_attention_2" 或自动选择——先确认你的注意力真的走了 FlashAttention（不是数学等价的朴素实现）
变长序列：batch 内序列长度不齐时用 varlen 接口（packing），不要 padding 到最长——padding 是隐形的算力浪费
精度：FA 默认 FP16/BF16；FP8 路径（FA3）进一步提速但要注意数值（见混合精度与 FP8）
与 CP 配合：超长序列用上下文并行的环形注意力，FA 是计算核心
调试：数值对拍用"小序列 + 单卡 + 数学等价的朴素实现"做参考；online softmax 的 rescale 是 bug 高发区

本章要点

标准注意力的 s² 矩阵三次 HBM 往返：平方显存 + 访存受限，两头都是灾难。
FlashAttention = 分块 + online softmax：显存 $O(s^2) \to O(s)$ ，速度 2–4 倍，数学完全等价。
FA2 重排循环提并行度，FA3 用 TMA/wgmma 把硬件占空比打到 75–85%。
融合的共同目标：访存受限的逐元素/reduction 算子；获得方式从免费（compile）到手写递增。
先确认注意力真走了 FA 路径；变长用 varlen/packing，别 padding。