做CUDA算子优化的人迟早会撞上同一个问题明明逻辑没错算力标得极高可一到长序列AttentionGPU利用率就掉到谷底。我第一次在A100上跑自研Attention时就是这个状态——kernel没报错结果也正确但SM占用只有30%左右。后来把Memory Throughput拉出来一看才反应过来瓶颈根本不在计算而在HBM。这就是所谓的GPU内存墙而FlashAttention恰恰是针对这堵墙做的IO感知设计。这篇是这个系列里我写起来最费劲的一篇因为要讲清楚的不只是算法还有它背后的物理直觉。1. 先说清楚内存墙是什么算力翻了三四十倍带宽只翻了不到十倍1.1 Roofline模型先给你的Attention Kernel定性要理解FlashAttention得先接受一个反直觉的事实GPU的瓶颈常常不是算不过来而是数据喂不过来。Roofline模型把这个讲得很明白。一条轴是算术强度单位是FLOP/Byte也就是每个字节的数据从内存进来后你能对它做多少次浮点运算另一条轴是实际能达到的算力。GPU有一个脊点当算术强度高于脊点计算受限你该去优化计算低于脊点带宽受限你优化计算一点用都没有。以我手头常用的A100 80GB为例FP16的Tensor Core理论算力大约312 TFLOPSHBM带宽约2 TB/s脊点大约在156 FLOP/Byte。也就是说如果你的kernel每个字节平均做不了156次浮点运算它就会被压在带宽这一侧。Attention在标准实现下算术强度远低于这个值。QK^T这个操作每读入一对元素能做的运算大约只有d/2量级head维度d常见是64或128也就是32到64 FLOP/Byte离156还有一大截。更不用说后面softmax和PV那几步还带着大量中间矩阵搬运。结论很直接标准Attention是典型的memory-bound kernel优化方向应该全部压在减少HBM访问上而不是压FLOPs。1.2 物理直觉GPU芯片上到底哪里快哪里慢把GPU的存储层次翻译成生活场景这幅画面就生动了。寄存器文件每个SM大概256KB相当于你手上的螺丝刀和扳手随拿随用几乎没有延迟。Shared MemorySMEM每个SM可用几十到一百多KB相当于操作台材料提前摆上来加工就在这完成。L2 Cache几十MB相当于后厨的大柜子拿一次比去仓库快但空间有限。HBM显存几十GB相当于几百米外的冷库。每次去冷库取一趟材料来回可能要大几百拍。GPU算力暴涨靠的是往芯片里塞更多计算单元但冷库到操作台之间的路并没有按同样比例拓宽。于是你写的代码如果频繁回显存取数据再强的计算单元也都在空转。FlashAttention的思路其实特别朴素算得再快也没用关键是减少去冷库的次数能在操作台上完成的事绝不搬回仓库。这个视角一旦建立整个优化方向就清楚了所有工作都应该围绕怎么让数据在片上多复用、少回HBM来设计。2. 标准Attention的HBM账单中间矩阵让数据做了四次无谓搬运2.1 QK^T、Softmax、PV三步各做了什么假设输入是Q、K、V形状都是[N, d]N是序列长度d是head维度。标准Attention分三步计算S Q K^T把N×N的分数矩阵算出来P softmax(S)按行做归一化O P V得到N×d的输出。每一步之间数据都要回到HBM里待着等下一个kernel来取。我们算算HBM的读写量QK^T这一步读Q和K各N×d写回SN×N个元素。Softmax这一步读SN×N写回P又是N×N。PV这一步读PN×N读VN×d写回ON×d。忽略掉d量级的小项总HBM访问大约在4N²这个量级。四次N²的搬运分别对应写S、读S、写P、读P。中间矩阵S和P就是罪魁祸首。有人可能会说N²读写一次也没多大吧我们拿数字说话。N4096head维度128FP16下S和P各是32MB来回搬四次总共128MB看起来还行。但序列长度到16K时S和P各占512MB光这两块中间数据的HBM流量就是2GB。再放大到长文本、多轮对话常用的32K单个head的S矩阵就有2GBP又2GB标准实现直接爆显存都有可能。这就是为什么长序列场景下标准Attention的延迟完全被内存搬运主导。你算没算多少个FLOPs反而不重要了反正HBM带宽就那么宽光搬运中间矩阵就已经把时间吃光了。2.2 为什么中间矩阵必须消除一个N²项引发的灾难Attention的FLOPs是O(N²·d)也就是计算量本身就带有N²项。但N²项在计算层面是必要的毕竟每个query都要跟每个key交互。问题在于标准实现把中间的N×N结果先完整落回HBM再读回来这就让N²项同时出现在IO维度上。更好的做法是什么只在计算的时候产生N²的中间数据但让它留在片上SRAM里不进HBM。让N²项从IO复杂度里消失只保留在FLOPs里。这个思路的本质是把存储中间结果的代价从HBM转移到了SRAM。SRAM没那么大装不下整个N×N矩阵但能装下其中一小块。于是就有了Tiling——把矩阵切成块一块一块在SRAM里算完算完就扔。但这里马上遇到一个麻烦softmax有全局依赖只切一块没法立刻算对。3. FlashAttention的物理直觉把全局Softmax改造成流式合并3.1 核心矛盾Softmax的全局依赖 vs Tiling的局部视角Softmax的公式本身很简短m(x) max(x_i) l(x) sum(exp(x_i - m(x))) softmax(x)_i exp(x_i - m(x)) / l(x)它需要先知道整行的最大值m才能算exp和分母。而在Tiling方案里一个block只看到了Q_i和一部分K_j算出来的只是局部S_tile还没看到整行所有key。你要是直接用局部最大值做了归一化后面再来新的key block之前的结果全都得推翻重算。FlashAttention的聪明之处在于它承认全局最大值目前还不知道但用一个可以持续更新的running state把结果维护住。每来一个key block就更新一次running max和running sum同时把之前已经累加好的输出结果按新的max做一次缩放。所有这一切都发生在SRAM里不写回HBM。这就是Online Softmax。3.2 数学原理running max和running sum如何精确合并我们来看两个block合并的核心公式。假设当前已经处理过一堆key维护了running max m_irunning sum l_i以及暂时的输出O_i。现在来了一个新的key block算出了局部分数S_ij拿到局部行最大值m_ij。新的全局最大值是m_new max(m_i, m_ij)然后P_ij exp(S_ij - m_new)注意是用新max来减。对于旧的那部分结果也不能直接扔掉而是要乘一个缩放系数调整过来scale exp(m_i - m_new)于是running sum更新为l_i l_i * scale rowsum(P_ij)输出更新为O_i O_i * scale P_ij V_j最后所有block处理完O_i再除以l_i得到最终的softmax结果。这里最关键的直觉就是每次看到更大的值就把历史输出整体打个折折扣率就是exp(旧max - 新max)。数学上是精确等价于全局softmax的不是近似。浮点运算顺序变了会有非常微小的数值误差但算法本身不牺牲精度。这是FlashAttention和一堆稀疏注意力、线性注意力最本质的区别。我当时第一次看到这个更新公式时脑子里冒出来的画面是你手里拿了一杯水每隔一段时间水位线上涨新的max变大你就得把杯里已有的水按比例倒掉一点再加新水最后整杯水的比例还是对的。这个画面虽然有点粗糙但比记公式牢固。3.3 复杂度变化为什么N²从IO里消失但FLOPs没变FlashAttention的HBM访问量大概可以近似写成O(N² · d² / M)其中M是SRAM可用容量d是head维度。直观理解整个N×N的分数矩阵被切成(N/B)²个小块每个块要从HBM读入对应的Q、K、V tile大小是B×d量级那么总读取量就是(N/B)² × B × d N²d/B。而SRAM能容纳的tile大小B受M约束大致M ≈ B·d于是得到N²d²/M。跟标准Attention的4N²对比一下。以d64、M大概100KB可用FP16下大约5万元素来估算N²d²/M大概只有0.08N²比标准实现的N²小了一个数量级以上。而且N越长中间矩阵占比越大收益越明显。更值钱的一点是FLOPs一点都没变变的纯粹是数据搬运量。这意味着FlashAttention在memory-bound的Attention上直接把你从带宽受限推到计算受限的边缘算力终于有机会被喂饱了。4. 从数学到CUDA算子Tiling、共享内存与Rescale的代码级细节4.1 为什么必须做算子融合一次Kernel解决所有问题理解了Online Softmax之后自然想到把QK^T、Softmax、PV这三步塞进同一个CUDA Kernel里。算子融合在这里有两个层面的收益。第一个层面是减少中间矩阵的HBM读写。标准做法是三个kernel分开跑每个kernel都要读HBM、写HBM。融合成一个kernel后中间S_tile和P_tile只存在于寄存器和共享内存里根本不出芯片。这是FlashAttention绝大多数加速的来源比省掉两次kernel launch的微秒级开销重要得多。第二个层面是减少启动开销和调度等待。虽然kernel launch的开销每次只有几微秒但Attention在长序列下如果跑几十个kernel累加起来也不小。融合后整个attention流程就一次launch后面所有block之间的协作都在GPU内部完成。4.2 Tile Size怎么定SMEM容量的数学账在写kernel之前必须先算清楚一个block能占多大地方。以A100为例每个SM最多可以配置约164KB的shared memory。假设head维度d64用FP16存储我们来看一个block处理B_r×B_c的tile需要多少空间Q_tileB_r × d × 2 bytesK_tileB_c × d × 2 bytesV_tileB_c × d × 2 bytesS_tile如果放shared memory还需要B_r × B_c × 2 bytes取B_rB_c128d64时Q/K/V三块各16KBS_tile是32KB加上一些统计量总共八九十KBA100的SMEM能扛住。但如果d128Q/K/V各32KBS_tile还是32KB直接逼近甚至超过可用的164KB了这时就必须把B_c降到64。这里有一个很容易忽略的点FlashAttention的官方实现里S_tile一般不放shared memory而是直接放在寄存器的fragment里通过warp协作直接读写。这样能省出32KB的SMEM给K/V块留更多空间或者允许更大的B_c。但代价是逻辑复杂度上升不少register的布局跟Tensor Core的mma指令强相关写起来很容易懵。我的建议是先实现一个纯CUDA Core版本让S_tile安安全全放在shared memory里跑通正确性再考虑Tensor Core和寄存器布局优化。逐级往上改不要一步到位。4.3 Kernel骨架伪代码Rescale的位置最容易写错核心循环长这样// 每个block处理一个Q_tile // 共num_kv_tiles个K/V块需要扫描 for (int j 0; j num_kv_tiles; j) { // 1. 将K_j、V_j异步拷贝到shared memory // 2. 计算 S_tile Q_tile * K_j^T // 3. 求S_tile每行的局部max: m_ij // 4. m_new max(m_i, m_ij) // 合并running max // 5. scale exp(m_i - m_new) // 6. P_tile exp(S_tile - m_new) // 7. l_i l_i * scale rowsum(P_tile) // 8. O_tile O_tile * scale P_tile * V_j // 9. m_i m_new } // 循环结束后: O_tile O_tile / l_i写回HBM这里面最坑的就是第5和第8步的scale。我第一次手写时鬼使神差地用了局部max和旧max的差结果数字歪得离谱。正确的scale一定是用running max m_i和这次合并后的m_new做差不是用S_tile里的max也不是用m_new和m_ij做差。想清楚你要补救的是历史结果这一点就不会错。另一个新手常踩的坑是m_i的初值。有人喜欢初始化为-inf结果第一步exp(-inf)乘到O_tile上变成0×inf直接出NaN。实际写的时候我通常把m_i初始化为一个比较大的负数比如-1e30l_i初始化为0O_tile初始化为0这样第一步的scale是exp(-1e30 - m_ij)在CUDA里正常下溢成0不会产生NaN。这属于很典型的工程细节数学公式里干干净净一落到浮点数上就全是陷阱。4.4 从FlashAttention-1到3每一步都在优化SMEM和流水线FlashAttention原始论文解决的问题是把HBM访问降下来。FlashAttention-2进一步的改进是把Q tile放到内层循环让K/V tile可以连续复用同时减少非矩阵乘法部分对SMEM的占用还顺带重新设计了work partition让并行度更高。这个改进在长序列上很明显因为之前K/V在外层循环时每个K/V块都要被所有Q块重新读一遍HBM流量还是有点冗余。FlashAttention-3则是针对Hopper架构做文章核心是warp specialization和异步流水线。一部分warp专门负责把数据从HBM搬到SMEM另一部分warp专门做矩阵乘法再把结果搬运消费掉让内存搬运和计算重叠起来。本质上也是冲着同一个内存墙去的数据加载的延迟如果能被计算完全掩盖SM的利用率才能拉满。这三个版本的发展路径其实也印证了前面说的物理直觉计算单元和存储带宽之间的鸿沟还在拉大所以每一代优化都在想方设法让数据流更顺畅。5. 实测验证与踩坑用Nsight Compute看数据而不是凭感觉调参5.1 一组实测数字长序列下Flash风格的收益有多大我在A100 80GB上跑过一组对比测试配置是batch size1head数32d128FP16测不同序列长度。对比对象是手写的标准三步kernel每一步单独一个kernel、一个粗糙的融合kernel不做Online Softmax但把中间矩阵留在SMEM、以及FlashAttention风格的精确实现。N1024时标准三步kernel大约0.15msFlash风格大约0.11ms差距不大。这是因为序列短时中间矩阵的N²项还不足以主导HBM流量kernel launch和启动开销占比高。但N8192时差距就拉开了。标准三步kernel大约3.8ms融合kernel大约1.9msFlash风格精确实现大约0.9ms。再往上到N16384标准实现已经要接近15ms了Flash风格只有2.1ms左右。长序列下收益非常可观几乎差了一个数量级。这个趋势背后的原因就是中间矩阵从N²级别暴涨而FlashAttention把N²项从HBM访问里剔除了。序列越长这堵墙越厚绕墙而走的收益自然越明显。5.2 Nsight Compute该看哪些指标如果你写了个Flash风格kernel怎么知道它到底有没有打破内存墙我的经验是开Nsight Compute简称ncu重点看这几个指标首先看DRAM Bytes。运行一次kernel之后ncu会在最前面列出dram__bytes_read和dram__bytes_write的总和。你会发现标准三步实现的DRAM流量大概在N²×8字节量级而Flash风格kernel比它少一个数量级。这一步基本就能确认你remove掉的是不是多余的HBM往返。其次看Memory Throughput vs Compute Workload。如果DRAM throughput已经到80%以上但Tensor Core利用率只有20%说明还在跟内存墙搏斗如果DRAM throughput降下来Compute Workload开始成为主导就说明kernel已经从memory-bound转向compute-bound。最后看SMEM Bank Conflicts和Stall原因。FlashAttention里常见的是矩阵转置导致shared memory bank冲突ncu会给出conflict的bank数量。我遇到过一行数据pad成72个element之后冲突直接消失的情况这种问题靠盲猜是猜不出来的。5.3 我踩过的几个坑列出来供你参考第一Tensor Core的fragment布局别硬啃。直接上m16n8k16的mma指令把Q_tile往寄存器里一放结果shuffle和ldmatrix全乱了。我的建议是先写纯CUDA Core版本用常规的FMA算矩阵乘确保Online Softmax逻辑完全正确再逐步引入Tensor Core。你可以在同一个kernel里用宏开关切换两种实现对比结果。第二FP16累加精度问题。QK^T和PV的累加必须用FP32否则长序列下误差会累积到不可接受。S_tile用FP16存没问题但参与累加时寄存器里一定是FP32。FlashAttention官方实现的数值稳定也在很大程度上依赖于这一点。第三异步拷贝别乱用。FlashAttention-3里的cp.async和warp specialization看起来很美但如果你没把依赖关系理清楚很容易出现K_tile还没拷完就开始计算的race condition。我的习惯是先加足够的__syncthreads()把正确性保护住再用barrier和pipeline逐步替换同步点。先正确再优化最后才谈性能。第四别忽略L2的复用。虽然FlashAttention主打SMEM但L2其实也在暗中帮忙。当一个Q block扫描多个K/V block时K/V块本身在L2里可能会被命中。写代码时保持循环顺序合理让L2多命中几次往往能再省不少HBM带宽。这些收益虽然不如主优化大但属于白捡的便宜。6. 最后再分享一点我个人的体会这个课题我前后啃了小半年。最大的感受是FlashAttention的代码实现确实琐碎但真正难的不是写代码而是建立IO感知的思考方式。你一旦习惯先问这步数据从哪来、到哪去、能不能留在片上再看其他算子优化就会豁然开朗。序列长度超过8K之后Flash风格kernel的收益非常稳定这个已经是生产环境的共识了。但如果你只是偶尔跑一次短序列Attention其实没必要非上这套复杂度标准实现可能就够了。做工程的人都知道选型不是越高级越好而是收益和成本匹配才是真的好。如果你最近也在写自己的Attention kernel我建议先别急着抄完整实现把Online Softmax的数学推导亲手推一遍再用Nsight Compute对比一下改进前后的DRAM流量。这个验证过程本身比任何一篇博客都更能让你记住什么叫内存墙。