发表机构
The Hong Kong University of Science and Technology; Independent Researcher（香港科技大学; 独立研究员）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
针对基于DiT的视频生成模型中注意力机制的计算瓶颈，提出RotateAttention框架，采用选择性FP16回退，引入RoPE感知旋转和范围优化P量化技术，提升速度并保持生成质量。
AI 中文摘要
在配备3D旋转位置嵌入（3D RoPE）的基于DiT的视频生成模型中，注意力机制因二次复杂度成为计算瓶颈。量化FlashAttention虽有前景，但现有低比特量化方法面临与RoPE不兼容及对称量化浪费动态范围等问题。本文提出RotateAttention框架，介绍两项核心技术，实验表明其能保持质量并加速。
英文摘要
In $\textbf{DiT-based video generation models equipped with 3D Rotary Position Embeddings (3D RoPE)}$, the attention mechanism remains a primary computational bottleneck due to its quadratic complexity with respect to sequence length. While quantized $\textbf{FlashAttention}$ offers a promising path toward hardware acceleration, existing low-bit quantization methods overlook two critical challenges in this setting: $\textbf{1)}$ applying online rotation matrices -- a widely used technique for mitigating outliers in Queries ($Q$) and Keys ($K$) -- is difficult to reconcile with $\textbf{RoPE}$; and $\textbf{2)}$ the non-negative attention matrix $P = \exp(QK - \max(QK))$ makes symmetric quantization waste half of the 4-bit dynamic range. In this work, we observe that the outlier distributions of $Q$ and $K$ are strongly affected by the dimensional partitioning of $\textbf{3D RoPE}$. Based on this finding, we propose $\textbf{RotateAttention}$, an efficient $\textbf{mixed-precision INT4 FlashAttention}$ framework tailored for $\textbf{DiT-based video generation models with 3D RoPE}$, using selective $\textbf{FP16 fallback}$ for accuracy-sensitive attention blocks and denoising steps. RotateAttention introduces two core techniques: $\textbf{1) RoPE-aware Rotation}$, which employs either mergeable rotation matrices that can be fused into RoPE or negligible-overhead matrices to mitigate RoPE-induced outliers in $Q$ and $K$; and $\textbf{2) Range-optimized $P$ Quantization}$, which uses fixed scales and zero-points to fully exploit the $\textbf{INT4 numerical range}$ with minimal computational overhead. Experiments show that $\textbf{RotateAttention}$ preserves video generation quality nearly identical to full-precision baselines while achieving up to 1.68$\times$ end-to-end speedup and 2.2$\times$ kernel-level acceleration.
