要解决的问题
注意力计算的瓶颈往往不在浮点运算,而在显存读写:朴素实现会把完整的注意力矩阵写到高带宽显存(HBM)再读回来。如何在不改变计算结果的前提下,把注意力算得更快、更省显存?
核心方法
目前最成功的路线是优化精确注意力的内核,而不是做近似。主要手段:
- 分块计算(tiling):把 Q、K、V 切成小块,放进片上的高速缓存(SRAM)中完成计算,配合在线 softmax,不必写出完整的注意力矩阵。
- 减少 HBM 访问:计算与访存融合,反向传播时重新计算而不是保存中间结果。
- 更好的并行划分:合理分配线程块(thread block)与线程束(warp)之间的工作。
- 利用新硬件与低精度:异步数据搬运、专用矩阵单元、FP8 计算。
代表工作
- FlashAttention(2022):提出 IO 感知的分块精确注意力,显著降低显存占用并提速。
- FlashAttention-2(2023):改进并行划分和线程束之间的工作分配,提高 GPU 利用率。
- FlashAttention-3(2024):针对 Hopper 架构,利用异步执行与 FP8 低精度进一步加速。
价值
FlashAttention 的精确性指没有用稀疏模式或低秩近似替换注意力运算,并不保证浮点结果逐位相同。不同规约顺序可能产生数值差异,FP8 路径还引入量化误差。实际收益取决于硬件、序列长度、头维度和实现,不能把它说成所有场景下都「免费」且优于其他路线。
参考资料
相关专题
- 稀疏注意力与局部注意力:另一条路线:通过近似减少注意力的计算量。
- PagedAttention与KV缓存管理:推理引擎在解码阶段使用的注意力内核需要配合分页的 KV 缓存。
- 优化器、学习率调度与混合精度:低精度计算在训练侧的对应话题。