一句话结论

FlashAttention 系列解决了"标准 attention 在 GPU 上跑不快"的系统问题——不改数学、不牺牲精度，靠 tiling（分块）和 IO-aware（感知内存层次）设计把 attention 推向硬件极限；Flash Linear Attention（FLA）系列则用同样的分块思想，把理论 O(N) 的线性 RNN 内核也提速到超越 FlashAttention。两条线在"长上下文高效序列建模"上汇合，共同构成了今天大模型长上下文能力的底层基础设施。

先说清一个名字：FlashAttention 是一系列论文（2022–2024），优化的是标准的 softmax attention；Flash Linear Attention（FLA） 是一个开源项目（2024 起），优化的是线性 attention / 线性 RNN。两者共享"分块 + 避免 HBM 往返"的思想，但不是同一个东西。

为什么 attention 慢：GPU 内存层次是第一瓶颈

标准 attention 需要计算 QK^T，时间与内存都随序列长度 N 二次增长。但很多人忽略一个更实际的问题：GPU 上跑的慢，往往不是算得慢，而是搬数据搬得慢。

GPU 内存是分层的：HBM（高带宽内存，容量大但慢）和片上 SRAM（容量小但快）。朴素实现把 N×N 的注意力矩阵完整写进 HBM 再读回来，大量时间花在内存搬运上，而不是矩阵乘上。

FlashAttention 系列盯住的就是这个"内存搬运"问题，用一句话概括它的核心策略：把中间结果留在离计算单元最近的地方，绝不落盘到 HBM。

先给出标准 attention 的数学形式，后面所有优化都是围绕它展开的：

$$ \text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right)V, \quad S = QK^\top \in \mathbb{R}^{N \times N} $$

其中 $Q, K, V$ 是查询、键、值矩阵，$N$ 是序列长度，$d$ 是每个头的维度。问题在于中间矩阵 $S$ 是 $N \times N$，序列变长时它主导了内存和时间。

FlashAttention（2022）：IO-aware 的精确注意力

要解决的问题

当时已有大量近似 attention 方法（稀疏、线性等）试图降低计算复杂度，但多数没有带来实际的 wall-clock 加速。原因是它们只优化了 FLOPs，忽略了内存访问开销。论文首次明确提出 IO-awareness 原则——算法设计必须把 HBM 与 SRAM 之间的读写次数作为一等公民。

核心机制

FlashAttention 用两个机制实现"精确 attention + 减少 HBM 访问"：

1. Tiling（分块）：把 Q、K、V 切成块，每个块的计算都在 SRAM 内完成。注意力矩阵（N×N）从不出现在 HBM 中，这是相比朴素实现最本质的差别。

2. Online softmax（在线 softmax）：标准 softmax 需要先看到整行才知道最大值和归一化项，分块后拿不到全局。FlashAttention 用"running maximum + running sum"的增量式更新：每个块到来时，用新的最大值重新缩放之前已累积的中间结果，最后统一归一化（Algorithm 1 第 10-12 行）。

用公式看更清楚。对向量 $x = (x^{(1)}, x^{(2)})$ 拆成两段分别做 softmax，再用各自的统计量合并：

$$ m(x) = \max(m(x^{(1)}), m(x^{(2)})), \quad \ell(x) = e^{m(x^{(1)}) - m(x)} \ell(x^{(1)}) + e^{m(x^{(2)}) - m(x)} \ell(x^{(2)}) $$$$ \text{softmax}(x) = \frac{f(x)}{\ell(x)}, \quad f(x) = \left[e^{m(x^{(1)}) - m(x)} f(x^{(1)}),\; e^{m(x^{(2)}) - m(x)} f(x^{(2)})\right] $$

这里 $m$ 是段内最大值，$\ell$ 是归一化和。分块计算时只需维护这两个标量统计量，每来一个新块就用新最大值重新缩放之前累积的 $f$ 和 $\ell$，最后统一除以 $\ell$。这就是"在线"的含义：不需要等整行算完，结果与一次性算完整 softmax 完全一致。

关键点：这个数学过程和标准 attention 完全等价，结果分毫不差。它不是一个近似方法，而是一个 IO 优化方法。

结果与意义

论文在摘要中给出的关键数字：

  • BERT-large（seq 512）端到端训练比 MLPerf 1.1 记录快 15%
  • GPT-2（seq 1K）快 3 倍
  • Long Range Arena（seq 1K–4K）快 2.4 倍
  • 首次让 Transformer 在 Path-X（16K，61.4%）和 Path-256（64K，63.1%）上超过随机水平

（上述数字出自论文摘要；正文实验部分给出了相同量级的结果，例如 GPT-2 相对 HuggingFace baseline 快 3 倍、相对 Megatron 快 1.8 倍。）

FlashAttention 没有改变 attention 的数学，只是让它跑得更快更省内存，因此成为后续几乎所有长上下文工作的底层依赖。

再补充一个重要的实现细节：重计算（recomputation）。反向传播通常需要注意力矩阵 $S$ 和 $P$，朴素实现得存下整个 $N \times N$ 的中间结果。FlashAttention 选择不存它们，反向时从 SRAM 里的 $Q,K,V$ 块重新算出来——这相当于把"存中间结果"换成了"多算一遍"，但因为省掉了大量 HBM 读写，反向反而更快（论文 Fig. 2）。

FlashAttention-2（2023）：把注意力推向 GEMM 效率

要解决的问题

FlashAttention 虽然比朴素实现快 2–4 倍，但只达到 GPU 理论峰值 FLOPs 的 25–40%，远低于矩阵乘（GEMM）的效率。瓶颈是工作划分不当：不同 thread block 和 warp 之间的负载不均，导致占用率低或共享内存读写过多。

三个改进

  1. 减少非矩阵乘 FLOPs：把 softmax 的缩放整合到后续运算里，避免重复读写。
  2. 跨 thread block 并行化：即使单个 head，也把 attention 计算分散到多个 thread block 上，提高占用率（利用 sequence 维度的天然并行性）。
  3. warp 间工作重分配：thread block 内部重新划分 warp 的工作，减少通过共享内存的通信。

结果

  • 相比 FlashAttention 约 2 倍加速
  • A100 上达到理论峰值 50–73%，接近 GEMM 效率
  • 端到端训练 GPT 模型时单卡 A100 达 225 TFLOPs/s（72% 模型 FLOPs 利用率）

FlashAttention-2 此后成为一年多里训练/推理框架的事实标准实现。

FlashAttention-3（2024）：为 Hopper 硬件重新设计

要解决的问题

FlashAttention-2 的算法假设"计算与内存搬运串行、按块同步推进"，没有利用 Hopper（H100）新增的硬件能力——Tensor Core 异步执行和 TMA（Tensor Memory Accelerator）。作者实测 FlashAttention-2 在 H100 上只有约 35% 利用率。

三个技术

  1. Warp specialization（warp 专用化）：让不同 warp 分别负责矩阵乘和内存搬运，用生产者-消费者模式把计算与数据移动重叠，而不是让每个 warp 同步做两件事。
  2. 块级 matmul 与 softmax 交错：把 softmax 的等待时间藏进矩阵乘的执行时间里，减少同步开销。论文为此重排了 FlashAttention-2 算法，绕开 softmax 与 GEMM 之间的顺序依赖（2-stage 版本里，softmax 处理一个分数块的同时，WGMMA 异步计算下一个块）。
  3. FP8 低精度 + block quantization（块量化）与 incoherent processing（非相干处理）：利用 Hopper 对 FP8 Tensor Core 的硬件支持，用按块量化和"去相干"技巧降低低精度带来的数值误差。

结果

论文在 H100 SXM5 上的实测：

  • FP16 相比 FlashAttention-2 前向快 1.5–2.0 倍，达 740 TFLOPs/s（75% 利用率）
  • FP8 接近 1.2 PFLOPs/s
  • FP8 版本（带块量化与非相干处理）比 per-tensor 量化的 baseline FP8 attention 数值误差低 2.6 倍

这些数字出自论文正文实验部分，是论文在特定硬件/配置下的声明。FlashAttention-3 展示了系统优化从"算法层"深入到了"硬件架构特性层"——同一数学，不同代 GPU 需要不同的内核设计。

Flash Linear Attention：从精确到线性的另一次提速

FLA（flash-linear-attention 仓库，Yang & Zhang，2024 起维护）把 FlashAttention 的分块思想移植到线性 attention / 线性 RNN 上，为 Mamba、RWKV、Gated DeltaNet、Mamba2、xLSTM 等一系列序列模型提供统一的高效 kernel 库。最接近论文形态的是 2025 年底的 Tiled Flash Linear Attention（Beck、Pöppel、Lippe、Hochreiter，arXiv 2503.14376）。

FLA 的核心思路

线性 RNN（linear attention）的隐藏状态可以写成 chunk-wise 形式：把序列切成块，块内用类似 attention 的并行计算（训练友好），块间用循环方式传递状态（推理友好）。这样既能并行计算，又能享受 O(N) 计算。

像 FlashAttention 一样，FLA 也把中间结果留在 SRAM 里，避免 HBM 往返——这正是它名字里"Flash"的来源。

用公式看 chunkwise 并行是什么。把序列 $T$ 切成 $N_c = \lceil T/L \rceil$ 个长度为 $L$ 的块，第 $k$ 块的查询、键、值为 $Q^{(k)}, K^{(k)}, V^{(k)}$。线性 RNN 的核心是块间循环（inter-chunk recurrence）：

$$ C_k = \bar{g}_k C_{k-1} + \bar{a}_k \odot K^{(k)\top} V^{(k)} $$

$C_k$ 是第 $k$ 块的隐藏状态，$\bar{g}_k$、$\bar{a}_k$ 是门控项（chunkwise gates）。这一步每个块只依赖上一个块的 $C_{k-1}$，可以顺序循环推进，只物化每块的起始状态，中间时刻的状态不用存。

块内是并行贡献（intra-chunk parallel contribution），对每个块内做类似 attention 的运算：

$$ S^{(k)} = \frac{1}{\sqrt{d_{qk}}} Q^{(k)} K^{(k)\top} \odot D^{(k)}, \quad H^{(k)}_{\text{intra}} = S^{(k)} V^{(k)} $$

$D^{(k)}$ 是块内的门控矩阵（对角占优、形状 $L \times L$），控制块内位置间的"遗忘"。最终的块输出是块内并行部分与块间循环部分之和：

$$ H^{(k)} = H^{(k)}_{\text{inter}} + H^{(k)}_{\text{intra}} $$

关键点：块内计算是二次的（$L^2$），但 $L$ 是常数（典型如 64），所以总复杂度是 $O(T \cdot L \cdot d)$，随序列长度线性增长——这就是"线性 RNN"名字的由来。chunkwise 公式把"并行训练"和"循环推理"统一到了一套表示里。

Tiled FLA：解决 FLA 的瓶颈

FLA 的 chunk size 受限于 GPU 片上 SRAM 的物理容量，不能无限增大。这导致必须物化很多中间状态到 HBM，算术强度（arithmetic intensity）低，长上下文预训练时内存和 IO 开销大。

TFLA 引入"chunk 内的额外一层序列并行"，让 chunk size 可以任意大，提高算术强度。论文把 TFLA 应用到 xLSTM 的矩阵记忆模块 mLSTM，并提出一个带 sigmoid 输入门、计算量更小的 mLSTM 变体。

论文声称新内核在速度基准上超过了高度优化的 Flash Attention、Linear Attention 和 Mamba 内核（代码在 NX-AI/mlstm_kernels）。这个"超过 FlashAttention"有具体前提：长上下文线性 RNN 场景，不是普遍意义上的"线性 attention 比 softmax attention 快"。

具体地，TFLA 的正文实验显示：在 H100 上，TFLA mLSTM 内核在训练（前向+反向）中对于长序列快于 FlashAttention-3，对所有序列长度比 Mamba 2 内核快 2 倍以上（Figure 5，embedding 4096、65536 tokens）。

两条线的关键差异

维度FlashAttention 1/2/3FLA / TFLA
目标让标准 softmax attention 更快让线性 attention / 线性 RNN 更快
数学精确，与标准 attention 一致本身是近似（线性化内核），追求与完整 attention 精度接近
瓶颈HBM 与 SRAM 的 IO同样 IO，外加 chunk 并行度与中间状态物化
核心技巧tiling + online softmaxchunkwise 并行 + 状态传递
适用短-中上下文 Transformer长上下文、可压缩/状态化模型
开源Dao-AILab/flash-attentionfla-org/flash-linear-attention（Triton 实现，跨 NVIDIA/AMD/Intel）

一句话总结关系：FlashAttention 证明了"把内存搬运优化到极致，精确 attention 可以非常快"；FLA 借用了这个方法论，证明"理论 O(N) 的线性模型，只要内核写得好，也能跑得比 FlashAttention 更快"——但代价是数学上不再精确，精度由模型本身（门控、状态设计）来弥补。

读论文顺序建议

  1. FlashAttention 1：理解 tiling + online softmax 这两个核心机制，这是整条线的地基。
  2. FlashAttention-2：理解 GPU 并行模型（thread block / warp 层次）如何影响实际速度，建立"理论 FLOPs ≠ 实际速度"的直觉。
  3. FlashAttention-3：可略读。它偏向特定硬件（Hopper）的工程技巧，对理解算法本质帮助不大，但对写 GPU 内核的人很有价值。
  4. FLA + Tiled FLA：此时你会自然理解"线性模型为什么需要专门内核"“chunkwise 并行是什么”，也能看懂 FLA 与 FlashAttention 的关系。

适用边界与待验证点

  • 本文所有具体数字均来自论文摘要和正文（FlashAttention-3 的 H100 实测来自正文实验部分），未做第三方独立复验。
  • FlashAttention 系列的"精确"指数学上与标准 softmax attention 等价；FP8 版本的 FlashAttention-3 是低精度近似，论文用块量化和非相干处理控制误差。
  • FLA / TFLA 的"超越 FlashAttention"是特定长上下文线性 RNN 场景下的结论，不能外推到所有 attention 变体。
  • FlashAttention 系列官方代码在 Dao-AILab/flash-attention；FLA 是 Triton 实现，支持跨厂商 GPU（NVIDIA/AMD/Intel），这是它与 CUDA 原生的 FlashAttention 的重要工程差异。

参考论文

  • FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（Dao et al., 2022, arXiv 2205.14135）
  • FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning（Dao, 2023, arXiv 2307.08691）
  • FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（Shah et al., 2024, arXiv 2407.08608）
  • Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels（Beck et al., 2025, arXiv 2503.14376）
  • FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism（Yang & Zhang, 2024, GitHub）