登录
首页 >  科技周边 >  人工智能

4比特量化三倍加速不掉点!清华即插即用的SageAttention迎来升级

时间:2024-12-31 22:18:56 165浏览 收藏

各位小伙伴们,大家好呀!看看今天我又给各位带来了什么文章?本文标题《4比特量化三倍加速不掉点!清华即插即用的SageAttention迎来升级》,很明显是关于科技周边的文章哈哈哈,其中内容主要会涉及到等等,如果能帮到你,觉得很不错的话,欢迎各位多多点评和分享!

清华大学陈键飞团队推出SageAttention2:实现4-bit即插即用注意力机制,显著提升大模型推理速度

图片

AIxiv专栏持续报道全球顶尖AI研究成果。 近年来,该专栏已发表2000余篇学术技术文章,涵盖众多高校和企业实验室的先进研究。 欢迎优秀研究者投稿或联系报道 (邮箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com)。

论文共同一作张金涛和黄浩峰分别来自清华大学计算机系和交叉信息研究院,通讯作者陈键飞副教授及其他合作者均来自清华大学计算机系。

大模型线性层的低比特量化已日趋成熟,但注意力模块仍普遍采用高精度计算(如FP16或FP32),尤其在长序列处理中,注意力机制的计算成本日益突出。

此前,陈键飞团队提出的8-bit即插即用注意力机制SageAttention (https://mp.weixin.qq.com/s?__biz=MzA3MzI4MjgzMw==&mid=2650939511&idx=3&sn=dabddded9914c9219df01f93576b1930&scene=21#wechat_redirect),通过将QK^T量化至INT8,保持PV精度为FP16,并结合Smooth K技术,实现了2倍于FlashAttention2的速度提升,同时保持了端到端精度。SageAttention已广泛应用于CogvideoX、Mochi、Flux、Llama3、Qwen等开源及商业大模型。

最新研究成果SageAttention2进一步将注意力机制量化至4-bit,相较于FlashAttention2和xformers分别实现了3倍和4.5倍的即插即用推理加速,并同样在各种大模型上保持了端到端精度。

图片

即插即用特性

SageAttention2提供高效的注意力算子,实现即插即用加速。 只需输入Q, K, V矩阵,即可快速获得注意力输出(O)。 克隆仓库 (git clone https://github.com/thu-ml/SageAttention) 并执行 python setup.py install 后,一行代码即可替换模型中的注意力函数:

图片图片

以CogvideoX-1.5-5B为例,SageAttention2实现了1.8倍的端到端加速,且视频生成质量无损:

4比特量化三倍加速不掉点!清华即插即用的SageAttention迎来升级全精度 Attention 4比特量化三倍加速不掉点!清华即插即用的SageAttention迎来升级SageAttention2

SageAttention2还扩展了硬件支持,在不同GPU平台上均有显著加速效果。

图片

(后续内容,由于篇幅限制,此处省略对前言、挑战、技术方案和实验效果的详细描述,但图片链接保留,读者可自行参考原文深入了解。)

(此处保留原文中所有图片链接)

文中关于工程的知识介绍,希望对你的学习有所帮助!若是受益匪浅,那就动动鼠标收藏这篇《4比特量化三倍加速不掉点!清华即插即用的SageAttention迎来升级》文章吧,也可关注golang学习网公众号了解相关技术文章。

相关阅读
更多>
最新阅读
更多>
课程推荐
更多>