高效注意力：FlashAttention 与线性注意力
本文基于模型知识整理（生成时未联网核对），FlashAttention 机制建议对照 Dao et al. 2022/2023 原文复核。
一句话定义
对付 O(n²) 瓶颈（kp-005）的两条技术路线：FlashAttention（不近似！通过 IO 感知的分块计算让 n² 注意力矩阵永不落显存——速度与省显存兼得的精确算法）与近似家族（线性注意力把 O(n²) 改写为 O(n)、稀疏注意力只算重要位置）——前者是"工程奇迹"，后者是"数学换血"，现代训练两者并用。
为什么重要
FlashAttention 是长上下文训练能落地的直接功臣（几乎所有现代 LLM 训练都用 FA2）；线性注意力/SSM 则是"下一代架构"的候选路线（kp-031）。分辨"精确加速"与"近似换血"两类的本质差异，是评估一切"高效注意力"论文的第一道过滤器。
前置知识
kp-002/005（注意力与 O(n²) 瓶颈）；GPU 显存层级（HBM vs SRAM）概念。
核心概念
- FlashAttention（精确加速）：
- 洞察：瓶颈不在计算在IO（n² 矩阵在 HBM 写读，kp-005 的带宽受限分析）； - 方法：Q/K/V 分块载入 SRAM（快速片上内存，~20MB）——分块算 softmax 需要全局分母的技巧：在线 softmax（维护 running max 与 running sum，块间增量更新）；反向传播不存 n² 矩阵、重算注意力（以算换存）； - 成绩：数学结果与朴素实现完全一致（精确），显存 O(n)（非 O(n²)）、速度提升 2~4 倍（H100 上 FA3 再优化）。
- 线性注意力：把 softmax(QKᵀ)V 中的"先点积后 softmax"改写为"先特征化 φ(Q)(φ(K)ᵀV)"——利用结合律把 n×n 项消掉：
$$\sum_{ij}(q_i^\top k_j)v_j=\sum_i q_i^\top\underbrace{\left(\sum_j k_jv_j^\top\right)}_{d\times d}$$ ——O(n) 计算、O(d²) 状态（本质是"维护一个滚动摘要矩阵"，与 RNN 状态同构，kp-031 的桥）。代价：φ 替代 softmax 是近似（softmax 的归一化竞争被拆散），质量通常逊于精确注意力。
- 稀疏注意力：只计算"局部窗口+全局锚点/随机连接"的位置对（Longformer/BigBird）——从 n² 降到 n·w；近似假设：多数 token 对无关（长程稀疏）。滑动窗口（Mistral 的 SWA）是工程最简版。
- 组合现状：现代长上下文训练 = FlashAttention（精确）+ 滑窗/稀疏（限范围）+ RoPE 改造（kp-025）——多技术叠加而非单点方案。
原理与机制
在线 softmax 为什么可行：softmax 分子分母可增量维护——新块的局部 max/sum 与历史 running 值合并（$m^{new}=\max(m_{old},m_{block})$，权重按 $e^{m_{old}-m^{new}}$ 重缩放）——数学上与全量 softmax 严格等价。"流式算法"思想在注意力上的精确落地。
线性注意力为什么"像 RNN"：滚动摘要矩阵 $S_t=S_{t-1}+φ(k_t)v_t^\top$——每步只更新一个 d×d 状态（与 kp-031 SSM 的状态递归同构）；这就是"线性注意力=可并行训练的 RNN"论断的出处。注意力与 RNN 的边界被消融——两大架构家族在数学上握手（kp-031 的深层原因）。
为什么"重算"划算：反向传播需要注意力矩阵——存它 O(n²) 显存；重算它只需前向的 1.5~2 倍计算——在"显存比算力贵"的 GPU 经济学下是净赚（与 kp-029 梯度检查点同思想：以算换存）。
图示
FlashAttention(精确):
QKV分块→SRAM → 在线softmax(running max/sum) → 块间合并
显存O(n) ; 精确 ; 以重算换IO
线性注意力(近似):
Σᵢⱼ(qᵢᵀkⱼ)vⱼ = Σᵢ qᵢᵀ(Σⱼkⱼvⱼᵀ) → O(n) + d×d状态
softmax被φ替换 = 近似 ; 与RNN状态同构(桥向kp-031)
稀疏: 局部窗+全局锚 → n·w (Longformer/滑窗)
现状: FA(精确)+滑窗(限程) 叠加 ; 线性/SSM=下一代候选直观类比
朴素注意力像"把全城合影先洗出一张巨幅照片再找人脸"（n² 显存爆）；FlashAttention 像"分街区拍、当场用算盘汇总统计"（照片从未整张存在）——结果分毫不差。线性注意力像"不拍照，只维护一本'人群画像总账'"（d×d 状态）——翻页极快，但账本是摘要（近似）。
实例或案例
- 训练提速实证：FlashAttention-2 在长序列上把注意力内核提速至接近理论峰值——现代训练栈（Megatron/PyTorch SDPA）默认集成。
- Mistral 的 SWA：滑动窗口+RoPE 滚动缓冲——8k 窗口模型处理长文本的性价比方案（kp-025 联动）。
- 线性注意力的落地：RWKV/Mamba 的推理（常数状态、常数延迟）——长流式场景（kp-031）。
常见误区
- 误区一："FlashAttention 是近似"。精确算法——"Flash"指快不是"闪approx"；混淆会导致"不敢用于需要精确的场合"的错误回避。
- 误区二："线性注意力可以无缝替换 softmax 注意力"。质量差距在推理/检索类任务可感；替换=架构变更需重训——不是即插即用。
- 误区三："有了 FA 就不怕无限长上下文"。FA 只解决 IO 与显存、不消除 O(n²) 计算——长上下文还需 kp-025 的配套（位置插值/数据）。
与其他知识点的关系
自测题
- FlashAttention 的两个关键技巧与"精确"的保证？
答：分块+在线 softmax（running max/sum 增量合并）保证数值等价；反向用重算避免存 n² 矩阵——数学结果与朴素实现一致。
- 线性注意力如何消除 n×n 项？
答：结合律重排 $\sum(q^\top k)v=\sum q^\top(kv^\top)$——先累计 d×d 的 Σkvᵀ 再一次乘 φ(Q)，O(n²d)→O(nd²)。
- 为什么说"线性注意力是可并行的 RNN"？
答：其滚动状态 $S_t=S_{t-1}+φ(k)v^\top$ 与 RNN 隐状态递归同构——状态有限、可并行训练（kp-031 的 SSM 同款性质）。
延伸阅读
- Dao 等, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning"（2023）。
- Katharopoulos 等, "Transformers are RNNs"（2020，线性注意力开山）。
- Tay 等, "Efficient Transformers: A Survey"（2022，变体全景）。