又快又准，即插即用！清华8比特量化Attention，两倍加速于FlashAttention2，各端到端任务均不掉点！

又快又准，即插即用！清华8比特量化attention，两倍加速于flashattention2，各端到端任务均不掉点！

AIxiv专栏是本站发布学术、技术内容的栏目。过去数年，本站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果+ o .您有优秀的工作想要分享，欢迎投稿o 5 C [ t Z y z A或者联系报道。投u : l稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

论文第一作者张l s % ~ n p W R金涛来自清华大学计算机系，论文通讯作者陈键飞副教授及其他合作作者均来自清华大学计算机系。

大模型中，线性层的低比特量化（例如 Iu B \ 6 = z , K 6NT8, INT4）已经逐步落地；对于注意力模块，目前几乎各个模型都还在用高精度（例如 FP16 或 FP32）的注意力运算进行训练和推理。然而，随着大型模型需要处理的$ ] 8 n p u d序列长度不断增加，Attention（注意力运算）的时间开销逐渐成为网络优s f z 9 i ! l 7化的主要瓶颈。

为了提高注意力运算的效率Z G / 5，清华大学陈键飞团队提出了 8Bit 的 Attention（SageAttention）。实现了 2 倍以及 2.7 倍相比于 FlashAttention2 和 xformers 的即插即用的推理加速，且在视频、图像、# l 3 z F文本生成等大模型上均没有端到端的精度损失。

又快又准，即插即用！清华8比特量化Attention，两倍加速于FlashAttention2，各端到端任务均不掉点！

论文标题：SageAttention: Accurate 8-Bit3 ? 4 R a g # O * Attention for Plug-and-play Inference AcI x / % b 9 vceleration
论文链接：https://ari 9 j – bxiv.org/abs/2410.02367
开源代码：https://github.com/thu-ml/SageAttention

即插即用Z ! i s s举例

SageAttention 可以一行代码轻松替换掉 torch 中当前最优的 Attention 接口（scaled_dot_product_attention），实现即插即用的推理加速。

具体来说，SageAttention 的使用非常( – – Q N 6 9方便，使用 pi– z K Hp install sageattention 后，C ] m R

只需要在模型的推理脚本前加入以下三行代码即可：

效果上，以开源视频生成模型 CogvideoX 为例，使用 SageAttention 可以端到端加速 35%，且生成的视频无损：* l \ $ ( T x

^{全精度 Attention}

^{SageAttention}

接下来，将从背景与挑战，技术方案，以及实验效果介绍 SageAttentionF m D % { 2 S M。

背景

随着大模型需要处理的序列长度越来越长（比如 Llama3.1 支持 128K 的序列长度），Attention 的速度优化? h = w 3 V变得越来越重要。下图展示了一个标准的 Transformer 模型中各运算随着序列长度变化的时间占比：

挑战

为了方便指代注意力元算中包含的矩阵，我们先回顾一下注意力的计算公式：

将神经网Y Z w络中各运算的数值类型从高V ) 5 z r J –比特量化至低比特是一种有效提升计算和访存效率的方法。然而，研究团队发现直@ M v r 4 S r `接将注意力运算中的 Q, K, P,H z n d \ y . Y V 从 FP16 量化为 INT8 或k A # & H 6 ^ 1 ?者 FP8 后将会导致在9 k m J M几乎所有模型和任务上都会得到极差的结果，例如，在 Unidiffusc h c &er 文生图模型中，会得到一张完全模糊的图像；在 Lla; x Vma2-7BW # ^ u K s f [ 进行四选一选择题6 / o s P a d _ j任务上得到 25.5% 的准确率。

经过仔细分析后，研究团队发现主要是两个原因导致了量化注意力的不准确：

大多视频、图像生成模型中，矩阵 K 表现出了极强的通道维度的异常值分布，直接使用 INT8 或者 FP8 数据类型对其进行| k U a M量化会导致巨大的误差。

在所有模型中，对矩阵 P, V 进行量化不能保证一个模型中所有层的精度。下表展示了对 P, V 量e V y w h y化后，Llama2-7B 和 Unidiffuser 模型所有层中，最差情况的层对应的量化注意力的准确度，（该准确度为量化注意9 F y力相比全精度注意力的误差），可以发现不管对 P, V 矩阵进行何种 8Bit （INT8，E4M3，E5M2）量化，] \ *总有些层的准确率非常差c ~ a g \ d ^ X，导致M E 8 0 , T q O了端到端效果的下降。

技术方案

为了解决上述的两个关键问题，研究团队提出了N Y 5对应的解决办法。

对 K 进行平滑处理。SageAttention 采用了一个简单但非常实用的方法来消除矩阵 K 的异常值：K = K – mean (K) 其中 mean (K) 是沿着通道维度求平均值。这个简单的做法不仅不会影响注意力计算的正确性 Softmax (QK^T) = Softmax (Q (K-mean (K))^T) ；且对整个 Attentio8 . c t $ d c in 速度I Z C的影响只有 0.2%；同时还3 \ ^ k 7保证了量化后的注意力运算的精度：

对 Q, K 进行分块 INT8 量化。对于矩阵 Qk ( i ] w $, K，SageAH Z & r ] {ttent1 Z , ! ^ 4ion 采用了以 FlashAttention 的分块大小为粒度的 INT8 量化。这是因为：1. 对 Q, K 矩阵进行 INT8 量化相比于进行 FP8 量化，注意力= . I ~ H \ U (的精度更高。2. 在一些常用卡上，比如 RTX4090，INT8 矩阵乘法（INT32 为累加器）的速度是 FP8（FP32 为累加器）的两倍。
对 P, V 采, & ~ h G z 3 C F用 FP16 数据类型的矩阵乘法累加器。对于矩阵 P, V，SageAtt5 ; r uention 采用了保留 P, V 为X b B + W h FP16 的类型，但进行矩阵乘法时采用 FP16 数据类型的累加器。这是因为：1. PV 矩阵乘法的数值范围始终/ f } I ? X * | #在 FP16 的表示范围内，且经过大量实验验证，FP16 作为累加器的数据类型不会带来任何精度损失（见下表）。2. 在一些常用卡上，比如 RTX4090，以 FP16 为累加. 2 U h : K j Q器数据类型的矩阵乘法的速度是 FP32 作为累加器的两倍。

SageAtten/ ` U 5tion 的流程图及算法如下所示：

实验效果

SageAttention 实现了底层的 GPU Kernel，在算子速度以及各个模型的端到端精度上都有十分不错的表现。

具体来说，算子速度d g ~ p | ;相比于 FlashAttention2 和 xformers 有 2.1 以及4 ] 4 O 2.7 倍的加速。以下 4 张图展示了在 RTX4090 上，不同的序列长度下 SageAttention 的各种 Kernel 与其他方法的速度比较。

以下 4 张图展示了在 RTXE S | @ y3090 上，不同的序列长度下 SageAttention 的各种 Kernel 与其他方法的速, O A f 4度比较。

下表展示了在 RTX4090 上，各模型中的注意力模块中 SageAtteu s v b o Wntion 相比于使用模型原始的注意力的加速比。

真实任务的精度上，下表展示了 SageAttentiop c $ En 在视频、图像、文q m H本生成等大模型上均没有端到端的精度损失：

以上就是又快又准，即插\ { u b B 0即用！清华8比特量化Attention，两倍加速于FlashAttention2，各端到端任务均不掉点！的详细内容！

百度未收录

本文由【好易之】整理自网络！
原创文章，作者:【好易之】如转载请注明出处：https://www.zhengjiaxi.com/zxwd/itzx/118292.html
如有侵权，请邮件联系 aoduoye@qq.com 删除。
本站发布的文章及附件仅限用于学习和研究目的;不得将上述内容用于商业或非法用途，否则后果请用户自负。
本站信息来自网络，版权争议与本站无关。您必须在下载后的24个小时之内，从您的电脑中彻底删除上述内容。
如果您喜欢该程序，请支持正版软件，购买注册，得到更好的正版服务。如有侵权请邮件与我们联系处理。