FlashAttention-1、2、3的区别

smartcat2010

已于 2025-02-24 08:28:06 修改

阅读量911

点赞数 12

文章标签： GPU 大模型

于 2025-02-24 08:25:05 首次发布

本文链接：https://blog.youkuaiyun.com/smartcat2010/article/details/145819517

版权

以下是FlashAttention-1、FlashAttention-2和FlashAttention-3的详细对比，涵盖核心目标、算法优化、硬件适配、性能提升等方面的相同点与不同点：

核心目标一致
所有版本均旨在优化Transformer注意力机制的计算效率和内存占用，通过减少对高带宽内存（HBM）的访问，避免存储中间矩阵（如 $S=QK^T$ 和 $P=\text{softmax}(S)$ ），从而解决注意力计算中序列长度的平方复杂度（ $O(N^2)$ ）问题。
分块计算（Tiling）
均采用分块策略将输入矩阵划分为小块（Tile），在GPU的快速SRAM（共享内存）中完成局部计算，通过逐块迭代累积全局结果，避免将完整注意力矩阵写入HBM。
重新计算（Recomputation）
在反向传播时，不存储前向传播的中间矩阵，而是通过保留的统计量（如softmax的归一化系数）重新计算中间结果，以节省显存。
精确注意力结果
所有版本均保持与标准注意力完全相同的计算结果，无需近似或牺牲精度。

核心优化
首次提出分块计算和中间结果融合，将显存占用从 $O(N^2)$ 降低到 $O (N)$ ，并减少HBM读写次数。通过分块Softmax和逐块累积输出矩阵 $O$ 实现高效计算。
并行化策略
在批次（batch）和注意力头（head）维度并行化，但对长序列的并行支持有限，导致GPU利用率较低（如A100上仅达理论FLOPs的25-40%）。
局限性
非矩阵乘法操作（如softmax中的指数和除法）未充分优化，且循环顺序（外层循环为KV，内层循环为Q）导致共享内存访问效率不足。

算法改进
- 减少非矩阵乘法操作：延迟Rescale操作（如softmax分母计算），减少除法运算次数，提升Tensor Core利用率。
- 循环顺序调整：外层循环改为Q，内层循环为KV，减少中间结果的同步和共享内存访问冲突。
- 并行化增强：在序列长度维度并行化，支持更细粒度的任务分配，尤其适合长序列场景（如处理8k以上长度的输入）。
性能提升
前向传播速度提升2倍，H100 GPU的FP16计算利用率达73%，反向传播达63%。
因果掩码优化
通过跳过无效块（如上三角区域），进一步减少计算量，加速因果注意力场景（如自回归生成）。

硬件特性适配
- 异步执行：利用Hopper GPU的Tensor Memory Accelerator（TMA）和Warpgroup异步性，将数据移动与计算重叠，提升吞吐量。
- 低精度支持：引入FP8计算（e4m3格式），结合块量化和非连贯处理（QuIP技术）减少量化误差，FP8计算速度达1.2 PFLOPs/s。
算法创新
- 生产者-消费者流水线：通过Warp专门化，分离GEMM和softmax到不同Warp，实现计算与内存访问的并行。
- 块级GEMM与softmax交错：将softmax的低吞吐操作与GEMM的异步指令重叠，减少等待时间。
性能表现
FP16计算利用率达75%（H100），FP8误差比基线降低2.6倍，适合超长上下文（如64k以上序列）。

维度	FlashAttention-1	FlashAttention-2	FlashAttention-3
核心优化	分块计算、中间结果融合	减少非Matmul操作、循环顺序调整	异步流水线、FP8支持
硬件适配	通用GPU（如A100）	优化Ampere架构	专为Hopper架构（H100）设计
并行化策略	批次和头维度并行	序列长度维度并行	Warp专门化与异步执行
性能指标	显存占用O(N)，速度提升2-4倍	H100前向利用率73%，速度再提升2倍	FP8达1.2 PFLOPs/s，FP16利用率75%
适用场景	常规长度序列（≤8k）	长序列（8k-32k）	超长序列（≥64k）、低精度需求