东吴北上赋

Mon, 06 Jul 2026 00:00:00 GMT

岁次壬寅，序属季夏。负箧吴江，栖迟久泳。时维高考新败，登阊门而北望，临胥江以长嗟。昔者子安命蹇，尚能奋藻于滕阁；今吾才疏，岂可沉沦于吴市？乃焚膏以继晷，立雪而追贤。自暑月即研电学之微，探数理之奥。但求踔厉侪辈，砥砺锋芒。

洎乎素商既至，庠序始开。选修五门皆得满绩，总评冠绝同侪。然尘世难料，人心叵测。交非其人，悔吝丛生。一载相交，竟转瞬成参商。当是时也，恍若屈子见放，贾生受谗。然仲尼厄而作春秋，左丘眇而厥有国语。经此砥砺，竟愈宿疾：昔之曲意逢迎，今则坦荡自持；向之汲汲人言，今则泰然自若。所谓塞翁失马，焉知非福耶？

若夫玄冥司节，青阳启序。既专课业，复骋赛场。数模电赛，殚精竭虑；车赛集创，呕心镂骨。常观吴江夜月，每伴姑苏晨星。悬梁刺股，非为功名之累；凿壁囊萤，实怀鸿鹄之志。终使课业蝉联榜首，竞赛累获殊荣。然形神俱瘁，犹记寒宵病骨，强支案牍；伏暑昏眩，犹自深研不辍。

苏大本科毕业

岁次甲辰，时逢白藏。绩点累年称冠，奖项盈箧成行。遂膺国奖之荣，如登龙门之津。乃遍谒名师，广求良策。拟鹏徙南冥，期凤鸣岐阳。幸得燕园李萌先生青眼，许列门墙。继入艾捷科芯实习，兼修毕业设计。两事交并，心力俱疲。承钟林峰师兄鼎力，终克难关。

今当辞别葑溪，将赴燕台。忆昔韩洪润学长，指迷津于暗夜；念吴圣洁挚友，伴苦读于寒窗。实乃天眷优渥，得遇诸君。昔者范公划粥，终成社稷之臣；欧母画荻，乃育文章之伯。予虽驽钝，敢不踵武前修？悟已往之不谏，知来者之可追。陶元亮归去来辞，实获我心；王子安穷且益坚，宁移素志？今将整装而北发，岂效楚囚之泣？当乘长风，破巨浪，展鸿图于未央！

北大研究生录取

推理优化 | FlashAttention 1-4 详解

Fri, 10 Jul 2026 00:00:00 GMT

本文受 Datawhale llm-algo-leetcode 活动启发而作，具体参考了其中的 FlashAttention 仿真实现与 Triton Flash Attention 示例两个章节。该项目系统整理了 LLM 推理/训练中的经典算法，推荐关注。

算法#

原算法：标准 Attention 与显存瓶颈#

Transformer 的标准缩放点积注意力定义为：

Attention(Q,K,V)=softmax ⁣(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) VAttention(Q,K,V)=softmax(dkQKT)V

其中 Q,K,V∈RN×dQ, K, V \in \mathbb{R}^{N \times d}Q,K,V∈RN×d，NNN 为序列长度，ddd 为头维度。计算流程如下：

S=QKT∈RN×N(矩阵乘)P=softmax ⁣(Sd)∈RN×N(逐行 softmax)O=PV∈RN×d(矩阵乘)\begin{aligned} S &= Q K^T &&\in \mathbb{R}^{N \times N} \quad \text{(矩阵乘)} \\ P &= \text{softmax}\!\left(\frac{S}{\sqrt{d}}\right) &&\in \mathbb{R}^{N \times N} \quad \text{(逐行 softmax)} \\ O &= P V &&\in \mathbb{R}^{N \times d} \quad \text{(矩阵乘)} \end{aligned}SPO=QKT=softmax(dS)=PV∈RN×N(矩阵乘)∈RN×N(逐行 softmax)∈RN×d(矩阵乘)

瓶颈分析：

SSS 和 PPP 均为 O(N2)O(N^2)O(N2) 大小的中间矩阵，必须写入 HBM 再读出，带宽严重不足
GPU HBM 带宽（A100 ~2TB/s）远低于片上 SRAM（~19TB/s），而 Attention 的计算量仅为 O(N2d)O(N^2 d)O(N2d)，使得整体呈 Memory-Bound
对长序列（如 N≥2048N \ge 2048N≥2048），超过 90% 的时间耗费在显存读写上

1
HBM 数据流视角（标准实现）：
2
  Load Q,K → Compute S → Write S to HBM
3
  Read S → Softmax → Write P to HBM
4
  Read P,V → Compute O → Write O to HBM
5
每一步跨越慢速 HBM，瓶颈显著。

FlashAttention 1（2022）：Tiling + Online Softmax#

论文：Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022.

核心思想：将注意力矩阵分解为 tile，在 SRAM 内完成关键计算，避免 SSS 和 PPP 写入 HBM。

Online Softmax 数学推导#

标准 softmax 需要两遍扫描——先求最大值做数值稳定，再求和做归一化：

softmax(xi)=exi−m∑jexj−m,m=max⁡kxk\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}, \quad m = \max_k x_ksoftmax(xi)=∑jexj−mexi−m,m=kmaxxk

Online Softmax 将其融合为单次遍历，不断用新的局部最大值修正历史结果。设当前处理第 jjj 个 tile，已累积的全局最大值为 mmm、归一化分母为 ℓ\ellℓ：

mnew=max⁡(m,rowmax(Sij))P~ij=exp⁡ ⁣(Sij−mnew)ℓnew=em−mnew⋅ℓ+∑P~ijOinew=em−mnew⋅Oi+P~ijVj\begin{aligned} m_{\text{new}} &= \max(m, \text{rowmax}(S_{ij})) \\[4pt] \widetilde{P}_{ij} &= \exp\!\left(S_{ij} - m_{\text{new}}\right) \\[4pt] \ell_{\text{new}} &= e^{m - m_{\text{new}}} \cdot \ell + \sum \widetilde{P}_{ij} \\[4pt] O_i^{\text{new}} &= e^{m - m_{\text{new}}} \cdot O_i + \widetilde{P}_{ij} V_j \end{aligned}mnewPijℓnewOinew=max(m,rowmax(Sij))=exp(Sij−mnew)=em−mnew⋅ℓ+∑Pij=em−mnew⋅Oi+PijVj

最终输出为 Oi=Oinew/ℓnewO_i = O_i^{\text{new}} / \ell_{\text{new}}Oi=Oinew/ℓnew。由于 mnew≥mm_{\text{new}} \ge mmnew≥m，修正因子 em−mnew≤1e^{m - m_{\text{new}}} \le 1em−mnew≤1，数值稳定。

PyTorch 实现#

1
def flash_attention_v1(Q, K, V, Br=128, Bc=128):
2
    N, d = Q.shape
3
    O = torch.zeros_like(Q)
4
    m = torch.full((N,), -float('inf'), device=Q.device)
5
    ell = torch.zeros((N,), device=Q.device)
6

7
    for i in range(0, N, Br):
8
        Qi = Q[i:i+Br]
9
        Oi, mi, li = O[i:i+Br], m[i:i+Br], ell[i:i+Br]
10

11
        for j in range(0, N, Bc):
12
            Kj, Vj = K[j:j+Bc], V[j:j+Bc]
13
            Sij = Qi @ Kj.T
14
            mi_new = torch.max(mi, Sij.max(dim=1).values)
15
            Pij = torch.exp(Sij - mi_new.unsqueeze(1))
16
            li = torch.exp(mi - mi_new) * li + Pij.sum(dim=1)
17
            Oi = torch.exp(mi - mi_new).unsqueeze(1) * Oi + Pij @ Vj
18
            mi = mi_new
19

20
        O[i:i+Br] = Oi / li.unsqueeze(1)
21

22
    return O

展开收起

FlashAttention 2（2023）：并行性重构#

论文：Dao, T., FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023.

FlashAttention 1 已经逼近 IO 理论下界，但并行效率仍有提升空间。FlashAttention 2 的改进思路是”让矩阵乘法占据更多 GPU 时间，减少其他操作的开销”。

内外循环交换#

这是 FlashAttention 2 最关键的改动：

1
FlashAttention 1 循环结构:                  FlashAttention 2 循环结构:
2
  for Qi in Q_blocks:  (外循环)              for Kj,Vj in KV_blocks: (外循环,可沿序列并行)
3
    for Kj,Vj in KV_blocks: (内循环)           for Qi in Q_blocks:   (内循环,独立可并行)
4
      Sij = Qi @ Kj^T                           Sij = Qi @ Kj^T

为什么这样改？

FlashAttention 2 的外循环沿 KV 序列维度划分，每个 KV block 交给不同的 thread block 并行处理
内循环的 QiQ_iQi 之间没有依赖，可完全并行
K,VK, VK,V 被加载到 SRAM 后可供多个 QiQ_iQi block 复用

减少非矩阵乘运算占比#

FlashAttention 1 中 softmax 相关的 rescale、mask、dropout 等 element-wise 操作虽然计算量小，但占据了 GPU 调度时间。FlashAttention 2 将这些操作尽可能后置、合并，让 Tensor Core 的 GEMM 时间占比接近 100%。

用公式表示，FlashAttention 2 将 softmax 的归一化步骤延迟到内循环全部结束之后，减少了每次迭代中的除法指令：

FlashAttention 1:Oi←rescale(Oi)+PijVj(每次迭代都 rescale)FlashAttention 2:Oi←Oi+PijVj,最终 Oi/ℓi(仅一次归一化)\begin{aligned} \text{FlashAttention 1:}&\quad O_i \leftarrow \text{rescale}(O_i) + P_{ij}V_j &&\text{(每次迭代都 rescale)} \\ \text{FlashAttention 2:}&\quad O_i \leftarrow O_i + P_{ij}V_j,\quad \text{最终 } O_i / \ell_i &&\text{(仅一次归一化)} \end{aligned}FlashAttention 1:FlashAttention 2:Oi←rescale(Oi)+PijVjOi←Oi+PijVj,最终 Oi/ℓi(每次迭代都 rescale)(仅一次归一化)

注：FlashAttention 3（Shah et al., 2024）的算法架构与 FlashAttention 2 一致，主要针对 Hopper (H100) 架构做了硬件级优化（见”其他优化”部分），此处不单独讨论算法。

FlashAttention 4（2026）：面向非对称硬件扩展的算法-Kernel 流水线协同设计#

论文：Zadouri, Hoehnerbach, Shah et al., FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling, arXiv:2603.05451, March 2026.

背景：Blackwell 架构（B200/GB200）带来了”非对称硬件扩展”——Tensor Core 的算力翻倍，但共享内存带宽、指数函数单元等其他部件的提升不同步甚至停滞。这意味着之前在 Hopper 上不是瓶颈的组件，现在变成了新的瓶颈。

FlashAttention 4 就是围绕如何应对这种”不对称”，重新设计算法-系统的协同优化方案。

三项核心技术#

（1）全异步 MMA 流水线 + 更大 tile

Blackwell 支持完全异步的矩阵乘加（MMA）操作。FlashAttention 4 重新设计了计算流水线，使用更大的 tile 尺寸配合异步执行，让 Tensor Core 在等待数据时不空闲：

Tile 大小从 (Br,Bc)→(Br′,Bc′),Br′≈2×Br\text{Tile 大小从 } (B_r, B_c) \to (B_r', B_c'), \quad B_r' \approx 2\times B_rTile 大小从 (Br,Bc)→(Br′,Bc′),Br′≈2×Br

（2）软件模拟指数函数 + 条件 softmax rescale

Blackwell 的指数函数单元带宽没有成比例增长。FlashAttention 4 用软件模拟的指数函数替代硬件 exp 指令，并结合”条件 rescale”策略——仅在确实需要时才执行 softmax 的修正步骤，减少非矩阵乘运算占比：

rescale 仅在 mnew>m 时触发，否则跳过\text{rescale 仅在 } m_{\text{new}} > m \text{ 时触发，否则跳过}rescale 仅在 mnew>m 时触发，否则跳过

（3）Tensor Memory + 2-CTA MMA 模式

利用 Blackwell 的 Tensor Memory 层级和 2-CTA MMA 模式，将反向传播中的 shared memory 流量和 atomic add 操作大幅减少。两个 CTA（Cooperative Thread Array）共用同一块 MMA 单元，减少了数据搬移开销。

性能#

相比 cuDNN 9.13 快 1.3×
相比 FlashAttention 3 快 2.7×（在同代 Blackwell 硬件上）

其他优化#

以下为各版本中超出核心算法的工程优化技术。

FlashAttention 2 的工程优化#

减少 warp 间通信：forward 中调整计算顺序（Q 与 K/V tile 的配对方式），减少 shared memory 读写，消除跨 warp 的同步等待
Bank conflict 消除：Q 按行优先、K 按列优先存储于 shared memory，避免同一 warp 内的线程访问同一 bank
Causal mask 下三角计算：利用下三角结构跳过被 mask 的 tile，将计算量减半至 ≈N2/2\approx N^2/2≈N2/2

FlashAttention 3 的硬件特性利用（Hopper / H100）#

WGMMA 指令：替代传统 mma.sync，支持 warp group 级别的异步矩阵乘，在共享内存和寄存器之间直接传输
TMA 硬件拷贝：Tensor Memory Accelerator 异步从 HBM 加载数据到 SRAM，完全卸载数据传输开销
GEMM-Softmax 流水线重叠：Tensor Core 做下一 tile 的矩阵乘时，CUDA Core 同时做当前 tile 的 softmax
FP8 低精度：forward 中使用 FP8 Tensor Core 加速矩阵乘，反向传播保持 BF16 精度

FlashAttention 4 的工程优化（Blackwell / B200）#

全异步 MMA 流水线：利用 Blackwell 的异步 MMA 指令，将 GEMM、指数运算、归一化三级流水覆盖，消除等待延迟
软件指数函数模拟：绕过硬件 exp 单元瓶颈，使用多项式近似 + 分段线性拟合替代
条件 softmax rescale：仅在全局最大值变化时才执行修正，减少冗余的除法/指数运算
2-CTA MMA 模式：两个 CTA 共享 MMA 计算单元，减少反向传播中的 atomic add 和 shared memory 流量
Tensor Memory 层级利用：直接使用 Blackwell 的 Tensor Memory 存储中间结果，绕过 shared memory 瓶颈

总结#

	FlashAttention 1	FlashAttention 2	FlashAttention 3	FlashAttention 4
论文标题	FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness	FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning	FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision	FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling
中文标题	具备 IO 感知能力的快速低显存精确注意力机制	通过更好并行性与工作划分实现更快的注意力计算	结合异步执行与低精度计算的快速精确注意力机制	面向非对称硬件扩展的算法与 Kernel 流水线协同设计
GPU 型号	A100 (Ampere)	A100 / H100	H100 (Hopper)	B200 / GB200 (Blackwell)
核心算法改进	IO 感知分块、增量 softmax	内外循环交换、序列维度并行	流水线异步执行、低精度计算	全异步流水线、软件指数模拟、条件 rescale、2-CTA MMA
关键工程优化	SRAM 驻留计算、反向传播重计算	减少 warp 通信、bank conflict 消除、causal 稀疏化	WGMMA 指令、TMA 硬件拷贝、FP8 低精度	异步 MMA 流水线、软件 exp 替代、Tensor Memory 层级利用

参考#

Datawhale llm-algo-leetcode：在线文档 | 代码仓库
FlashAttention 1：Dao, T., et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. doi:10.48550/arXiv.2205.14135
FlashAttention 2：Dao, T. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. doi:10.48550/arXiv.2307.08691
FlashAttention 3：Shah, J., et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. 2024. doi:10.48550/arXiv.2407.08608
FlashAttention 4：Zadouri, T., Hoehnerbach, M., Shah, J., et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. arXiv:2603.05451, 2026. doi:10.48550/arXiv.2603.05451

HaoyangPing