FlashAttention：少搬运注意力中间结果
通过分块、融合与重计算，让注意力计算更适应 GPU 存储层次。
相关概念
工程问题#
普通注意力实现物化 N×N 中间矩阵，消耗显存并反复读写 HBM。
实际实现#
将 Q、K、V 分块搬入片上 SRAM，维护分块 softmax 的最大值与归一化统计，累积输出；反向时重算部分中间量。
处理流程#
- 选择可放入片上存储的 tile
- 加载 Q、K、V 块并计算局部分数
- 更新 softmax 统计并合并输出
- 写回结果，反向按需重算
源码片段#
固定版本 · 原文摘录csrc/flash_attn/src/fmha_fprop_kernel_1xN.h765741c1eeb86c96ee71a3291ad6968cfbf4e4a1 · 2022-06-14 · L355–L374
355 // Load over the entire sequence length.356 for( int l = 0; l < steps; l++ ) {357 if((begin + l) * Cta_tile_p::M >= binfo.actual_seqlen) break;358359 // Declare the accumulators for the 1st gemm.360 fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];361 fmha::Clear_accumulator<typename fmha::Accumulator_type, Cta_tile_p::WARPS_K>::apply(acc_p);362363 // Do this part of P = Q * K^T.364 gemm_q_k(acc_p);365366 uint4 out[Gmem_tile_o::STGS_PER_LOOP];367 if (!Is_first) { gmem_o_tmp.load(out, 0); }368369 // Trigger the load for the next Q values.370 if( l < steps - 1) {371 gemm_q_k.smem_q.move_to_next_write_buffer();372 gmem_q.move();373 gmem_q.load();374 }device_1xN_
循环按 CTA tile 处理 Q 与当前 K 块的局部乘积；随后推进 Q 的共享内存写缓冲，并触发下一块加载。softmax 归一化和输出合并在同函数后续完成，不能将片段单独当成完整 Attention。代码来自 2022 年原始实现时期；该文件头单独采用 BSD-3-Clause，不能直接套用仓库根 Apache-2.0。
CPPBSD-3-Clause
论文为 2022 年 arXiv v2；固定代码 commit 为 2022-06-14 的同期实现，不是 FlashAttention-2/3 等后续版本。
收益与取舍#
工程分析带来的收益
- 避免保存完整注意力矩阵
- 降低 HBM 与 SRAM 间的数据搬运
付出的代价
- 重计算增加部分浮点运算
- 分块大小与内核效率受硬件和形状约束
可比较的选择
稀疏或近似注意力可进一步减少计算量，但改变计算模式或近似精度，需要验证模型效果。
版本与适用范围#
原始 FlashAttention 精确注意力设计，数学语义保持不变但浮点舍入可能不同；不将其称为降低稠密注意力算术复杂度的近似算法。
参考来源#
源码、官方文档与原始论文01
FlashAttention 原论文
论文2022 原论文 · arXiv v2§3.1 / Algorithm 1 · Tiling / Recomputation / Online softmax