★ FEATURED ARTICLE 1. 从一次显存爆炸说起为什么要啃Flash Attention前向去年年底我在一台单卡24G的机器上跑一个中等规模的序列建模任务序列长度拉到4096batch size只敢开到2。结果训练到第三个epoch显存直接爆了报错信息里那一长串的中间张量分配记录看得我头皮发麻。当时我第一反应是去查注意力模块的实现发现用的是最朴素的写法先算QK转置softmax再乘V每一步都老老实实把中间结果写回显存。序列长度一上去那个注意力矩阵就是O(N²)的显存占用4096的序列单头单batch的注意力矩阵就是1600万个浮点数多头多batch一叠加显存不炸才怪。后来我花了两周时间把Flash Attention的前向实现从头到尾啃了一遍又结合CUTLASS的模板库自己动手复现了一版。这篇文章就是那段时间的学习记录重点放在前向传播这一块因为前向是理解整个融合思路的入口搞懂了前向反向的推导就是顺水推舟的事。这篇文章适合谁看如果你已经写过基础的注意力模块知道Q、K、V是什么也大致了解softmax的计算流程但一看到Flash Attention论文里那些分块、重计算、在线softmax的公式就犯晕那这篇内容就是为你准备的。我会从最朴素的注意力实现讲起一步步推导到分块融合的版本中间穿插CUTLASS里的核心概念比如tile、warp、shared memory的分配策略。代码层面我会给出关键片段的伪代码和CUDA实现思路但不会贴完整的工程代码因为CUTLASS的模板参数太多贴全了反而看不清主线。先给一个全局的印象Flash Attention前向的核心思想就一句话——把注意力计算拆成小块在片上高速缓存里完成softmax和加权求和避免把完整的注意力矩阵写回显存。听起来简单但要做到高效需要解决三个问题分块之后softmax的归一化因子怎么累积、QK转置的结果怎么在不完整的情况下参与后续计算、以及怎么利用CUTLASS的warp-level原语把矩阵乘法和softmax流水线起来。下面我按自己的理解顺序把这几个问题逐个拆开。2. 朴素注意力实现的瓶颈到底在哪2.1 标准注意力前向的计算流程回顾先快速过一遍标准注意力的前向公式不然后面分块版本的对比就没有参照系。给定输入Q、K、V三个矩阵形状分别是[N, d]、[N, d]、[N, d]其中N是序列长度d是头维度。标准计算分三步第一步计算注意力分数矩阵S Q × K^T形状是[N, N]。这一步是矩阵乘法计算量是O(N²d)。第二步对S的每一行做softmax得到P softmax(S)形状还是[N, N]。softmax需要先求每行的最大值做数值稳定再求指数和做归一化。第三步计算输出O P × V形状是[N, d]。这一步又是矩阵乘法计算量是O(N²d)。整个前向的计算量是O(N²d)显存占用的大头在S和P这两个[N, N]的矩阵上。当N4096时单精度下S矩阵就是64MB如果batch size是8、头数是16那就是8GB的中间显存占用这还没算反向传播需要保存的激活值。2.2 显存墙与带宽墙的双重夹击朴素实现的问题不只是显存占用大更关键的是显存带宽成为了瓶颈。现代GPU的算力增长远远快于显存带宽的增长以某代数据中心GPU为例FP16算力能到300 TFLOPS以上但显存带宽只有不到2 TB/s。这意味着做一次矩阵乘法如果数据从显存读写的次数太多算力根本喂不饱。具体到注意力计算S矩阵写完要读出来做softmaxsoftmax的结果P写完又要读出来做第二次矩阵乘法。这中间的读写次数是写SN²、读SN²、写PN²、读PN²总共4N²次显存访问。而实际的计算量只有2N²d次浮点运算。当d64时计算访存比只有32远低于GPU的算力带宽比所以瓶颈在带宽上。Flash Attention的解决思路就是把中间结果留在片上。GPU的shared memory和寄存器带宽比显存高一个数量级如果把S和P的计算都放在片上完成只把最终的O写回显存那显存访问量就降到O(Nd)级别计算访存比大幅提升。2.3 为什么不能直接把整个S放在shared memory有人可能会想既然shared memory快那把整个S矩阵放进去不就行了。问题在于shared memory的容量有限一块GPU的shared memory通常只有几十KB到一百多KB。N4096时S矩阵是64MB差了三个数量级。所以必须分块每次只处理一小块S算完就丢掉只保留必要的累积结果。这就引出了Flash Attention的核心设计分块计算 在线softmax。分块好理解就是把Q、K、V按行或按列切成小块每次加载一块到片上。在线softmax是难点因为softmax需要全局的最大值和全局的指数和但分块之后你只能看到局部怎么在只看到局部的情况下算出正确的全局softmax结果这是整个算法的关键。3. 在线softmax分块归一化的数学推导3.1 softmax的数值稳定版本先回顾一下数值稳定的softmax怎么算。对于向量x标准softmax是exp(x_i) / sum(exp(x_j))。当x_i很大时exp会溢出所以实际实现会先减去最大值m max(x)变成exp(x_i - m) / sum(exp(x_j - m))。这样指数部分最大是0不会溢出。这个m和分母的sum就是softmax的两个关键统计量。在完整计算时这两个量都是一次性算出来的。但在分块场景下你只能看到一部分x算出来的m和sum都是局部的不能直接用。3.2 分块累积的递推公式假设我们把x分成两块x1和x2分别算出局部的最大值m1、m2和局部的指数和l1、l2。全局的最大值m max(m1, m2)。全局的指数和l l1 * exp(m1 - m) l2 * exp(m2 - m)。这个递推关系就是在线softmax的基础。推广到多块的情况每处理一个新块就用当前块的局部最大值和指数和去更新全局的统计量。更新公式是新全局最大值 m_new max(m_old, m_block)新全局指数和 l_new l_old * exp(m_old - m_new) l_block * exp(m_block - m_new)注意这里l_old需要乘一个修正因子exp(m_old - m_new)因为之前累积的指数和是基于旧的全局最大值算的现在最大值变了需要重新缩放。3.3 输出累积的同步修正softmax的输出是P exp(x - m) / l。在分块计算中输出O P × V也是分块累积的。假设当前已经累积了前几块的输出O_old现在来了新的一块算出局部的P_block和对应的O_block那么新的输出是O_new O_old * (l_old * exp(m_old - m_new) / l_new) O_block * (exp(m_block - m_new) / l_new)这个公式看起来复杂但本质就是旧的输出需要按新的归一化因子重新缩放新的输出按当前块的归一化因子缩放然后相加。实际实现时通常会维护一个未归一化的累积输出O_acc最后再统一除以l这样可以减少中间的重缩放次数。注意在线softmax的递推公式里exp的修正因子是容易写错的地方。我一开始实现的时候忘了给l_old乘修正因子结果训练loss一直不收敛排查了两天才发现是归一化因子算错了。建议在实现时先用小规模数据跟朴素softmax的结果做逐元素对比确认无误后再上大规模。4. CUTLASS视角下的分块矩阵乘法4.1 CUTLASS的层级抽象device、kernel、threadblock、warpCUTLASS是某公司开源的CUDA模板库专门用来做高性能矩阵乘法。它的抽象层级从上到下是device层负责管理多个SM上的kernel调度kernel层对应一个gridthreadblock层对应一个CUDA blockwarp层对应32个线程的warp最底下是thread层。做Flash Attention前向时我们主要关注threadblock和warp这两层。threadblock负责加载Q、K、V的分块到shared memorywarp负责具体的矩阵乘法和softmax计算。CUTLASS提供了很多warp-level的原语比如warp-level的矩阵乘法mma指令、shared memory的加载和存储、以及寄存器之间的数据交换。4.2 tile的划分策略Q按行切K和V按列切Flash Attention前向的分块策略是把Q按行切成若干块每块大小是Br × d把K和V按列切成若干块每块大小是Bc × d。对于每个Q块遍历所有的K块和V块计算注意力分数和加权输出。为什么Q按行切、K和V按列切因为注意力计算是Q的每一行跟K的每一列做点积得到S的一个元素。Q按行切之后每个Q块跟所有的K块都能算出完整的S行。而V是按列切的因为输出O的每一行是P的每一行跟V的所有列做加权求和V按列切之后每个V块贡献输出的一部分列。具体块大小的选择需要权衡。Br和Bc越大片上缓存的利用率越高但shared memory的占用也越大。通常Br取64或128Bc取64或128d取64或128。在CUTLASS里这些块大小对应到threadblock的tile形状需要跟warp的tile形状匹配。4.3 shared memory的分配与双缓冲shared memory的分配是性能调优的关键。每个threadblock需要缓存Q块、K块、V块以及中间的S块和P块。以Br64、Bc64、d64为例Q块是64×64的FP16占8KBK块和V块各8KBS块和P块各8KB。总共32KB加上一些辅助变量大概40KB左右。这在大多数GPU的shared memory容量范围内。双缓冲是指加载下一块K和V的同时计算当前块的S和P。这样可以隐藏显存加载的延迟。CUTLASS的pipeline机制就是做这个的通过cp.async指令实现异步拷贝配合多级缓冲让计算和加载重叠起来。实操心得shared memory的bank conflict是容易被忽略的性能杀手。S块和P块在shared memory里的布局要特别注意如果按行存储warp里的线程访问同一列时会产生bank conflict。解决办法是加padding比如每行多分配一个元素让列访问错开bank。这个padding的大小通常是8或16具体要看访问模式。5. 前向融合kernel的完整实现拆解5.1 整体流程外层循环Q块内层循环K/V块整个前向kernel的结构是两层循环。外层循环遍历Q的块每个Q块加载一次到shared memory然后初始化该块的累积统计量全局最大值m初始化为负无穷全局指数和l初始化为0累积输出O_acc初始化为0。内层循环遍历K和V的块。对于每个K块计算S_block Q_block × K_block^T得到局部的注意力分数。然后对S_block的每一行求局部最大值m_block和局部指数和l_block。接着用在线softmax的递推公式更新全局的m和l并计算P_block exp(S_block - m_new)。最后计算O_acc O_acc * scale_old P_block × V_block其中scale_old是旧输出的缩放因子。内层循环结束后O_acc除以全局的l得到最终的输出O_block写回显存。5.2 QK转置的warp-level实现QK转置的计算是矩阵乘法CUTLASS提供了warp-level的mma指令。以FP16输入为例一个warp可以做16×8×16的mma即16行Q、8列K、16维d。多个warp组合起来覆盖整个threadblock的tile。具体实现时Q块和K块从shared memory加载到寄存器然后调用mma指令。mma的结果是S_block的一部分存在寄存器里。注意S_block的布局要跟后续softmax的访问模式匹配否则寄存器之间的数据交换会很频繁。5.3 softmax的warp-level归约softmax需要求每行的最大值和指数和这是行内归约操作。在warp层面一行S_block可能分布在多个线程的寄存器里需要做warp-level的归约。CUDA提供了__shfl_xor_sync等指令来做warp内的数据交换通过蝶形归约可以在log2(32)5步内完成32个线程的归约。如果一行跨越多个warp还需要跨warp的归约这通常通过shared memory来做。把每个warp的局部归约结果写到shared memory然后一个warp读取所有结果做最终归约再把结果广播回去。5.4 P×V的累积与重缩放P_block × V_block也是矩阵乘法同样用mma指令。但这里有个细节P_block是softmax之后的概率需要先做重缩放再乘V。重缩放的因子是exp(m_block - m_new)这个因子对P_block的每一行是相同的所以可以在寄存器里直接乘。累积输出O_acc的重缩放因子是l_old * exp(m_old - m_new) / l_new这个因子也是每行相同的。实际实现时为了减少除法通常会维护未归一化的O_acc最后再统一除以l。这样重缩放因子就变成l_old * exp(m_old - m_new)少了一次除法。常见坑P_block在乘V之前需要转成跟V匹配的数据类型。如果Q、K、V是FP16S_block和P_block通常用FP32计算以保证精度但乘V时需要转回FP16。这个转换的时机和位置会影响性能建议在寄存器里转换避免多一次shared memory的往返。6. 性能调优与常见问题排查6.1 块大小选择的经验法则块大小的选择没有万能公式但有一些经验法则。Br和Bc通常取64或128d取64或128。如果shared memory够大可以尝试Br128、Bc128这样每个threadblock的计算量更大kernel launch的开销摊薄。但如果shared memory不够就要减小块大小或者增加threadblock的数量。另一个考虑是warp的数量。每个threadblock通常用4个或8个warpwarp太少并行度不够warp太多shared memory和寄存器的竞争会加剧。我实测下来Br64、Bc64、4个warp的配置在多数场景下比较均衡。6.2 数值精度的陷阱Flash Attention前向涉及大量的exp运算和累积求和数值精度容易出问题。FP16的exp范围有限如果S_block的值太大exp之后可能溢出。虽然在线softmax减去了最大值但分块场景下局部最大值和全局最大值的差异可能导致修正因子exp(m_old - m_new)非常小累积时被截断。解决办法是S_block和P_block用FP32计算只在乘V时转FP16。另外累积输出O_acc也用FP32最后写回时再转FP16。这样精度损失最小性能影响也在可接受范围内。6.3 常见问题速查表问题现象可能原因排查方法解决方案loss不收敛在线softmax的修正因子算错用小规模数据对比朴素softmax输出检查l_old的修正因子是否遗漏输出全为0exp下溢或上溢打印S_block的最大值和最小值确保减去全局最大值后再exp性能远低于预期shared memory bank conflict用profiler查看bank conflict计数给S块和P块加padding结果有微小误差FP16精度不足对比FP32参考实现S和P用FP32O_acc用FP32kernel launch失败shared memory超限查看编译时的shared memory用量减小块大小或减少缓冲级数6.4 跟朴素实现的性能对比我在同一台机器上对比了朴素实现和Flash Attention前向的性能。序列长度4096、头维度64、batch size 8、16个头的情况下朴素实现的前向耗时约120ms显存占用约8GBFlash Attention前向耗时约35ms显存占用约1.2GB。加速比约3.4倍显存节省约85%。这个加速比主要来自两个方面一是减少了显存读写二是提高了计算访存比。显存节省则是因为不再需要保存完整的S和P矩阵。序列越长这个优势越明显因为朴素实现的显存占用是O(N²)而Flash Attention是O(N)。7. 从理解到复现我的学习路径建议如果你也想啃下Flash Attention前向我建议按这个顺序来。第一步先把朴素注意力实现写一遍确保能跑通理解每一步的输入输出形状。第二步手动实现一个不分块的在线softmax用两块数据模拟分块累积的过程验证递推公式的正确性。第三步把分块逻辑加上先用小块数据在CPU上模拟确认结果跟朴素实现一致。第四步移植到CUDA先用最简单的实现不考虑性能只求正确。第五步引入CUTLASS的mma指令和pipeline机制逐步优化性能。这个过程中最容易卡住的地方是在线softmax的递推公式和warp-level的归约。前者建议多推导几遍用具体数值代入验证后者建议先看CUTLASS的example理解warp-level原语的用法再自己写。最后分享一个我踩过的坑一开始我为了省事把S_block和P_block都放在shared memory里结果bank conflict严重性能还不如朴素实现。后来改成S_block在寄存器里算完直接做softmaxP_block算完直接乘V只在必要时才写shared memory性能才上来。所以能放寄存器就别放shared memory这是CUDA优化的基本原则。
阅读完成 · 觉得有帮助？
