去年我在某个推理服务上排查性能问题时，发现一件特别打脸的事：GPU 利用率显示接近满核，但端到端吞吐就是上不去。监控面板里 HBM 流量和 L2 命中率项亮着红灯，计算核心却有大把周期在空转。查到最后，问题出在我亲手写的一个注意力算子上——按教科书实现，中间矩阵老老实实写回显存，又从显存读回来。那个算子既不算错，也不算慢，但它是最典型的内存墙受害者。
这篇文章是系列第 12 篇。前面讲了不少 GPU 底层调度和算子性能分析的基础，这篇集中把 FlashAttention 和 CUDA 算子优化那层“物理直觉”讲透。网上讲 FlashAttention 的文章很多，但大部分停留在“分块避免 O(N²) 中间矩阵”这个层面，至于为什么要分块、分块为什么偏偏那样分、内核里那些奇怪的循环和 shared memory 布局到底在干嘛，很少有人说清楚。我争取这篇讲完，你能直接照着思路去写自己的融合算子，而不是只搬运现成库。
1. 内存墙的物理本质：搬运比计算贵一个数量级
1.1 GPU 算力增长和带宽增长根本不是一条斜率
先看一组很朴素的数字。每一代 GPU，片上的 FLOPS（浮点算力）基本能做到接近两倍增长，但显存带宽通常只增长 20% 到 50%。这几年 HBM 系列迭代之后带宽确实提上来了，但依然追不上算力的膨胀速度。
这意味着一个很现实的问题：一块卡的“极限吞吐”早就不是由计算峰值决定的，而是由数据搬运速度决定。我们经常说某段代码是内存密集（memory-bound）还是计算密集（compute-bound），判断依据就是算术强度（Arithmetic Intensity）——每个字节数据被读进片上之后，能支撑多少次浮点运算。如果这个值低于机器本身的平衡点，吃得再饱的 ALU 也会因为等数据而空转。
拿当代常见的数据中心 GPU 来算，单卡 FP16 算力动辄几百 TFLOPS，HBM 带宽大概在 1 到 3 TB/s，算下来平衡点大约在 100 到 200 FLOPS/byte。也就是说，你搬一个字节进来，至少要让它干上千一百次运算才算“划算”。可我见过太多算子，尤其是 Transformer 里那些中间步骤，大部分时候搬 1 个字节只干几十次甚至几次运算，算术强度低得可怜，整卡性能自然被死死按住。
这个现象在体系结构里叫“内存墙”（Memory Wall）。这个名字很形象，你不是卡在计算能力，而是卡在 CPU/GPU 核心到主存之间那条“马路”的通行能力上。就好比一个快递分拣中心，分拣机器人一年比一年快，但门口进货车的那条路还是老宽度，那总吞吐量就归道路管，不归机器人管。
1.2 注意力机制为什么会精准踩中这堵墙
Transformer 里的标准注意力公式是：
O = softmax(Q Kᵀ) V
其中 Q、K、V 的形状都是 [N, d]，N 是序列长度，d 是每个头的维度。很多人第一次跑大模型时，注意力模块最容易爆显存，原因就在中间步骤。
Q Kᵀ 算完是一个 [N, N] 的分数矩阵 S。接着 softmax 之后又得到一个 [N, N] 的权重矩阵 P。这两个矩阵都是中间结果，教科书实现会把它完整写到 HBM（显存）里，然后下一个算子再把它从头读一遍。
N 一旦长起来，这个中间矩阵有多恐怖？以 N=4096、batch size=32、32 个头为例，每个头的 P 矩阵是 4096×4096，FP16 下占 32MB。32 个头乘以 32 个 batch 样本，总共就是 32×32×32MB，相当于 32GB 的中间数据。而这些数据只是“中转”了一下，最后还是要被读回来做乘法。
这个例子能说明一个关键现象：标准注意力算子搬运的字节数，和序列长度 N 是平方关系，不是线性关系。中间结果写一次读一次，HBM 流量就是 O(N²) 级别。相比之下，真正的浮点计算量虽然也是 O(N² d)，但中间矩阵的存在让内存流量大到失衡。每算一次矩阵乘法，还要额外搬几份同样大小的中间数据，算力再强也被内存带宽拖死。
2. FlashAttention 的物理直觉：把中间矩阵“焊”在芯片上
2.1 核心思想不是省计算，而是省搬运
很多人第一次听 FlashAttention，以为它像某种数学魔法一样把注意力计算量降低了。实际上它的 FLOPs 几乎没变，注意力本身的数学公式也没变，变的只是中间结果的“居住地”。
原来的做法：中间矩阵 S 和 P，生成之后先送回显存，然后下一次再读出来。显存离计算核心远，来回一次非常贵。
FlashAttention 的做法：把计算切成小块，让每一小块在计算时，相关数据都尽量留在片上 SRAM /寄存器里。算完一小块，结果直接累加到最终输出，中间分数矩阵不离开芯片。
用个生活化的比喻。你要做一顿饭，厨房台面很小，食材仓库在楼下。普通做法是：把一整车食材拉上楼，摆满台面，做完一部分，剩菜剩饭又端下楼放仓库，下一道菜需要时再上楼拿。来回折腾。FlashAttention 的做法是：规划好菜单，每趟只拿当前这一步需要的食材，加工完立即合并到大锅里，台面上永远只有正在处理的这一小批。
2.2 分块 softmax 究竟怎么做到“和原来一样”的
这里最大的疑问通常是：softmax 的分母是全局的，需要对整行所有元素求和，你分块之后怎么保证结果正确？答案是维护两个在线统计量：行最大值 m 和指数和 l。
看伪代码会更清楚：
输入 Q, K, V，形状都是 [N, d] 输出 O，形状是 [N, d] 把 Q 按行切成块 Q_i，每块 B_r 行 把 K、V 按行切成块 K_j、V_j，每块 B_c 行 for i in range(块数): m_i = -inf l_i = 0 O_i = 0 for j in range(块数): S_ij = Q_i @ K_j^T # B_r × B_c m_new = max(m_i, rowmax(S_ij)) P_ij = exp(S_ij - m_new) l_i = l_i * exp(m_i - m_new) + rowsum(P_ij) O_i = O_i * exp(m_i - m_new) + P_ij @ V_j m_i = m_new O_i = O_i / l_i整个过程里，S_ij 和 P_ij 都只存在于当前这一层循环里，用完就释放。而最终的 O_i 除以 l_i 之后，数学上和全局 softmax 的结果完全一致。
为什么一致？因为 softmax 的分子和分母可以分块累加：分子部分是 exp(x - 全局最大值) 的加权和，分母是 exp(x - 全局最大值) 的总和。每次遇到一个新块，先用新的局部最大值把旧统计量缩放到同一尺度，再累加。这相当于你分批统计最高分和总分数，最后归一化，结果自然一样。
还有一点容易被忽略：这个在线算法是数值稳定的。传统 softmax 为了防止 exp 溢出，会减去整行最大值再算。FlashAttention 在分块累加过程中不断更新最大值，同样避免了 exp 爆炸问题。所以它不是“牺牲精度换速度”，而是稳定且精确的。
2.3 复杂度分析：为什么收益是“数量级”的
标准实现里，中间矩阵 S 和 P 分别写一次、读一次，HBM 流量大约正比于 N²。FlashAttention 分块之后，每一块从 HBM 里加载 Q 的一个子块，然后按顺序扫过整列 K、V 块，把需要的部分累加到输出块里。由于每个输出块只需要遍历一遍所有 K/V 块，K/V 从 HBM 被读取的总次数从“每个位置都要读一次”变成“每块数据被不同输出块分时复用”。整体 HBM 流量从 O(N²) 降到大约 O(N²·d/B) 的量级，其中 B 是块大小。
如果 d 取 64 或 128，B 取 128 或 256，这个常数因子非常可观。再加上反向传播环节 FlashAttention 还避免了“保存 S/P 中间矩阵供梯度使用”的需求，而是重算一遍，显存峰值能降一个数量级。这就是为什么它能同时解决“显存不够”和“带宽不够”两个问题。
2.4 用能量视角再看一次
我说这套方案“物理直觉”强，还有一个原因：它能用能耗的视角解释。从 HBM 读一个字节的能量开销，比在 SRAM/寄存器里做一次浮点运算高出一两个数量级，延迟也高一个数量级。优化任何 GPU 算子，本质都是在减少“远端搬运”，而不是斤斤计较多了几条指令。
你去看那些写得好的 CUDA 库，很多算子在寄存器层面大量复用数据，尽可能把数据留在离 ALU 最近的地方。FlashAttention 的核心贡献就是把这种“搬运要少、计算要近”的通用原则应用到了注意力机制上，然后用一套漂亮的分块统计方法把正确性也保证了。
3. CUDA 算子优化落地：从算法到内核里的每个循环
3.1 块大小怎么定：SRAM 容量和 tile 形状的博弈
把 FlashAttention 思想变成 CUDA 内核，第一步是决定每个 block 负责多大的 tile。通常每个线程块负责一个 Q 子块（比如 128×64），内部再循环 K/V 子块（比如 128×64 或 256×64）。
这里的块大小选择不是拍脑袋。shared memory 总共就那么大（常见配置从几十 KB 到两百多 KB），你不仅要放 K、V 子块，还要留出 S_ij 和 P_ij 的空间，甚至要做流水线双缓冲，那就得再翻一倍。块太小则复用率低，每个 K/V 子块会被太多 Q 子块重复读取；块太大则 shared memory 放不下，或者 occupancy 掉得没法看。
经验上是让 K/V 块在 shared memory 中占主流。比如 B_c=256、d=64，FP16 下 K 和 V 各 256×64×2 字节，两个加起来就是 64KB，再乘上双缓冲直接 128KB。这已经接近某些卡的 shared memory 上限。所以很多实现会退回到 B_c=128，甚至用寄存器缓存一部分数据，就是为了在流水线和共享容量之间找平衡。
3.2 共享内存银行冲突：一个让新手怀疑人生的坑
CUDA 内核里计算逻辑对了，性能却可能烂得离谱，最常见的原因之一就是 shared memory bank conflict。
shared memory 被分成 32 个 bank，每个 bank 每个周期可以服务一个线程的访问。如果同一组线程同时访问的地址落在同一个 bank 里，就会发生冲突，硬件只能串行处理，相当于把带宽砍成几分之一。
具体到 FlashAttention 这种以矩阵乘为主的算子，K 块和 V 块会频繁按行读写。一个朴素做法是把 K 块按 [B_c][d] 放进 shared memory，线程按行读取。d 通常是 64，每个 float 4 字节，按行间隔恰好等于 bank 宽度整数倍时，同一列的数据会全部落到同一个 bank 上。于是你内存里看起来连续的数组，实际访问时满屏冲突。
解决手法也经典：给每个行长加 padding，比如从 [B_c][d] 改成 [B_c][d+1]。这样下一行的起始地址偏移了一个字，不同行的同一列就不会再落到同一 bank。成本几乎可以忽略，但性能提升经常超过 30%。
贴一段示意：
// 坑版本：B_c=128, d=64 __shared__ float Ks[128][64]; // 正确版本：尾部补 1 个 float 错开 bank __shared__ float Ks[128][65];看似简单，但很多人内核拆到一半就卡在这种细节上。你要在做完基本优化之后，专门跑一遍 bank conflict 检查，别靠猜。
3.3 软件流水线：把延迟藏起来
内存墙上还有另一半问题：延迟。HBM 访问延迟非常高，线程发出加载指令后，要等很久数据才到。要让计算核心不空等，就得提前把下一块数据加载进来，同时计算当前这一块。
这就是双缓冲（double buffering）加异步拷贝。CUDA 里有cp.async这类指令，可以让全局内存到 shared memory 的拷贝不经过寄存器、不阻塞计算。内层循环里你一边算 K_{j} 和 V_{j} 的乘积，一边让 DMA 异步搬运 K_{j+1} 和 V_{j+1}。
骨架大概长这样：
// 假定循环在 Q_i 块内遍历 K/V 块 for (int j = 0; j < num_tiles; ++j) { // 异步加载下一块到 buf[(j+1) % 2] cp_async_to_shared(buf[(j + 1) % 2], K_global + (j + 1) * tile, bytes_per_tile); // 计算当前 buf[j % 2] 对应的 S/P O compute_block(buf[j % 2]); // 等待当前异步加载完成，保证下一轮计算安全 commit_group_and_wait(); }写这个循环最烦的是同步。异步拷贝发出去之后，你并不知道它什么时候完成。所以要用cp.async.commit_group和cp.async.wait_group来“提交”和“等待”某一批拷贝。只要等待条件设置得准，计算和加载几乎完全重叠。
很多人在自己写融合算子时忽略这一条，觉得反正 shared memory 快，直接同步 load 也行。结果就是内存延迟直接暴露在关键路径上，性能立刻打回原形。正确的软件流水线往往是“最后一公里”性能的关键。
3.4 别只盯着外存，寄存器级的复用才是真功夫
共享内存是离计算单元第二近的存储，但最快的存储其实是寄存器文件。矩阵乘法里最狠的优化是在寄存器里做小块累积，而不是每次从 shared memory 里反复取数。
比如每个线程负责输出一个 8×8 甚至 16×16 的小 tile。加载一份 K、V 数据到寄存器之后，能连续做多次乘加操作，中间再也不碰 shared memory。这种寄存器分块（register tiling）能极大提高算术强度。FlashAttention 的 CUDA 实现中，核心的矩阵乘部分正是依赖这种“寄存器里滚雪球”式的累加。
配合向量化加载，比如一次float4读 16 字节，事务数少、带宽利用率高。你要是只用标量 load，任何内核的带宽都吃不满。说到底，这套优化逻辑和做 GEMM 内核没有本质区别：把数据往寄存器里压，把访存次数压到极限。
4. 真实项目里的分析和调优：数据比感觉可靠
4.1 先找一个让你信服的基线
我改注意力算子时，第一步永远是先把标准实现的性能特征拉出来，而不是直接埋头重写内核。怎么拉？很简单，用分析工具跑一遍并记录几类指标：
- 内核执行时间占比
- DRAM 吞吐和计算吞吐
- shared memory 占用率
- bank conflict 计数
- 全局加载事务数
看清这些数字之后，你才有资格说“瓶颈在内存”还是“瓶颈在计算”。有些时候你以为问题在访存，结果发现是 kernel launch 太密集，光函数调用开销就吃掉一半性能。这种误判特别常见。
以我的经验，标准注意力算子用分析工具测出来的 DRAM 吞吐几乎总是贴着峰值，而计算吞吐远低于峰值。这时候就能下结论：它就是一个纯内存瓶颈算子。后面做的任何优化，目标都应该是减少 DRAM 读写的字节数，而不是减少乘法指令。
4.2 优化前后怎么对照才科学
同一段算法，我建议至少做三版对照。
- 版本 A：标准库或教科书实现
- 版本 B：简单分块注意力（只把中间矩阵放到 shared memory，伪代码完整实现）
- 版本 C：加了异步拷贝、寄存器分块、padding 的完整优化版
不要一上来就写终极版。三版对比能让你清晰看到每一类优化的贡献占比。比如我实测过一个 4096 长度的注意力算子，版本 A 的 DRAM 流量巨大；版本 B 因为把中间矩阵留在片上，DRAM 流量骤降，但 bank conflict 和同步开销明显，性能只提升了不到一倍；版本 C 把流水线做上去之后，才真正拉开数量级差距。
记录数据时，除了跑得快，还要看数值一致性。FlashAttention 分块之后，结果和标准实现的理论误差应该在浮点精度范围内，不是“差不多”，而是“严格等价”。如果你发现误差明显偏大，优先检查在线统计的缩放因子写错没有，比如exp(m_old - m_new)的符号对不对。
4.3 显存峰值下降是意外之喜
很多人只盯着速度，忽略了一个附带收益：FlashAttention 不用存储中间 S/P 矩阵，反向传播时也只需要重算，显存峰值大幅下降。这在长序列场景里可能比速度快更重要。序列长度翻倍，标准实现显存占用直接平方增长，没多久就把显存顶爆了。分块实现让显存占用变成线性级别，很多原来需要模型并行的场景，单卡就能跑。
所以评估这套方案，不要只拿“耗时下降”一个指标说事。把显存峰值、端到端吞吐、甚至功耗都记下来，这些数据会让你的技术方案更有说服力。
5. 常见问题与排查技巧实录
5.1 块越大就一定越快吗
不是。块大小与 shared memory、寄存器数量、线程数强耦合。块太大，shared memory 溢出或 occupancy 下降，线程数变少，矩阵乘所需的并行度就不够，反而更慢。块太小，K/V 被重复读取的次数变多，DRAM 流量回升。正确做法是分析工具配合矩阵尺寸做网格扫描，取一个峰值区间，而不是套用别人内核里的固定值。
5.2 只优化内核，却忘了前后算子
FlashAttention 通常只替换注意力核心，但注意力前后还有 QKV 投影、残差连接、LayerNorm 等等。如果你只把核心算子改了，前后算子依然在频繁读写全局内存，甚至还要做 layout 转换（比如从 [B, H, N, d] 转成 [B, N, H, d]），那整体收益会被吞掉很多。
更彻底的思路是算子融合（kernel fusion），把多个访存密集的小算子合并成一个大算子，让中间数据留在寄存器或 shared memory 里。FlashAttention 本身也只是这套思想下的一个典型例子。
5.3 同步指令用多了，性能照样崩
异步拷贝用起来之后，很多人习惯性在每个循环末尾调用同步等待，结果流水线被打断，加载和计算完全失去重叠。正确做法是让等待发生在“真正需要使用数据”的前一刻。说白了，这种微调需要你静下来算清数据依赖，而不是靠感觉堆同步。
5.4 小序列上过度优化属于浪费感情
FlashAttention 的收益在大序列上才明显。N=128 或者 N=256 这种小规模，中间矩阵也就几万到几十万个元素，搬起来没那么痛。强行上分块和异步拷贝，反而引入额外同步与调度开销。所以做优化前先想清楚你的线上负载到底什么规模，不要为了技术满足感去优化一个根本不在瓶颈上的场景。
5.5 分析工具常见指标一图流
| 指标 | 它告诉你什么 | 怎么用 |
|---|---|---|
| DRAM Throughput | 显存带宽利用率 | 接近 90% 以上，就是内存瓶颈 |
| Compute Throughput | 计算单元利用率 | 接近峰值但 DRAM 不高，才算计算瓶颈 |
| Achieved Occupancy | 活跃线程占比 | 太低就检查 shared memory 和寄存器 |
| Bank Conflicts | shared memory 冲突次数 | 次数高就加 padding 或改布局 |
| Global Load Transactions | 全局加载事务数 | 和理论最小事务数对比，看是否有浪费 |
这套表格基本覆盖了我每次排查 GPU 算子性能时最先看的五个项目。真正干活时，我会再从这些指标倒推定位到具体代码段。
6. 一点个人总结
每次要向别人介绍 FlashAttention，我都不会只讲它“快了多少”，而是强调它把注意力计算从“搬中间矩阵”重构成“在片上在线聚合”的思维转变。这个转变背后是 GPU 体系结构里一直存在的矛盾：算力增长远超带宽增长，数据搬运永远是第一瓶颈。CUDA 优化的本质，其实就是围绕这个矛盾做三件事：减少搬的字节数、让被迫搬的那些字节走最近的路径、在等待数据时做点别的计算。FlashAttention 是这三个原则的完美示范。
如果你在业务里也遇到 GPU 利用率很高但吞吐上不去的怪现象，先别急着怀疑代码写得不够花哨，用分析工具看一眼 DRAM Throughput 和 Compute Throughput 的比例，八成能找到问题上，要么是某个算子没融合，要么是中间张量在显存里来回倒腾。把这类“搬运问题”解决掉，收益往往比换一个更高级的算法模型还要明显。
