企业官网建设流程全解析

做CUDA算子优化的人，迟早会撞上同一个问题：明明逻辑没错，算力标得极高，可一到长序列Attention，GPU利用率就掉到谷底。我第一次在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 TFLOPS，HBM带宽约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 Memory（SMEM）：每个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分三步计算：

1. S = Q @ K^T，把N×N的分数矩阵算出来；
2. P = softmax(S)，按行做归一化；
3. O = P @ V，得到N×d的输出。

每一步之间，数据都要回到HBM里待着，等下一个kernel来取。我们算算HBM的读写量：

• QK^T这一步：读Q和K，各N×d，写回S，N×N个元素。
• Softmax这一步：读S，N×N，写回P，又是N×N。
• PV这一步：读P，N×N，读V，N×d，写回O，N×d。

忽略掉d量级的小项，总HBM访问大约在4N²这个量级。四次N²的搬运，分别对应：写S、读S、写P、读P。中间矩阵S和P就是罪魁祸首。

有人可能会说，N²读写一次也没多大吧？我们拿数字说话。N=4096，head维度128，FP16下S和P各是32MB，来回搬四次总共128MB，看起来还行。但序列长度到16K时，S和P各占512MB，光这两块中间数据的HBM流量就是2GB。再放大到长文本、多轮对话常用的32K，单个head的S矩阵就有2GB，P又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_i，running 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²对比一下。以d=64、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维度d=64，用FP16存储，我们来看一个block处理B_r×B_c的tile需要多少空间：

• Q_tile：B_r × d × 2 bytes
• K_tile：B_c × d × 2 bytes
• V_tile：B_c × d × 2 bytes
• S_tile：如果放shared memory，还需要B_r × B_c × 2 bytes

取B_r=B_c=128，d=64时，Q/K/V三块各16KB，S_tile是32KB，加上一些统计量，总共八九十KB，A100的SMEM能扛住。但如果d=128，Q/K/V各32KB，S_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初始化为一个比较大的负数，比如-1e30，l_i初始化为0，O_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 size=1，head数=32，d=128，FP16，测不同序列长度。对比对象是手写的标准三步kernel（每一步单独一个kernel）、一个粗糙的融合kernel（不做Online Softmax，但把中间矩阵留在SMEM）、以及FlashAttention风格的精确实现。

N=1024时，标准三步kernel大约0.15ms，Flash风格大约0.11ms，差距不大。这是因为序列短时，中间矩阵的N²项还不足以主导HBM流量，kernel launch和启动开销占比高。

但N=8192时，差距就拉开了。标准三步kernel大约3.8ms，融合kernel大约1.9ms，Flash风格精确实现大约0.9ms。再往上到N=16384，标准实现已经要接近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流量。这个验证过程本身，比任何一篇博客都更能让你记住什么叫内存墙。