去年我在某个推理服务上排查性能问题时发现一件特别打脸的事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 算力动辄几百 TFLOPSHBM 带宽大概在 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 一旦长起来这个中间矩阵有多恐怖以 N4096、batch size32、32 个头为例每个头的 P 矩阵是 4096×4096FP16 下占 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 或 128B 取 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_c256、d64FP16 下 K 和 V 各 256×64×2 字节两个加起来就是 64KB再乘上双缓冲直接 128KB。这已经接近某些卡的 shared memory 上限。所以很多实现会退回到 B_c128甚至用寄存器缓存一部分数据就是为了在流水线和共享容量之间找平衡。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][d1]。这样下一行的起始地址偏移了一个字不同行的同一列就不会再落到同一 bank。成本几乎可以忽略但性能提升经常超过 30%。贴一段示意// 坑版本B_c128, d64 __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_{j1} 和 V_{j1}。骨架大概长这样// 假定循环在 Q_i 块内遍历 K/V 块 for (int j 0; j num_tiles; j) { // 异步加载下一块到 buf[(j1) % 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 的收益在大序列上才明显。N128 或者 N256 这种小规模中间矩阵也就几万到几十万个元素搬起来没那么痛。强行上分块和异步拷贝反而引入额外同步与调度开销。所以做优化前先想清楚你的线上负载到底什么规模不要为了技术满足感去优化一个根本不在瓶颈上的场景。5.5 分析工具常见指标一图流指标它告诉你什么怎么用DRAM Throughput显存带宽利用率接近 90% 以上就是内存瓶颈Compute Throughput计算单元利用率接近峰值但 DRAM 不高才算计算瓶颈Achieved Occupancy活跃线程占比太低就检查 shared memory 和寄存器Bank Conflictsshared memory 冲突次数次数高就加 padding 或改布局Global Load Transactions全局加载事务数和理论最小事务数对比看是否有浪费这套表格基本覆盖了我每次排查 GPU 算子性能时最先看的五个项目。真正干活时我会再从这些指标倒推定位到具体代码段。6. 一点个人总结每次要向别人介绍 FlashAttention我都不会只讲它“快了多少”而是强调它把注意力计算从“搬中间矩阵”重构成“在片上在线聚合”的思维转变。这个转变背后是 GPU 体系结构里一直存在的矛盾算力增长远超带宽增长数据搬运永远是第一瓶颈。CUDA 优化的本质其实就是围绕这个矛盾做三件事减少搬的字节数、让被迫搬的那些字节走最近的路径、在等待数据时做点别的计算。FlashAttention 是这三个原则的完美示范。如果你在业务里也遇到 GPU 利用率很高但吞吐上不去的怪现象先别急着怀疑代码写得不够花哨用分析工具看一眼 DRAM Throughput 和 Compute Throughput 的比例八成能找到问题上要么是某个算子没融合要么是中间张量在显存里来回倒腾。把这类“搬运问题”解决掉收益往往比换一个更高级的算法模型还要明显。