论文

RotateAttention:视频生成中 INT4 量化注意力的 RoPE 感知旋转和范围校正

RotateAttention: RoPE-Aware Rotation and Range Rectification for INT4 Quantized Attention in Video Generation

摘要

在配备 3D 旋转位置嵌入 (3D RoPE) 的 $\textbf{基于 DiT 的视频生成模型}$ 中,注意力机制由于其相对于序列长度的二次复杂度,仍然是主要的计算瓶颈。虽然量化的 $\textbf{FlashAttention}$ 为硬件加速提供了一条有希望的途径,但现有的低位量化方法忽略了此设置中的两个关键挑战:$\textbf{1)}$ 应用在线旋转矩阵 - 一种广泛使用的用于减轻查询 ($Q$) 和键 ($K$) 中异常值的技术 - 很难与 $\textbf{RoPE}$ 协调一致; $\textbf{2)}$ 非负注意力矩阵 $P = \exp(QK - \max(QK))$ 使得对称量化浪费了一半的 4 位动态范围。在这项工作中,我们观察到 $Q$ 和 $K$ 的异常值分布受到 $\textbf{3D RoPE}$ 维度划分的强烈影响。基于这一发现,我们提出了$\textbf{RotateAttention}$,一个高效的$\textbf{混合精度INT4 FlashAttention}$框架,专为$\textbf{具有3D RoPE的基于DiT的视频生成模型}$量身定制,使用选择性$\textbf{FP16后备}$用于精度敏感的注意块和去噪步骤。 RotateAttention 引入了两项核心技术:$\textbf{1) RoPE 感知旋转}$,它采用可融合到 RoPE 中的可合并旋转矩阵或可忽略的开销矩阵来减轻 $Q$ 和 $K$ 中由 RoPE 引起的异常值; $\textbf{2) 范围优化$P$量化}$,它使用固定比例和零点以最小的计算开销充分利用$\textbf{INT4数值范围}$。实验表明,$\textbf{RotateAttention}$ 保持了与全精度基线几乎相同的视频生成质量,同时实现高达 1.68$\times$ 端到端加速和 2.2$\times$ 内核级加速。