AI 术语表 · 进阶 · 效率与基础设施
FlashAttention
FlashAttention 是一种计算方法,它按照 GPU 的内存结构重新组织 Transformer 的注意力计算,在结果不变的前提下运行得更快、占用内存更少。
工作原理
注意力会计算句子中所有 token 两两之间的相关程度。如果简单实现,就要反复把巨大的中间表写入 GPU 内存(HBM)再读出来。在 GPU 上,瓶颈往往不是计算本身,而是这种内存往返的时间。FlashAttention 把输入切成小块(tile),在 GPU 芯片内部小而快的内存(SRAM)中完成计算,不把巨大的中间表整体存储下来。关键在于它不是近似,而是得出相同结果的精确注意力。
简单注意力 vs FlashAttention
简单注意力整体存储中间表
- 把大型中间表写入 GPU 内存
- 频繁往返内存
- 输入一长,内存很快耗尽
FlashAttention切成小块在芯片内计算
- 利用快速的芯片内部内存
- 内存往返少
- 能处理更长的输入
공통 两者都是得出相同注意力结果的精确计算
为什么重要
内存占用减少后,可以用更长的上下文进行训练和推理,速度也更快。因此它被视为支撑当今拥有长上下文窗口的 LLM 的技术之一。PyTorch、Hugging Face Transformers 和主要推理引擎都默认使用或以选项形式提供这类经过优化的注意力计算。随着版本迭代,针对新一代 GPU 的改进也在持续。
注意事项
用 Hugging Face Transformers 加载模型时的 attn_implementation 选项,或安装阶段的 flash-attn 包,都会出现这个名字。是否支持取决于 GPU 代际和精度(FP16·BF16 等),在较旧的 GPU 上可能无法安装,或被替换为其他注意力实现。它不是提升模型质量的技术,而是让同样的计算更快、更省内存的技术。在训练或读取长提示词的阶段效果尤其明显,如果以简短对话为主,体感差异可能不大。
示例
假设一个团队用公司内部 GPU 微调长报告摘要模型,刚把输入长度调大就遇到了内存不足错误。团队先确认所用 GPU 代际是否支持 FlashAttention,安装 flash-attn 包,然后在加载模型时更改注意力实现选项。结果不变而内存占用下降,于是可以在同一块 GPU 上用更长的输入继续训练。