FlashAttention-4 通俗解读 —— 给大众的硬核科普

一句话抓住重点

大模型每读一句话，都要做一种叫 「注意力（Attention）」的计算——它决定「每个字该参考上下文里的哪些字」。这个计算很费显卡。新显卡（英伟达 Blackwell）变强后出现一个怪现象：「算乘法」的部件快了一倍多，但「算指数」和「搬数据」的部件几乎没变快。于是瓶颈变了。FlashAttention-4 就是专门绕开这两个新瓶颈的重写版，让最强显卡 B200 上的注意力计算又快了 1.3 倍。

读前必看 · 5 个词

先花 1 分钟，把 5 个关键词搞懂

下面这 5 个词会反复出现。先用大白话过一遍，后面就轻松了。

注意力 Attention

大模型理解语言的核心机制。读「它很好吃」时，模型要算出「它」到底指前文哪个东西——给上下文里每个字打一个「相关分」，再按分加权汇总。这就是注意力。

GPU / 显卡

跑大模型的硬件。它内部有很多种「专用工人」：有的专管乘法、有的专管指数、有的专管搬运数据。整体快不快，取决于最忙的那种工人。

张量核 Tensor Core

显卡里专门算「矩阵乘法」的工人，是 AI 算力的主力。新显卡里它被堆得最猛，一代就快了一倍多。

Softmax / 指数运算

把一堆「相关分」换算成「百分比权重」的一步，中间要大量算 $e^x$（指数）。算指数的工人（SFU）很金贵，数量一直没怎么涨。

内存带宽

显卡内部「搬运数据的马路宽度」。工人算得再快，原料运不进来、成品运不出去，照样堵车。这条马路也一直没拓宽。

Kernel / 内核

指一段在显卡上跑的高度优化的程序。FlashAttention 就是一个专门算注意力的 kernel。「优化 kernel」≈ 重新编排工人的活法，让它们别闲着、别打架。

1.3×

比英伟达官方库 cuDNN 更快

2.7×

比开源方案 Triton 更快

71%

把 B200 显卡算力榨到 71%

第 4 代

FlashAttention 系列最新版

01 · 问题的根源

怪事：显卡更强了，老办法却变慢了

从上一代显卡 H100 到新一代 B200，算乘法的「张量核」快了 2.25 倍，但是算指数的工人、搬数据的马路，几乎原地没动。这就好比一家工厂——主力车间产能翻倍，但供料的叉车和打包的工位数量没变。结果就是：原来不起眼的环节，现在成了堵点。

📌 一句话：英伟达把钱主要砸在了「算乘法」上，「算指数」和「搬数据」被晾在一边。这种「一条腿猛长、其他腿不动」的失衡，原文叫 「非对称硬件扩展」。

🚀

算乘法的：快了 2.25 倍

张量核（Tensor Core）大升级，是这一代显卡的主角

🐌

算指数的：原地踏步

算 softmax 要用的 $e^x$ 工人（SFU），数量、速度都没涨

🐌

搬数据的马路：没拓宽

显卡内部「共享内存」的带宽不变，成了反向计算的堵点

🧩

注意力 = 两次乘法夹一个 softmax

先算「相关分」($S=QK^T$)，做 softmax 转成权重，再加权汇总 ($O=PV$)

❌ 最常见的误解

很多人以为「注意力快不快，全看算乘法多快」。在新显卡上恰恰相反：乘法已经快到不是瓶颈了，真正卡脖子的变成了「算指数」（前向计算）和「搬数据」（反向计算）。FA4 的全部聪明之处，就是认清了这个转变。

H100 vs B200 各资源速度对比 — **图1：**新旧两代显卡各类「工人」的速度对比。最左边算乘法的大幅变高，而算指数、搬数据的几乎一样高——一眼看出失衡在哪。

02 · 来龙去脉

四代 FlashAttention：每代都在「追当时最堵的那个点」

这个系列从 2022 年至今出了四代，主线其实特别简单：哪个环节最堵，就重做哪个环节。看懂这条主线，比记住任何术语都重要。

版本	年份 / 显卡	当时最堵的点	它的解法（大白话）
FA1	2022 · A100	反复读写显存太慢	不把中间结果存回显存，分小块在「桌面」上算完再走，速度 2-4 倍
FA2	2023 · A100	工人没排满、有人闲着	更好地分工，减少工人之间的等待与沟通成本
FA3	2024 · H100	算乘法和算 softmax 在排队	用新显卡的「异步」能力，让两件事同时干、不互相等
FA4	2026 · B200	改成「算指数」和「搬数据」堵 ✓	把算指数的活分一部分给闲着的工人，并大幅减少搬运量

💡 我的看法：看懂这一点，就看懂了整个系列

FlashAttention 从来没有「发明新数学」——四代算的都是同一个注意力。它真正做的，是每出一款新显卡，就把算法重新「贴」到新硬件最值钱、最稀缺的资源上。早年显存慢就省显存，后来有了异步指令就抢着重叠，如今乘法过剩、指数稀缺，就把指数的活挪给闲人。「谁稀缺就优化谁」——这一句话能解释所有高性能计算的优化，不只是 AI。

03 · 新显卡给了什么

新显卡 B200 的三件「新工具」

FA4 的很多招数，是直接建立在 B200 新增的硬件能力上的。这里用大白话讲清三件最关键的，后面就不会卡壳。

🗄️

TMEM：紧挨车间的小仓库

新显卡给算乘法的工人配了一块专属的高速暂存区，中间结果就近放，不用挤占工人随身的小口袋（寄存器）

⚡

异步张量核：交代完就走

算乘法可以「下单后不必干等」，同时去张罗别的活，多个任务并行推进

🤝

2-CTA：两个班组拼一单

两个相邻的工作小组合伙完成一次大乘法，各搬一半原料，省掉重复搬运

🐍

用 Python 写显卡代码

FA4 用一种叫 CuTe-DSL 的工具，能用 Python 写底层显卡程序，编译速度比传统 C++ 快 20-30 倍

📌 名词对照：TMEM = 「张量内存」，工人身边的高速小仓库；CTA = 显卡里的一个「工作小组」(一批一起干活的线程)，2-CTA 就是两组合作；寄存器 = 工人随身口袋，又快又小，装不下就得往外挪。

💡 为什么「小仓库 TMEM」是关键

上一代（FA3）只能把中间结果塞进工人随身的小口袋（寄存器），口袋一满，工人就只能一件一件串着干、没法并行。新显卡有了就近的小仓库 TMEM，等于给「同时挂好几件活」松了绑——这是 FA4 能做更深流水线、更大批量的物理前提。

2-CTA 协作示意 — **图2：**「两个班组拼一单」（2-CTA）示意。两个相邻小组把一次大乘法拆开分担，各自只准备一半原料，从而把数据搬运量近乎减半。

04 · 第一个瓶颈

前向计算：敌人是「算指数」

大模型「读」一句话叫前向计算。这一步要做大量 softmax，而 softmax 离不开算指数 $e^x$。偏偏算指数的工人最金贵、又没变多，于是它从「中间一小步」变成了整个计算的堵点。FA4 用了两个巧招。

📌 前向 / 反向是什么：前向 = 模型「读题、给答案」(推理时只有这一步)；反向 = 训练时「对答案、回头改参数」的过程，计算量更大。

巧招一：算指数的活，分一半给闲着的工人

显卡里有一批专算指数的工人（叫 MUFU），人少活多忙不过来；另一批算普通加减乘（叫 FMA）的工人却常常闲着。FA4 的做法是：用闲着的 FMA 工人，靠一道数学近似公式「凑」出指数的结果，和专业工人一起分担，总产能就上去了。

$$ e^x \;\approx\; 2^{n}\cdot\big(p_0 + p_1 f + p_2 f^{2} + p_3 f^{3}\big) $$

大白话：把「算指数」拆成「好算的整数部分 $2^n$」+「用一个小三次多项式凑出来的零头」。整数部分对计算机几乎免费，零头用普通乘加就能算——这样不靠专业工人也能算指数。（系数由数学工具拟合，误差极小）

巧招二：能省则省的「重新校准」

注意力是分小块累加的，每加一块，理论上都要把之前的结果按新的最大值重新校准一次（这是 softmax 的数学要求）。FA4 发现：只有当数值「跳变」很大时才真有必要校准，跳变小的时候直接跳过，反正最后会用真实数据统一校准一次，结果分毫不差。这样就省掉了大量无用功。

📌 打个比方：就像称重时，只有放上很重的东西才需要重新调零，放颗糖就不用每次都校准秤——最后总账一算，结果一样准。

FA4 前向流水线 — **图3：**FA4 前向的「流水线」编排。把「算指数」和「重新校准」这两件累活，从主流程里挪出去交给专门的小组并行处理，主线就不被它们拖慢。

05 · 第二个瓶颈

反向计算：敌人是「搬数据」

训练时的反向计算，活儿比前向多约 2.5 倍。但在新显卡上，卡住它的不是「算不过来」，而是「数据搬不过来」——搬运的马路太窄。原文有个很妙的比喻：优化反向就像往小房间塞一张超大地毯，按平这个角，那个角又翘起来。FA4 从两头同时下手减少搬运。

🗄️

中间结果就近放小仓库

把要反复用的中间数据直接留在 TMEM 小仓库里，省掉来回搬运

🤝

两组拼单，各搬一半

用「2-CTA 拼单」，两个小组各准备一半原料，搬运量直接减半

➗

少做一半「抢着登记」

多个小组写同一份结果时要排队登记（原子操作），既慢又乱；FA4 让登记次数减半

🔄

小组之间「互换半成品」

拼单后两组各缺对方一半数据，就在小组间直接对调补齐，不绕远路

📌 「原子操作」是什么：当多个工人要往同一个本子上记账，必须排队一个一个写、不能同时下笔，否则数字会乱。这种「排队登记」就是原子操作——又慢，而且每次顺序还可能不一样。

⚡ 一个体现功力的细节

FA4 在反向时故意把数据算成「转置」（行列对调）的摆法，让算出来的中间结果，刚好就是下一步乘法需要的摆放格式，省掉一道「翻转重排」的搬运。这种「顺手就把下一步的料备好」的编排，是高性能代码里最见功力的地方。

FA4 反向流水线 — **图4：**反向的流水线编排——在算这一块的同时，提前把上一块的后续乘法发出去，让「算」和「搬」始终都不闲着。

06 · 两个实用功能

额外两个对训练很有用的功能

🎯

「可复现」模式

默认情况下，反向算两遍结果会有微小差异（因为「排队登记」顺序会变）。FA4 提供一个强制固定顺序的模式，让同样输入永远得到同样结果——而且只损失 10-15% 速度

📊

更聪明的「派活」

不同句子长短不一，活儿有多有少。FA4 像食堂排队让活多的先开工，避免「快的早早干完、慢的拖后腿」，整体更快收工

💡 「可复现」为什么越来越重要

过去大家觉得「训练反正有随机性，结果有点抖动无所谓」。但现在做强化学习对齐（就是让模型更听话、更安全的那套训练）、以及排查模型训练为什么崩时，「同样的输入跑两次得到不同结果」会让 bug 极难复现。FA4 把「可复现」做到只损失一成速度，意味着追求确定性，终于不用再牺牲一半性能——这对做大模型训练的团队是实打实的好消息。

07 · 到底快多少

跑分：比官方库还快

在最强显卡 B200 上，和几种主流方案对比，FA4 在「读」（前向）和「训练」（反向）两个方向都领先，尤其在处理长文本时优势明显。

FA4（把显卡算力榨到）71%

比英伟达官方库 cuDNN 快~1.3 倍

比开源方案 Triton 快~2.7 倍

FA4 前向跑分 — **图5：**「读」（前向）的速度对比。横轴是文本长度，纵轴是速度，越高越快。蓝色的 FA4 在各种长度下都稳居最前。

🔥 联网补充：英伟达官方已经「抄作业」了

原文有个很大方的「自曝」：从某个版本起，FA4 团队主动和英伟达官方库 cuDNN 团队合作，把 FA4 的这些招数搬进了官方库。所以最新版 cuDNN 已经能跑出和 FA4 接近的速度。这说明 FA4 真正的价值，不只是「又一个更快的程序」，而是它探明的思路正在被官方吸收，最终让所有用大模型的人都受益——哪怕你根本不知道它的存在。

另外，FA4 也不是永远最快：在某些特殊配置下，反而是上一代 FA2 更优。选哪个，仍要看具体场景。

08 · 一句话总结

这件事真正教会我们什么

💡 三点收获（哪怕你不写代码也值得记住）

① 瓶颈会「搬家」。硬件越堆越猛的今天，卡脖子的环节正从「算得快不快」转向「喂不喂得动」（数据搬运、特殊运算）。这个判断几乎适用于所有现代芯片，不止 AI。

② 闲置的资源就是免费的午餐。FA4 最妙的一招，就是把忙不过来的「算指数」分一部分给闲着的工人。哪里有过剩产能，就把活挪过去——这是最朴素也最有效的优化哲学。

③ 好的优化，是「软件」和「硬件」一起设计。省校准（软件层）+ 用小仓库和拼单（硬件层）必须配合才有效。这就是论文标题里 "co-design（协同设计）" 的分量——单看任何一边都不完整。