1. 从一次显存爆炸说起：为什么要啃Flash Attention前向
去年年底我在一台单卡24G的机器上跑一个中等规模的序列建模任务，序列长度拉到4096，batch 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]的矩阵上。当N=4096时，单精度下S矩阵就是64MB，如果batch size是8、头数是16，那就是8GB的中间显存占用，这还没算反向传播需要保存的激活值。
2.2 显存墙与带宽墙的双重夹击
朴素实现的问题不只是显存占用大，更关键的是显存带宽成为了瓶颈。现代GPU的算力增长远远快于显存带宽的增长，以某代数据中心GPU为例，FP16算力能到300 TFLOPS以上，但显存带宽只有不到2 TB/s。这意味着做一次矩阵乘法，如果数据从显存读写的次数太多，算力根本喂不饱。
具体到注意力计算，S矩阵写完要读出来做softmax，softmax的结果P写完又要读出来做第二次矩阵乘法。这中间的读写次数是：写S（N²）、读S（N²）、写P（N²）、读P（N²），总共4N²次显存访问。而实际的计算量只有2N²d次浮点运算。当d=64时，计算访存比只有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。N=4096时，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、warp
CUTLASS是某公司开源的CUDA模板库，专门用来做高性能矩阵乘法。它的抽象层级从上到下是：device层负责管理多个SM上的kernel调度，kernel层对应一个grid，threadblock层对应一个CUDA block，warp层对应32个线程的warp，最底下是thread层。
做Flash Attention前向时，我们主要关注threadblock和warp这两层。threadblock负责加载Q、K、V的分块到shared memory，warp负责具体的矩阵乘法和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或128，Bc取64或128，d取64或128。在CUTLASS里，这些块大小对应到threadblock的tile形状，需要跟warp的tile形状匹配。
4.3 shared memory的分配与双缓冲
shared memory的分配是性能调优的关键。每个threadblock需要缓存Q块、K块、V块，以及中间的S块和P块。以Br=64、Bc=64、d=64为例，Q块是64×64的FP16，占8KB；K块和V块各8KB；S块和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是FP16，S_block和P_block通常用FP32计算以保证精度，但乘V时需要转回FP16。这个转换的时机和位置会影响性能，建议在寄存器里转换，避免多一次shared memory的往返。
6. 性能调优与常见问题排查
6.1 块大小选择的经验法则
块大小的选择没有万能公式，但有一些经验法则。Br和Bc通常取64或128，d取64或128。如果shared memory够大，可以尝试Br=128、Bc=128，这样每个threadblock的计算量更大，kernel launch的开销摊薄。但如果shared memory不够，就要减小块大小，或者增加threadblock的数量。
另一个考虑是warp的数量。每个threadblock通常用4个或8个warp，warp太少并行度不够，warp太多shared memory和寄存器的竞争会加剧。我实测下来，Br=64、Bc=64、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的修正因子是否遗漏 |
| 输出全为0 | exp下溢或上溢 | 打印S_block的最大值和最小值 | 确保减去全局最大值后再exp |
| 性能远低于预期 | shared memory bank conflict | 用profiler查看bank conflict计数 | 给S块和P块加padding |
| 结果有微小误差 | FP16精度不足 | 对比FP32参考实现 | S和P用FP32，O_acc用FP32 |
| kernel launch失败 | shared memory超限 | 查看编译时的shared memory用量 | 减小块大小或减少缓冲级数 |
6.4 跟朴素实现的性能对比
我在同一台机器上对比了朴素实现和Flash Attention前向的性能。序列长度4096、头维度64、batch size 8、16个头的情况下，朴素实现的前向耗时约120ms，显存占用约8GB；Flash 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在寄存器里算完直接做softmax，P_block算完直接乘V，只在必要时才写shared memory，性能才上来。所以能放寄存器就别放shared memory，这是CUDA优化的基本原则。
