简介
本文详细介绍了Flash-Attention库从v2.0到v2.7的主要更新,包括v2.2的推理优化、v2.3的滑动窗口注意力、v2.4的ALiBi支持、v2.5的Paged KV cache、v2.6的Softcapping功能以及v2.7的torch compile兼容性。这些更新显著提升了大模型推理效率,支持更长序列和新型模型架构,是大模型开发者必备的技术知识。
在翻阅自己过往笔记的时候发现一个小的变化值得一提,可能业内的大佬都知道 flash-attn 仓库的变化过程。
很久之前尝试在某国产卡上实现 FA kernel 并与官方对比的时候就怎么也对不上精度,只好去读仓库的源码,最后发现 flash-attn-2.0.x 版本与 v2.1 相比,之后 causal mask 的规则变了,如下图所示。
https://github.com/Dao-AILab/flash-attention/tree/v2.8.0
上图中主要分为两个情况,seqlen_q <=seqlen_k 和 seqlen_q>seqlen_k。共同变化就是从注意力矩阵的左上角变成了右下角对齐。
举个简单例子,熟悉 LLM decode 过程的同学,应该理解:在 batchsize=1 时,Query 只有 1 个 token(即当前的第 N+1 个 token),Key 有 N+1 个 tokens(之前缓存的 N 个 k/v cache + 当前的 1 个 token)。
当计算 causal attention 时,当前 query 要与之前的所有 key 计算的,所以这个变化是顺理成章的。
借着这次机会,再整理下截止至 v2.7.0 版本,官方的 README 中都指出了哪些变化。
v2.2: Optimize for inference(split-kv)
针对长文本推理(decoding 阶段)的场景进行优化。我理解的就是这个版本增加了 flash-decoding 的实现。
主要思路是:不再让一个线程块(Block)处理所有的 KV,而是把 KV Cache 在序列维度上切分开。比如,Block 1 处理前 1000 个 KV,Block 2 处理后 1000 个 KV。
最后再用一个额外的 Kernel 把结果合并(Reduce)。可以参考 pytorch 的博客:
https://pytorch.org/blog/flash-decoding/v2.3: Local (Sliding Window) attention (滑动窗口注意力)
适配 Mistral 等模型,如果不原生支持 SWA,FlashAttention 会计算很多无用的 0(被 mask 掉的部分),浪费算力。原生支持后,它只计算窗口内的部分,进一步提升速度并减少显存访问。
v2.4: ALiBi & Deterministic backward
ALiBi:为了外推性(Extrapolation)。有些模型(如 MPT, Bloom)使用 ALiBi 来处理比训练时更长的序列。
FlashAttn 需要支持这种“非标准”的 Attention 计算。Deterministic backward:确定性的反向传播(这里不太理解,可以在评论区补充下)。
v2.5: Paged KV cache (PagedAttention 支持)
支持非连续内存的 KV Cache,即 vLLM 的核心技术 PagedAttention。这里的代码我可以指个路,可以在 forward 时传入 pagedKVcache 的 block_table 信息。至于 block_table,是在框架层面维护的。
https://github.com/Dao-AILab/flash-attention/blob/6f8f0406eea522735d590c2d7b46139167b95b6e/csrc/flash_attn/flash_api.cpp#L1213C44-L1213C56v2.6: Softcapping
适配 Google Gemma-2 和 xAI Grok 模型。这些新模型为了训练稳定性,防止 Attention Score 过大(导致 Softmax 后的梯度消失/熵崩塌),引入了 tanh 截断。
FlashAttn 需要在 kernel 内部融合这个操作,否则用户得切回 PyTorch 原生实现,速度会大打折扣。
v2.7: Compatibility with torch compile
PyTorch 2.0 引入了图编译模式。v2.7 让 FlashAttn 能更好地配合 PyTorch 的计算图捕获,方便用户在生产环境部署时进行全图优化。
暂时结束啦,后续可以在这里持续跟进 flash-attn 仓库的新特性~
如何学习AI大模型?
如果你对AI大模型入门感兴趣,那么你需要的话可以点击这里大模型重磅福利:入门进阶全套104G学习资源包免费分享!
这份完整版的大模型 AI 学习和面试资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】
这是一份大模型从零基础到进阶的学习路线大纲全览,小伙伴们记得点个收藏!
第一阶段:从大模型系统设计入手,讲解大模型的主要方法;
第二阶段:在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用;
第三阶段:大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统;
第四阶段:大模型知识库应用开发以LangChain框架为例,构建物流行业咨询智能问答系统;
第五阶段:大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型;
第六阶段:以SD多模态大模型为主,搭建了文生图小程序案例;
第七阶段:以大模型平台应用与开发为主,通过星火大模型,文心大模型等成熟大模型构建大模型行业应用。
100套AI大模型商业化落地方案
大模型全套视频教程
200本大模型PDF书籍
👉学会后的收获:👈
• 基于大模型全栈工程实现(前端、后端、产品经理、设计、数据分析等),通过这门课可获得不同能力;
• 能够利用大模型解决相关实际项目需求: 大数据时代,越来越多的企业和机构需要处理海量数据,利用大模型技术可以更好地处理这些数据,提高数据分析和决策的准确性。因此,掌握大模型应用开发技能,可以让程序员更好地应对实际项目需求;
• 基于大模型和企业数据AI应用开发,实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能, 学会Fine-tuning垂直训练大模型(数据准备、数据蒸馏、大模型部署)一站式掌握;
• 能够完成时下热门大模型垂直领域模型训练能力,提高程序员的编码能力: 大模型应用开发需要掌握机器学习算法、深度学习框架等技术,这些技术的掌握可以提高程序员的编码能力和分析能力,让程序员更加熟练地编写高质量的代码。
LLM面试题合集
大模型产品经理资源合集
大模型项目实战合集
👉获取方式:
😝有需要的小伙伴,可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓