1. 这不是一篇讲“Attention有多快”的科普文，而是一份在MI50上跑通FlashAttention的真实工程手记
我第一次在MI50显卡上把FlashAttention的kernel跑起来时，盯着nvidia-smi里那条几乎拉满的GPU Util曲线，心里没想“哇好快”，而是松了口气：“终于没报错”。这和网上那些“FlashAttention提速2倍”的宣传稿完全不同——它根本不是开箱即用的魔法盒，而是一套需要你亲手拧紧每一颗螺丝的精密流水线。AIInfra这个领域最常被忽略的真相是：理论峰值带宽和实际数据吞吐之间，隔着编译器、内存布局、硬件访存模式三道深沟。这篇笔记不谈公式推导，只记录我在MI50上为FlashAttention做数据流优化时，踩过的坑、调过的参数、验证过的每一条数据路径。核心关键词很明确：AIInfra是底座，FlashAttention是载体，数据流优化是动作，工程验证是结果。如果你正打算在老一代AMD GPU（比如MI50）上部署大模型推理，或者需要在有限显存下压榨Attention层的吞吐，又或者正在搭建自己的AI基础设施栈——那你不是来学概念的，你是来抄作业的。下面所有内容，都来自我连续三周在MI50服务器机柜前反复编译、profile、重写kernel的实操现场。
2. 为什么非得在MI50上折腾FlashAttention？数据流瓶颈的物理本质
2.1 MI50不是“过时硬件”，而是暴露数据流问题的绝佳探针
很多人看到MI50（基于Vega架构，16GB HBM2，带宽1TB/s）就默认它“跑不动新模型”，这是个严重误解。MI50的HBM2带宽其实比很多A100的HBM2版本还高（A100 PCIe版是2TB/s，但MI50的1TB/s是在更小die面积下实现的），它的瓶颈从来不在总带宽，而在访存粒度与计算单元的匹配度。FlashAttention的核心思想是把Attention计算拆成多个小块，在SRAM（也就是GPU的L1 cache）里完成softmax归一化和partial sum累积，从而把原本需要反复读写显存的O(N²)访存，压缩成O(N)级。但这个“压缩”能否生效，完全取决于你的硬件是否能让数据真正“留在SRAM里”。
我在MI50上做的第一组对比实验就很说明问题：
- 原生PyTorch的
torch.nn.functional.scaled_dot_product_attention（SDPA）在序列长度8192时，GPU Util只有42%，显存带宽占用率却高达93%； - 同样配置下，直接编译官方FlashAttention v2（针对NVIDIA CUDA优化），根本无法加载——报错
PTX JIT compilation failed，因为MI50不支持PTX 7.5指令集； - 换用ck flashattention（即Compute Kernel FlashAttention，专为AMD ROCm生态重构的版本），编译能过，但首次运行时GPU Util飙升到98%，而实际吞吐反而比PyTorch SDPA低15%。
提示：这不是代码写错了，而是数据流在MI50上发生了“错位”。CK FlashAttention默认按NVIDIA的warp size（32线程）分块，但MI50的wavefront size是64。当一个block试图用32线程协同处理一块QKV数据时，另外32个空闲线程在等，导致计算单元闲置，而显存控制器却在疯狂搬运——这就是典型的“计算等数据”现象。
2.2 数据流优化的本质：让数据“走最短的路”，而不是“算最快的数”
很多人把“优化”等同于“换更快的算法”，但在AIInfra层面，真正的优化是重新设计数据在芯片内部的移动路径。我画了一张MI50上Attention计算的数据流草图贴在工位上，每天看三遍：
[Global Memory (HBM2)] ↓ 读取Q/K/V（非连续stride，需gather） [VGPRs (Vector General Purpose Registers)] ← 这里是第一个瓶颈：MI50的VGPR总量仅128KB，远小于A100的L1 cache（192KB） ↓ 搬入LDS（Local Data Share，相当于shared memory） [LDS: 64KB / CU] ← 第二个瓶颈：LDS带宽虽高，但bank conflict极易发生，尤其当Q/K/V的tile尺寸不是64的整数倍时 ↓ wavefront内64线程并行计算 [SGPRs (Scalar GPRs)] ← 存放block metadata、scale factor等标量参数 ↓ partial softmax + partial sum → 写回LDS或VGPRs ↓ 最终reduction → 写回Global Memory关键发现：MI50的LDS bank数量是32，每个bank宽度64字节。如果Q tile的width设为128（常见于NVIDIA方案），那么每次load会跨2个bank，而K/V tile若用不同width，就会引发bank conflict，实际带宽跌到理论值的35%。这不是软件bug，是硬件物理限制。所以“数据流优化”第一步，就是把所有tile width强制对齐到LDS bank width的整数倍——我们最终选定width=64，牺牲一点计算并行度，换来LDS带宽稳定在82%以上。
2.3 工程验证不是“跑通就行”，而是用三类指标交叉印证
在MI50上验证FlashAttention效果，我坚持用三个维度交叉验证，缺一不可：
- 硬件级指标：
rocm-smi --showmemuse看HBM2实际带宽利用率，rocprof --stats抓wavefront occupancy和LDS bank conflict rate； - 框架级指标：用
torch.utils.benchmark.Timer测单次forward耗时，同时记录torch.cuda.memory_allocated()显存峰值； - 业务级指标：在真实LLM推理pipeline中（我们用的是Llama-2-7B的decoder layer），测端到端token生成延迟（ms/token）和batch吞吐（tokens/sec）。
注意：只看“耗时降低”是危险的。有一次我把tile size从128改成64，forward耗时降了18%，但
rocm-smi显示HBM2带宽利用率从93%降到61%，而rocprof显示LDS bank conflict rate从41%降到8%——这说明优化有效；但另一次我把block size从256调到512，耗时降了22%，HBM2带宽却升到97%，LDS conflict rate跳到53%，这意味着显存成了新瓶颈，所谓“加速”只是把压力从LDS转移到了HBM2，整体能效其实在下降。
3. CK FlashAttention在MI50上的四步数据流重构实操
3.1 第一步：源码级适配——不是编译通过，而是理解每一行汇编
CK FlashAttention的ROCm分支（v1.2.0）提供了MI50支持，但默认配置是为MI250调的。我做的第一件事，是把csrc/flash_attn_rocm/flash_api.cpp里的kernel launch参数全打出来：
// 原始代码（MI250 optimized） int block_size = 256; int grid_size = (seqlen_q + block_size - 1) / block_size; hipLaunchKernelGGL((void*)flash_fwd_kernel, dim3(grid_size), dim3(block_size), 0, 0, args, 0, 0);问题在于：MI50的CU数量是64，每个CU有4个SIMD引擎，每个SIMD引擎处理64线程（wavefront）。当block_size=256时，一个block要占满4个CU，但MI50的LDS是按CU分配的（每个CU 64KB），这就导致LDS碎片化——一个block的tile数据被分散在4个CU的LDS里，跨CU访问LDS会产生额外延迟。
我的修改方案：
- 将
block_size固定为64（一个wavefront），这样每个block严格绑定在一个CU内； grid_size相应放大为(seqlen_q + 63) / 64；- 在kernel里用
__syncthreads()替换__syncthreads_block()，因为MI50的wavefront sync原语更高效。
实测结果：LDS bank conflict rate从37%降至5%，HBM2带宽利用率稳定在68%±3%。
3.2 第二步：内存布局重排——让Q/K/V在HBM2里“肩并肩”站队
MI50的HBM2控制器对连续地址访问最友好。但原始FlashAttention的Q/K/V输入是分别传入的三个tensor，它们在内存里是分离的。我用torch.cat([q, k, v], dim=-1)强行拼成一个大tensor，再在kernel里用指针偏移解包：
# before q, k, v = qkv.unbind(dim=2) # 三个独立tensor # after qkv_cat = torch.cat([q, k, v], dim=-1) # shape: [b, s, 3h] # kernel receives qkv_cat, then: # q_ptr = qkv_cat_ptr # k_ptr = qkv_cat_ptr + s * h * sizeof(float) # v_ptr = qkv_cat_ptr + 2 * s * h * sizeof(float)这个改动看似简单，但效果惊人：HBM2读取QKV的cache line miss rate从28%降到9%。原因在于——MI50的L2 cache line是128字节，而Q/K/V的head_dim通常是128，拼接后每次读取正好填满一个cache line，避免了跨line split。
实操心得：不要迷信“tensor contiguous()”。
qkv_cat.contiguous()在MI50上反而增加一次内存拷贝。正确做法是用qkv_cat = torch.empty(..., device='cuda', dtype=torch.float16)预分配，再用torch.stack([q, k, v], dim=2).transpose(1,2).reshape(...)直接构造，全程零拷贝。
3.3 第三步：LDS使用策略重构——从“大缓存”到“精准管道”
MI50的LDS不是越大越好，而是越“专”越好。原始CK FlashAttention把整个Q/K/V tile都搬进LDS，但MI50的LDS bandwidth是600GB/s，而HBM2是1TB/s——当LDS load成为瓶颈时，它反而拖慢整体。我做了三组LDS使用实验：
| LDS策略 | Q tile搬入 | K tile搬入 | V tile搬入 | L2 cache miss | HBM2带宽利用率 |
|---|---|---|---|---|---|
| 全搬入（默认） | 128x64 | 128x64 | 128x64 | 31% | 93% |
| 只搬Q/K | 64x64 | 64x64 | — | 22% | 76% |
| Q/K分块搬入+V流式 | 32x64 | 32x64 | 不搬 | 14% | 61% |
最终选择第三种：Q和K按32x64分块搬入LDS（刚好占满1个LDS bank），V保持流式读取（因为V只参与一次matmul，不参与softmax）。这样LDS占用从48KB降到12KB，但L2 cache miss率大幅下降，HBM2带宽压力释放，整体吞吐提升23%。
3.4 第四步：混合精度调度——FP16不是终点，而是起点
MI50的FP16 tensor core性能很强，但有个隐藏陷阱：它的FP16 accumulator是32位，而NVIDIA的A100是FP32 accumulator。这意味着在softmax归一化时，FP16累加会产生显著舍入误差。我最初用torch.float16训练的模型，在MI50上用FlashAttention推理时，top-k accuracy掉了0.8%。
解决方案是引入FP32 accumulator for softmax：
- 在kernel里，softmax中间变量（如
exp_scores,sum_exp）用FP32存储； - 但输入Q/K和输出O仍用FP16，保持带宽优势；
- 用
__builtin_amdgcn_fmed3指令做FP16→FP32转换，实测开销<0.3μs。
这个改动让accuracy回归基线，且forward耗时只增加1.2%，属于可接受范围。更重要的是，它暴露了一个AIInfra基本原理：精度策略必须和硬件微架构深度耦合，不能照搬NVIDIA方案。
4. 工程验证全流程：从单kernel到全pipeline的压测方法论
4.1 单kernel验证：用rocprof抓“真实世界”的执行痕迹
很多人用time.time()测kernel耗时，这在MI50上误差极大。正确做法是用ROCm官方工具链：
# 1. 编译时加debug info hipcc -g -O3 -I./csrc/flash_attn_rocm/ ... -o flash_attn.so # 2. 用rocprof抓硬件事件 rocprof --stats --hsa-trace --timestamp --unified-memory --sys-trace \ --hsa-trace --hsa-trace-file hsa_trace.csv \ --stats-file stats.csv \ python test_flash_attn.py # 3. 关键指标解读 # - "Wavefronts"：实际执行的wavefront数量，应接近理论值（grid_size * block_size） # - "LDS Bank Conflicts"：>10%就要调tile size # - "HBM Read/Write Bandwidth"：对比理论带宽（1TB/s），>80%说明显存是瓶颈 # - "VALU Utilization"：计算单元利用率，<60%说明数据没喂饱我遇到过一次诡异问题：rocprof显示Wavefronts=0，但程序没报错。查到最后是HIP runtime版本不匹配——MI50需要ROCm 5.6.0，但我装了5.7.0。降级后问题消失。这再次证明：AIInfra的“工程验证”，首先是环境验证。
4.2 模块级验证：在Llama-2 decoder layer中隔离测试
把FlashAttention嵌入真实模型时，最容易出问题的是梯度回传路径。MI50的HIP autograd engine对自定义kernel的支持不如CUDA成熟。我的验证步骤：
- Forward-only验证：先禁用
torch.no_grad()，确认forward输出数值正确（与PyTorch SDPA比对max diff < 1e-3）； - Backward验证：启用grad，用
torch.autograd.gradcheck验证梯度数值一致性； - Memory Leak验证：循环1000次forward+backward，用
torch.cuda.memory_summary()检查allocated memory是否稳定增长； - Multi-head一致性验证：确保不同head的output variance在合理范围（std < 0.01），排除LDS bank skew。
踩坑记录：MI50上
gradcheck失败率很高，原因是FP16梯度计算的数值不稳定。解决方案是临时将kernel的grad计算切换到FP32，验证通过后再切回FP16——这增加了15%显存占用，但保证了梯度正确性。
4.3 全pipeline压测：用真实负载暴露系统级瓶颈
我们用Llama-2-7B的完整decoder stack（32 layers）做端到端测试，输入序列长度从128到2048递增：
| 序列长度 | PyTorch SDPA (ms/token) | CK FlashAttention (ms/token) | 加速比 | 显存峰值 (GB) |
|---|---|---|---|---|
| 128 | 18.2 | 15.6 | 1.17x | 12.4 |
| 512 | 42.7 | 31.3 | 1.36x | 13.1 |
| 2048 | 158.4 | 92.6 | 1.71x | 14.8 |
关键发现：加速比随序列长度增大而提升，证明FlashAttention的数据流优化在长序列场景下价值最大。但显存峰值也同步上升——这是因为LDS使用策略改变了memory allocation pattern。我们不得不调整torch.cuda.set_per_process_memory_fraction(0.9)，预留更多显存给LDS。
4.4 稳定性验证：72小时无故障运行的“土办法”
在生产环境部署前，我做了72小时压力测试：
- 每5分钟启动一个新batch（batch_size=4, seqlen=1024）；
- 随机注入1%的bad input（如nan in Q）；
- 监控
rocm-smi --showtemp温度（MI50结温上限95℃，我们设定阈值85℃自动降频）； - 记录
dmesg | grep -i "gpu"是否有硬件错误。
结果：72小时内出现2次LDS timeout（MI50硬件bug，ROCm 5.6.0已知问题），触发自动fallback到PyTorch SDPA，服务无中断。这说明：工程验证的终极目标不是“最优”，而是“可控退化”——当优化路径失效时，系统要有确定性的降级方案。
5. 常见问题与避坑指南：MI50上FlashAttention的实战速查表
5.1 编译失败类问题
| 现象 | 根本原因 | 解决方案 | 验证方式 |
|---|---|---|---|
error: unknown type name 'hipStream_t' | HIP SDK路径未加入CMAKE_PREFIX_PATH | export CMAKE_PREFIX_PATH=/opt/rocm | hipconfig --version返回5.6.0 |
undefined reference to 'hipModuleLaunchKernel' | 链接时未加-lhip_hcc | 在setup.py的extra_link_args里添加'-lhip_hcc' | `nm -D flash_attn.so |
PTX JIT compilation failed | kernel用了MI50不支持的指令（如__shfl_down_sync） | 替换为__builtin_amdgcn_ds_permute | 查csrc/flash_attn_rocm/flash_fwd_kernel.hip里无shfl字样 |
5.2 运行时性能异常类问题
| 现象 | 根本原因 | 解决方案 | 验证方式 |
|---|---|---|---|
| GPU Util < 50% 但HBM2带宽 > 90% | 数据加载慢于计算，wavefront大量stall | 减小block_size至64，增加grid_size | rocprof --stats看VALU Utilization是否>75% |
| forward耗时波动大（±20%） | L2 cache污染，不同batch间cache line冲突 | 在每次forward前加torch.cuda.empty_cache() | rocprof --sys-trace看cache miss是否稳定 |
| 多卡训练时loss nan | FP16 accumulator overflow in softmax | 启用FP32 accumulator for softmax | 检查flash_attn_rocm/flash_fwd_kernel.hip中sum_exp类型为float |
5.3 数值精度类问题
| 现象 | 根本原因 | 解决方案 | 验证方式 |
|---|---|---|---|
| top-k accuracy下降 >0.5% | softmax归一化FP16舍入误差累积 | 在kernel中对exp_scores做clamping（exp_scores = min(exp_scores, 60000.0f)） | 用torch.allclose(output1, output2, atol=1e-3)比对 |
| grad norm爆炸 | backward kernel中FP16梯度累加溢出 | backward kernel中dQ,dK,dV用FP32 accumulator | torch.autograd.gradcheck通过率100% |
5.4 系统级稳定性问题
| 现象 | 根本原因 | 解决方案 | 验证方式 |
|---|---|---|---|
| 连续运行24h后GPU hang | MI50的LDS timeout硬件bug（ROCm 5.6.0已知） | 在flash_attnwrapper里加timeout handler，超时自动fallback | watch -n 1 'rocm-smi --showbus'确认PCIe link稳定 |
| 温度突升至90℃+ | 散热硅脂老化，GPU die与散热器接触不良 | 更换液态金属导热膏，压力控制在35psi | rocm-smi --showtemp结温稳定在75℃±3℃ |
实操心得：MI50的“稳定性”不是靠软件调优，而是靠硬件维护。我们给每台MI50服务器加装了机架式风冷增强模块，把进风温度从35℃降到25℃，结温直接降了12℃，LDS timeout概率从每24h 1.2次降到0.1次。
6. 数据流优化的边界在哪里？——来自MI50实践的三条硬经验
我在MI50上做完这套FlashAttention优化后，最深刻的体会不是“技术多精妙”，而是看清了AIInfra里几个血淋淋的现实边界：
第一条经验：没有银弹，只有trade-off。把tile size从128砍到64，LDS bank conflict率从41%降到5%，但计算并行度损失了2倍。我们最终选择64，是因为MI50的LDS带宽瓶颈比计算瓶颈更致命——这和A100的选择完全相反。所以当你看到某篇论文说“XX优化提速2.3倍”，一定要问：在什么硬件上？用什么指标衡量？trade-off是什么？
第二条经验：工程验证的终点不是“跑通”，而是“可解释的失败”。我们花最多时间的不是调优，而是写fallback逻辑：当LDS timeout发生时，自动切到PyTorch SDPA；当FP16 grad overflow时，自动切到FP32 accumulator。真正的工程能力，体现在你能否让系统在任何异常下，都给出确定性的、可追溯的行为。
第三条经验：AIInfra的“基础设施”属性，决定了它必须向后兼容。MI50是2018年的卡，但它还在产线上跑着。我们的FlashAttention patch必须能在ROCm 5.4.0到5.7.0之间无缝工作，因为产线不可能为一个kernel升级整个ROCm栈。所以所有优化都建立在HIP 5.4.0的最小API集上，宁可少用新特性，也要保证向下兼容。
最后分享一个小技巧：在MI50上做任何kernel优化前，先用rocminfo确认你的GPU compute unit数量（Num Compute Units: 64），再用rocm-smi --showhw确认HBM2带宽（HBM Memory Bus Width: 4096-bit），然后打开计算器——把理论带宽（4096/8*1000=512GB/s？不对，MI50是1TB/s，因为HBM2是双通道）算清楚。所有“优化”的起点，都是对硬件物理极限的诚实认知。
