Flash-Attention

这是一篇硬核的优化Transformer的工作。众所周知,Transformer模型的计算量和储存复杂度是 O ( N 2 ) O(N^2) O(N2) 。尽管先前有了大量的优化工作,比如LongFormer、Sparse Transformer、Reformer等等,一定程度上减轻了Transformer的资源消耗,但对Transformer的性能有所折损,且扩展性不强,不能泛化到其它领域、以及复杂结构的叠加。

这篇工作从底层对Transformer的计算和读写进行了优化,主要有三个贡献:

  1. 加速了模型计算:现在GPU的计算速度已经远远超过了内存读写速度,当GPU完成计算后,内存确还在读取数据,造成GPU闲置而内存繁忙读(消费者早就消费完了,生产者还在缓慢生产)的现象,也就是内存墙问题。FlashAttention通过tiling和算子融合计算,将复杂操作放到SRAM中计算,并减少从HBM读取次数,加快了模型计算速度。而之前的工作虽然减少了Transformer的计算复杂度,却并没有减少模型计算时间。
  2. 节省了显存:FlashAttention通过引入全局统计量,避免实例化大注意力矩阵,减少了显存占用。
  3. 精确注意力:FlashAttention从底层优化了Transformer的计算,但是任务指标上没有任何折损,与普通的Transformer结果是完全等价。

现代GPU内存分级


参考

相关推荐
沛沛老爹5 小时前
Web开发者进阶AI Agent:LangChain提示词模板与输出解析器实战
人工智能·ai·langchain·llm·agent·提示词·web转型
骚戴5 小时前
LLM API Gateway:LLM API 架构、AI 聚合与成本优化全解(2025深度指南)
人工智能·python·大模型·llm·gateway·api
七宝大爷5 小时前
GPU的硬件架构:SM(流式多处理器)剖析
硬件架构·gpu·sm流式多处理器
snoopy_215 小时前
LLM中位置编码
llm
Robot侠14 小时前
极简LLM入门指南4
大数据·python·llm·prompt·提示工程
HyperAI超神经20 小时前
【Triton 教程】triton_language.load
人工智能·学习·大语言模型·cpu·gpu·编程语言·triton
智泊AI1 天前
为什么通用Agent很难实现企业化落地?
llm
大模型教程1 天前
14天速成LLM高手!大佬开源学习笔记,GitHub狂揽700星
程序员·llm·agent
AI大模型1 天前
大模型相关术语和框架总结|LLM、MCP、Prompt、RAG、vLLM、Token、数据蒸馏
程序员·llm·agent
AI大模型1 天前
OpenAI官方出品 : 从0到1构建AI Agent实战指南, 解锁智能自动化新范式
程序员·llm·agent