先说结论：FlashAttention的价值不只是省显存，它把Attention从“显存搬运型”算子变成了“计算密集型”算子，而工程验证的意义，恰恰在于证明这种数据流重组在不同硬件上真的能稳定兑现理论收益。我最近在给一块老架构加速卡（某厂商的MI50级别GPU）适配FlashAttention时，把数据流优化重新捋了一遍，踩了不少坑，也积累了一些可复用的验证方法。这篇笔记就沿着“为什么这么设计→怎么落到具体硬件→怎么验证它真的对了”这条线展开，希望给做AIInfra或者算子优化相关工作的朋友一些参考。
FlashAttention这两个词在AIInfra圈子里已经被聊了很多，但大多数讨论停留在“它能省显存、能加速”这一层。真正动手适配过就会发现，数据流怎么安排、块大小怎么选、片上内存预算怎么算、数值累加顺序怎么处理，每一点都直接影响最终能不能跑得起来、跑得快不快。
1. FlashAttention为什么值得从数据流角度重新审视
1.1 标准Attention的显存瓶颈其实不在“计算”
标准的Attention计算路径大致是：QK^T算出分数矩阵S，对S做softmax得到P，再用P和V做矩阵乘得到输出O。如果按PyTorch eager模式一步步来，每一步都会把中间结果写回显存。假设序列长度L=4096，head dim d=128，单头单样本的S矩阵就是4096×4096×4字节约64MB，这还只是一个头。实际训练时batch和head数量叠上去，这个中间矩阵在显存里反复写、反复读，带宽消耗是灾难性的。
我用一个很直白的类比来理解：标准Attention像流水线工位，每个工位把半成品搬回仓库，下一个工位再从仓库搬出来加工。中间矩阵就是堆在仓库里的半成品，搬运时间远大于加工时间。对于计算单元来说，ALU在等数据；对于显存来说，带宽被中间结果占满。
1.2 FlashAttention改变了什么：搬运变成就地加工
FlashAttention的核心做法是kernel fusion，把QK^T、softmax、PV合并到一个kernel里，S矩阵和P矩阵不落显存，只在SRAM这种片上内存里按块计算、按块消费。听起来不复杂，但这里有个数学上的关键障碍：softmax需要全局的row max和row sum，而分块计算时每块的局部统计量不能直接决定最终结果。所以FlashAttention引入了online softmax，用带running max和running sum的方式边算边更新，既不需要完整的S矩阵，也不会丢失数值正确性。
从数据流的角度看，这个改动把Attention从“显存搬运型”变成了“计算密集型”。每个数据从显存加载到SRAM后，会被尽量复用到极致，中间不会跑一趟DRAM。
1.3 计算强度的变化是量级级的
Arithmetic intensity这个概念值得引入：它定义为单位字节的运算次数，单位是FLOP/Byte。标准Attention因为中间矩阵反复读写，arithmetic intensity很低，性能上限基本由显存带宽决定。FlashAttention通过分块和融合，把每个字节的利用率拉高，理论上可以把Attention推送到计算单元的算力上限。
我在适配时测过一个直观数据：同样的序列长度和head数，标准eager Attention的kernel耗时里，显存读写占比超过70%，而FlashAttention版本kernel耗时里，计算单元活跃占比才是大头。这个变化意味着，硬件矩阵乘单元越强、片上内存越大，FlashAttention的收益越明显。
提示：评估一个Attention优化方案是否适合某块硬件，先别看峰值算力，先算arithmetic intensity和硬件实际带宽是否匹配。如果方案还是以大量中间结果写显存为代价换计算效率，那在带宽受限的机器上很容易翻车。
2. 数据流拆解：把Attention看成两阶段GEMM的调度问题
2.1 块大小怎么定：先算SRAM预算，再定Tile尺寸
FlashAttention的计算结构，本质上是两个GEMM中间夹一个softmax变换。QK阶段是Q块和K块的GEMM，PV阶段是P块和V块的GEMM。所以工程实现上完全可以复用GEMM领域成熟的tiling思路，但Tile尺寸的设定比普通GEMM更复杂，因为SRAM里要同时放下Q、K、V、S、P和O块，以及online softmax的统计量。
我曾经在一款LDS约64KB级别的加速卡上做适配，head dim是128，如果BLOCK_M和BLOCK_N都取128，光S块就是128×128×4字节约64KB，直接爆掉。后来把BLOCK_M和BLOCK_N都压到64，Q块64×128×2字节=16KB（按fp16输入），K块16KB，V块16KB，S块64×64×4字节=16KB，再加上输出O块和各种中间变量，总预算才勉强压进64KB以内。
这给出一条通用计算式：总预算 = BLOCK_M×d×Q字节 + BLOCK_N×d×(K字节+V字节) + BLOCK_M×BLOCK_N×(S字节+P字节) + BLOCK_M×d×O字节 + 统计量。实际操作中先按这个公式估算，再留出10%~20%余量给寄存器变量和bank padding，否则编译器很容易因为LDS溢出或者寄存器溢出把kernel搞崩。
2.2 循环嵌套顺序：外层遍历Q块，内层遍历KV块
FlashAttention的典型循环结构是：外层遍历序列维度上的Q块，内层遍历对应的KV块。内层循环里Q块驻留在SRAM，K块和V块按顺序加载，计算得分并更新输出。为什么不是外层遍历KV、内层遍历Q？因为外层遍历Q可以让每个Q块只加载一次，而内层KV块是顺序流式读取，能被缓存机制较好利用。反过来会让Q块反复进出SRAM，数据流就乱了。
我见过有人把这个循环顺序搞反，结果性能掉了将近一倍。原因就是LDS的复用率下降，大量重复加载把带宽重新打满。这个顺序问题在纸面上看只是循环嵌套，实际对性能的影响非常直接。
提示：在片上内存里，Q块是“常驻数据”，K/V块是“流式数据”。任何让常驻数据反复加载的设计都是数据流上的失败设计。
2.3 Online softmax的数据依赖：O块必须随时“将就”统计量
FlashAttention的数据流有个隐蔽的依赖关系：当前块的最大值如果比之前的running max大，之前已经累加进O块的数值必须按比例缩放，否则累加结果就错了。具体来说，每次更新running max后，O块要乘上exp(running_max_old - running_max_new)的缩放因子，然后再累加新的P×V贡献。
这个缩放操作让O块不再是个简单的累加器，而是要随时响应统计量变化。工程实现上，这个分支在每次内层循环更新时都会走一遍，虽然计算量不大，但会打断纯GEMM那种线程束均匀执行的节奏。如果代码里对这块处理不谨慎，会造成线程束内divergence，性能损失在5%~10%之间都很正常。
当时我在适配时为了让这个分支更平滑，把O块的缩放操作从逐元素循环改成向量化操作，并且保证所有线程束在执行到该分支前已经同步完成running max的计算，避免部分线程用了旧max、部分线程用了新max。
3. MI50级别硬件适配落地：CK实现与关键参数取舍
3.1 为什么同一份FlashAttention代码换个硬件就“水土不服”
网络上有个热词是“ck flashattention适配mi50”，解释一下背景：CK是某厂商开源的Composable Kernel库，用来拼装高性能算子。MI50级别加速卡属于较老的CDNA架构，和常见游戏卡、新架构加速卡在缓存层次、LDS容量、线程束调度上有很大差异。通用FlashAttention实现往往是针对新架构调过的，直接把编译参数换个架构名，很容易出现LDS溢出、bank conflict爆炸、或者某些内置函数根本不支持的情况。
我这次适配就遇到一个问题：通用实现里默认wavefront是64线程，适配后性能反而低于eager。后来用profiling定位发现，绝大多数耗时耗在LDS bank conflict上。因为S块在按列方向读取时，访问模式正好踩在bank冲突的节奏上，这在新架构上因为显存通路宽不明显，在老架构上直接暴露。
3.2 CK库的价值：把算子拆成可拼装的流水线
CK库把GEMM类算子拆成prologue、body、epilogue三段，FlashAttention的两阶段GEMM结构正好可以复用这套骨架。prologue负责把Q/K/V从显存加载到LDS，body负责核心矩阵乘和softmax统计更新，epilogue负责把O块写回显存。这样抽象有几个好处：加载逻辑、计算逻辑、写回逻辑可以分别优化，不会互相牵扯。
在MI50级别硬件上，我重点调整了prologue的加载宽度。把K块和V块的加载从逐元素改成16字节向量加载，让每次内存请求的字节数贴近硬件的峰值带宽。实测加载阶段的耗时下降了30%以上，整个kernel端到端也有明显提升。
3.3 双缓冲、bank padding、寄存器调度三个关键点
双缓冲是这次优化里收益最大的一项。内层循环加载下一个K块/V块时，计算单元还在处理当前块，如果不做双缓冲，加载和计算必须串行，SRAM和寄存器都在空等。用了双缓冲后，加载下一块和计算当前块重叠起来，kernel的整体耗时基本被计算时间主导，不再被显存延迟拖住。
bank padding是解决上述S块读取冲突的直接手段。做法很简单：给LDS里的行数据尾部加一个占位字（dummy），让行首地址错开bank位，读取列向量时不再所有线程同时打同一个bank。这个操作只牺牲一点点LDS容量，但能解决大量冲突，性价比非常高。
寄存器的分配需要留意。BLOCK_M和BLOCK_N变大时，O块累加器占用寄存器数量会快速上升。如果寄存器溢出到local memory，性能会断崖式下跌。在MI50级别硬件上，我最终把BLOCK_M从128降到64，换来了没有寄存器溢出的稳定性能。这里没有统一答案，只能拿实际profiling数据说话。
// 伪代码示意：FlashAttention内层循环的数据流要点 for (int m = 0; m < num_q_blocks; ++m) { load_Q_to_sram(m); // Q块常驻 running_max = -inf; running_sum = 0; zero_O_block(); for (int n = 0; n < num_kv_blocks; ++n) { load_KV_to_sram(n); // 双缓冲模式下与计算重叠 compute_scores(S_block, Q_block, K_block); update_online_softmax(running_max, running_sum, max_s, sum_s); scale_O_block_by_max_change(); // 关键：O块随max缩放 compute_pv_contribution(O_block, P_block, V_block); } store_O_to_gmem(m); // epilogue }3.4 感性认知：硬件每差一代，优化阈值就变一次
适配MI50级别硬件的体验让我强烈感受到：数据结构设计得再漂亮，最终还是要回到硬件的容量、延迟、带宽、调度规则上来验证。FlashAttention那种“Q块常驻、KV流式”的思想在任何硬件上都成立，但BLOCK_M取多少、LDS padding做到什么程度、向量加载用几字节，必须逐参数试。
我整理了一个简单的适配检查单：
- 确认片上内存容量，按公式估算各块总占用。
- 确认硬件支持的数值类型，不要在bf16支持不全的老卡上硬上。
- 确认wavefront或线程束大小，调整LDS访问步幅，减少bank conflict。
- 确认寄存器堆容量，必要时缩小Tile尺寸。
- 确认是否实现双缓冲，否则内层循环的加载延迟无法隐藏。
每个检查项背后都是实际翻车过的问题。比如数值类型那项，如果老硬件没有原生的bf16计算单元，编译器会生成软件模拟的bf16乘加，速度能慢一个数量级，而且精度损失还不好控制。
4. 工程验证：数值正确性、性能基线定位与避坑实录
4.1 数值正确性验证三件套
写优化算子最怕的是“看起来跑得快，但结果悄悄错了”。FlashAttention因为涉及online softmax、缩放因子和累加顺序变化，数值行为跟标准Attention本来就不同，所以验证要分三层做。
第一层是最大绝对误差和余弦相似度。把FlashAttention的输出和标准Attention的fp32参考输出做对比，用max_abs_diff和cosine_similarity两个指标。我的经验阈值是：fp32下max_abs_diff应小于1e-5量级，fp16下应小于1e-2量级，cosine similarity至少0.999以上。如果误差比这个大很多，先看累加顺序和缩放因子的实现是不是有bug。
第二层是中间统计量的对比。直接打印每块的running max和running sum，跟标准softmax的结果逐块比对。这层能快速定位是online softmax更新逻辑的问题，还是GEMM累加的问题。
第三层是端到端对比。把FlashAttention接进一个真实模型里，跑几步训练或推理，观察loss曲线跟标准Attention是否一致。这一步能发现前两层发现不了的问题，比如梯度计算路径里的数值异常。
4.2 性能验证：别只看端到端时间，先看Roofline位置
性能验证最忌讳拿一个端到端时间就下结论。我习惯先算理论FLOPs，再测kernel实际耗时，得出achieved FLOPs，然后对照硬件的峰值算力看效率。举个例子：sequence length=2048、head数=12、head dim=128、batch=8的Attention，单次forward的计算量大约是2×4×L×L×d×heads×batch≈2×4×2048×2048×128×12×8约41.2GFLOPs。如果kernel耗时是2ms，achieved FLOPs就是约20.6TFLOPS。对照硬件标称fp16算力，如果只有40%左右的利用率，说明数据流还有优化空间；如果达到70%~80%，基本算到位了。
另外要记录显存读写量。可以通过硬件计数器或者profiling工具直接看到。优化前后对比这个数字，能非常直观地验证FlashAttention的核心逻辑有没有真正落地：中间矩阵不落显存的话，读写的量级应该远远低于标准Attention。
4.3 实测中踩过的几个坑
第一个坑：LDS预算超标时编译器不会报错，而是悄悄把数据spill到local memory。性能从“正常”变成“极其拉垮”，但日志里看不出任何异常。只能靠profiling工具去查local memory访问次数，一旦发现明显偏高，立刻减少Tile尺寸。
第二个坑：online softmax的缩放顺序。如果先累加P×V再做缩放，结果误差会随着序列长度递增，长序列下直接偏离。必须严格按照“先缩放O块再累加新贡献”的顺序。这个bug在短序列下几乎测不出来，长序列就会暴露。
第三个坑：老硬件上的bf16支持问题。MI50级别硬件的原生支持是fp16，bf16即使编译器允许编译，执行时也是软件模拟或者低效路径，性能和精度双双受损。因此我在适配时最终保持fp16计算路径，只在必要的统计量上保留fp32累加。
4.4 常见问题速查表
| 症状 | 可能原因 | 解决思路 |
|---|---|---|
| Kernel编译通过但性能极低 | LDS溢出导致spill到local memory | 缩小BLOCK_M/BLOCK_N，腾出LDS余量 |
| 输出误差随序列长度增大 | online softmax缩放顺序错误 | 调整为先缩放O块再累加新P×V |
| Bank conflict严重 | LDS存储布局未做padding | 每行末尾增加占位字，错开bank |
| 端到端速度不升反降 | 旧硬件bf16是软件模拟 | 改用fp16计算+fp32累加 |
| 多线程束结果不一致 | running max更新不同步 | 在更新统计量前做块内同步 |
| 显存读写量没降下来 | 中间S矩阵仍被写回显存 | 检查工程实现是否真的做了kernel fusion |
4.5 性能基线怎么设：先跑通，再调优，最后固化
工程验证的节奏建议是“先跑通，再调优，最后固化”。第一步不追求极致性能，先确保结果数值正确；第二步以标准Attention为基线，逐步调整Tile尺寸、双缓冲、向量加载宽度等参数，每次只改一个变量，并记录对应性能；第三步把最优配置固化成默认参数，并加一个回归测试用例，防止后续改动把性能调回去。
我在这块习惯把每次实验的配置和性能数据记录成表格，而不是只记最终结论。因为很多参数之间存在耦合，比如BLOCK_M调大可能让寄存器压力上升，双缓冲的效果反而被抵消；没有完整实验记录，很难找到最优组合。
5. 三个最容易被忽略的工程细节
5.1 同步粒度：统计量更新前必须做块级同步
online softmax的running max是跨整个线程块共享的。如果有的线程束已经算到新块，有的线程束还在用旧统计量，最终输出就会不一致。这个问题在调试时非常隐蔽，它不会导致编译错误，但会导致结果时好时坏。解决方法是在每次内层循环更新统计量前后，插入块级同步。
我当时排查一个“输出偶发性偏差”的问题用了整整一天，最后才发现是某条路径上漏了一个同步原语。加上之后问题彻底消失。这个经验值得分享：做并行算子时，先保证同步正确，再谈性能优化。
5.2 向量化加载有个“前提”：对齐
向量化加载虽然能提升带宽利用率，但它要求内存地址对齐到向量宽度。如果Q/K/V矩阵在显存里的存储顺序跟Tile访问顺序不一致，贸然用16字节加载轻则性能没提升，重则触发段错误或未定义行为。所以prologue阶段通常要先把数据从全局内存搬运到LDS的连续buffer里，再做向量化访问，而不是直接向量化读全局内存。
5.3 精度策略：fp16算、fp32包
FlashAttention的GEMM累加阶段必须用fp32。fp16累加在长序列和较大head数下误差会迅速累积，训练阶段的梯度更加敏感。我见过一个案例，别人把累加器也改成fp16，单次推理结果看起来还行，但训练时loss曲线发散了。所以累加器用fp32是底线，不能为了速度牺牲这层精度保护。
结尾想说的话
FlashAttention这套数据流优化思想之所以值得反复研究，不只是因为它解决了Attention的带宽瓶颈，而是它示范了一种“面向硬件特性设计算法”的思考方式。在MI50这类老架构硬件上做适配，恰恰能把这种方式逼到极限——你没法依赖新架构的宽松缓存和强大计算单元，必须老老实实算预算、排循环、抠同步。
如果你手头也要做类似适配，我建议先把“数据流”这三个字刻在脑子里：每个数据从哪来到哪去，在片上内存待多久，被复用几次，加载和计算是否重叠。这比直接抄一份代码，然后对着参数表盲目调，要有用得多。最后再分享一个小技巧：在调优遇到瓶颈时，把每个关键参数都做一次“极端取值”压测，比如把BLOCK_M直接翻倍、把向量宽度改成最小，很多时候问题的原因会自己浮现出来。
