← 返回归档

SpargeAttention: 准确且无训练稀疏注意力加速任意模型推理

一、论文摘要

随着大模型中序列长度不断增加(如视频生成和语言模型中的45K-128K),注意力机制的二次时间复杂度成为推理延迟的主要瓶颈。幸运的是,注意力图(attention map)本身具有稀疏性,即注意力矩阵中许多值接近于零,这使得跳过相应计算成为可能。许多研究利用这种稀疏模式来加速注意力计算。然而,现有工作大多聚焦于通过利用特定模型的稀疏模式来优化特定模型内的注意力,一种能够保证各种模型加速效果和端到端性能的通用稀疏注意力仍然难以实现。本文提出了SpargeAttention(SpargeAttn),一种适用于任意模型的通用稀疏和量化注意力方法。该方法采用两阶段在线过滤器:第一阶段,我们快速准确地预测注意力图,从而跳过注意力计算中的部分矩阵乘法;第二阶段,我们设计了一种在线softmax感知过滤器,不引入额外开销并进一步跳过部分矩阵乘法。实验表明,我们的方法在语言、图像和视频生成等多种模型上实现了显著加速,且不牺牲端到端指标。代码已开源:https://github.com/thu-ml/SpargeAttn


二、基本信息

论文 ID: 2502.18137

标题: SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference(无训练稀疏注意力加速任意模型推理)

作者: Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, Jianfei Chen

单位:

  • 清华大学计算机科学与技术系、清华-博世机器学习中心、BNRist Center、THBI Lab、Institute for AI (Tsinghua University)
  • 清华大学交叉信息研究院 (Institute for Interdisciplinary Information Sciences, Tsinghua University)
  • 加州大学伯克利分校EECS (EECS, University of California, Berkeley)

会议/期刊: ICML 2025

原文保存位置: ~/.openclaw/workspace/papers/20260324_SpargeAttention/source/

报告生成日期: 2025-03-24


三、论文主体分析

1. 引言(Introduction)

随着大型模型中序列长度变得越来越长(例如视频生成和语言模型中的45K-128K),注意力机制的时间消耗占据了推理延迟的很大一部分。幸运的是,注意力图 $P = \mathrm{Softmax}(QK^\top / \sqrt{d}$ 呈现出固有的稀疏性,因为softmax操作通常会产生许多接近零的值。稀疏注意力通过以下方式利用这种稀疏性来加速注意力计算:(1) 构建一个"稀疏掩码",指示注意力图 $P$ 中应该计算的重要非零位置;(2) 仅对稀疏掩码对应的部分计算注意力。根据稀疏掩码的生成方式,稀疏注意力方法可分为三类:基于模式的方法(如H2O、InfLLM、DUOAttention等依赖滑动窗口或注意力sink)、动态稀疏注意力(如SparQAttn、LOKI、Minference、FlexPrefill等)和基于训练的方法(如Reformer、FastAttention)。

局限性包括:(L1) 通用性受限——现有稀疏注意力方法的通用性仍然有限,通常针对特定任务(如语言建模)开发,使用任务特定的模式(如滑动窗口或注意力sink),但注意力模式因任务而异;(L2) 实用性不足——难以实现既准确高效的稀疏注意力,因为准确性要求精确预测注意力图中的稀疏区域,而效率要求预测开销最小化。例如,Minference需要较长的序列长度(如100K)才能实现明显的加速。

目标:设计一种无训练的稀疏注意力算子,在不损失指标的情况下加速所有模型。

方法:本文开发了SpargeAttention,一种无训练的稀疏注意力,可通用应用于各种任务,包括语言建模和文本到图像/视频生成,以及各种序列长度。提出了三项主要技术来提高通用性、准确性和效率:

  1. 通用稀疏掩码预测算法:通过将每个Q和K块压缩为单个token来构建稀疏掩码,且选择性地基于块内token的相似性进行压缩
  2. 稀疏在线softmax算法(在GPU warp级别):通过利用在线softmax中全局最大值和局部最大值之间的差异,进一步省略一些 $PV$ 乘积
  3. 集成到8位量化SageAttention框架中以进一步加速

图 1:SpargeAttention在L40 GPU上对Mochi模型实现了1.83倍加速,且无视频质量损失。

图 1:SpargeAttention在L40 GPU上对Mochi模型实现了1.83倍加速,且无视频质量损失。

:::

图 2:视频、图像和语言生成模型中注意力图P的一些采样模式示例。

图 2:视频、图像和语言生成模型中注意力图P的一些采样模式示例。

:::


2. 相关工作(Related Work)

根据稀疏掩码的构建方式,稀疏注意力方法可分为三类:

(1) 基于模式的方法:依赖注意力图的某些固定模式,如滑动窗口或注意力sink。H2O、InfLLM和DUOAttention依赖滑动窗口模式;SampleAttention、MOA和StreamingLLM依赖滑动窗口和注意力sink模式;DitFastAttn依赖滑动窗口模式和不同注意力图之间的相似性。但DitFastAttn仅限于简单的扩散transformer,与语言模型和MMDiT模型(如Flux、Stable Diffusion3和3.5、CogVideoX)不兼容。由于模式因模型而异,这些方法可能无法普遍适用于不同模型。

(2) 动态稀疏方法:基于输入动态构建稀疏掩码,无需预设模式,因此可能更具通用性。现有的可进一步分为通道压缩和token压缩两类。通道压缩方法包括SparQAttn和LOKIAttn,通过降低维度的完整注意力来构建掩码,但由于常用注意力维度已经很小(如64、128),加速潜力可能有限。Token压缩方法包括Minference和FlexPrefill,通过将每个token块压缩为单个token并在此缩短序列上计算注意力来进行构建。然而,这种近似过于激进:如果它们在压缩序列上没有较大的注意力得分,可能会遗漏P的重要块。SeerAttention需要训练额外参数,成本较高。此外,它们都专为语言模型设计,对其他模型类型(如扩散模型)的适用性仍不确定。

(3) 基于训练的方法:修改注意力计算逻辑,需要重新训练整个模型,如Reformer和FastAttention。这些方法比无训练方法昂贵得多。

其他加速注意力的方法包括优化内核实现(如FlashAttention)、量化、分发工作负载、设计线性时间注意力等。这些与本文方法正交。

图 3:各种模型注意力中query和key的示例模式。

图 3:各种模型注意力中query和key的示例模式。

:::


3. 方法(SpargeAttention)

SpargeAttention包含两阶段在线过滤器以实现稀疏FlashAttention。如图4的Step1和Step2所示,首先设计了一种快速准确的方法来预测注意力图中的稀疏块,从而跳过相应的 $Q_iK_j^\top$ 和 $\widetilde{P}_{ij} V_j$ 乘积。如图4的Step3所示,进一步设计了稀疏在线softmax方法来跳过 $\widetilde{P}_{ij} V_j$ 乘积。

图 4:SpargeAttention的工作流程。

图 4:SpargeAttention的工作流程。

:::

3.1 稀疏FlashAttention

SpargeAttention采用FlashAttention的平铺策略,跳过被过滤掉的块的计算。设 $S = Q K^\top/\sqrt{d},~P = \sigma(S),~O = P V$,其中 $\sigma(S)_{ij} = \exp(S_{ij})/\sum_{k} \exp(S_{ik})$ 是softmax操作。令 $N$ 为序列长度,$d$ 为每个头的维度;矩阵 $Q$、$K$、$V$ 的维度均为 $N \times d$,矩阵 $S$ 和 $P$ 为 $N \times N$。

FlashAttention提出从token维度将 $Q$、$K$、$V$ 划分为块 $\{Q_i\}$、$\{K_i\}$、$\{V_i\}$,块大小分别为 $b_q$、$b_k$、$b_k$。然后使用在线softmax逐步计算 $O$ 的每个块 $O_i$:

$$ S_{ij} = Q_i K_j^\top / \sqrt{d},~~(m_{ij}, \widetilde P_{ij}) = \tilde\sigma(m_{i,j-1}, S_{ij}), $$ $$ l_{ij} = \exp(m_{i,j-1}-m_{ij}) l_{i,j-1} + \mathrm{rowsum}(\widetilde P_{ij}), $$ $$ O_{ij}=\diag{\exp(m_{i,j-1}-m_{ij})} O_{i,j-1} + \widetilde P_{ij} V_j $$

其中 $m_{ij}$ 和 $l_{ij}$ 是 $b_q \times 1$ 向量,初始值分别为 $-\infty$ 和 $0$。

定义1(块掩码):设 $M_{g}$ 和 $M_{pv}$ 是维度为 $\lceil N/b_q \rceil \times \lceil N/b_k \rceil$ 的二进制掩码,每个值为0或1。这些掩码决定稀疏注意力机制中哪些计算被跳过。

定义2(稀疏FlashAttention):稀疏FlashAttention的计算规则定义为:

  • 如果 $M_{g}[i,j] = 0$,则跳过 $Q_i K_j^\top, \widetilde{P}_{ij} V_j$
  • 如果 $M_{pv}[i,j] = 0$,则跳过 $\widetilde{P}_{ij} V_j$

3.2 选择性Token压缩用于稀疏预测

关键思想:虽然不同模型的注意力图各异,但我们观察到一个共同特征:query和key矩阵中大多数相邻token表现出高相似性(参见图3)。因此,对于由高度相似token组成的块,我们可以将这些token合并为块的单个代表性token。

具体而言,我们首先压缩 $Q$ 和 $K$ 中表现出高自相似性的块。然后,使用压缩后的 $Q$ 和 $K$ 快速计算压缩注意力图 $\hat P$ 。最后,选择性地计算 $\hat P$ 累积得分高的那些 $(i, j)$ 对的 $\{Q_iK_j^\top, \widetilde{P}_{ij} V_j$。值得注意的是,块选择仅在自相似性高的块(称为"选择性块")中进行。对于非自相似块(称为"固定块"),因为难以找到能代表整个块的好token,我们选择在注意力操作中始终计算这些非自相似块。

预测算法: 首先计算 $Q$ 和 $K$ 每个块的平均余弦相似度。然后,通过计算跨token的平均值将每个块压缩为单个token。接着,使用压缩的 $Q$ 和 $K$ 计算压缩的 $QK^\top$。最后,为防止来自非自相似块的干扰,我们将 $S$ 中对应的值设置为 $-\infty$,然后通过softmax获得压缩注意力图:

$$ q = \{q_i\} = \{\mathrm{mean}(Q_i, \mathrm{axis}=0)\} $$ $$ k = \{k_j\} = \{\mathrm{mean}(K_j, \mathrm{axis}=0)\} $$ $$ s_{qi} = \mathrm{CosSim}(Q_i),~ s_{kj} = \mathrm{CosSim}(K_j) $$ $$ \hat S[i] = q_ik^\top ; ~~~ \hat S[:,j] = - \infty, ~\mathrm{If}~ s_{kj} < \theta $$ $$ \hat P[i] = \mathrm{Softmax}(\hat S[i]) $$

其中 $\mathrm{CosSim}(X) = \mathrm{mean}(\frac{XX^\top}{|\max(XX^\top)|})$ 测量块内的余弦相似度。

对于 $\hat P$ 的每一行,选择累积和达到 $\tau \cdot \sum \hat P[i]$ 的顶部值位置,这些位置在 $M_{g}[i,:]$ 中设置为1,其他位置设置为0:

$$ M_{g}[i,:] = \mathrm{TopCdf}(\hat P[i], \tau) $$

最后,确保不遗漏涉及 $Q$ 或 $K$ 非自相似块(固定块)的计算: $$ M_{g}[i,:] = 1, ~\mathrm{If}~ s_{qi} < \theta ;~~~M_{g}[:,j] = 1, ~\mathrm{If}~ s_{kj} < \theta $$

3.3 第一阶段掩码

$M_{g}$ 可直接应用于FlashAttention以节省一些计算。在FlashAttention的内循环中(即计算 $Q_i$ 和 $\{K_j\}, \{V_j\}$ 之间的注意力时),当 $M_{g}[i,j] = 0$ 时可以跳过 $\{Q_iK_j^\top, \widetilde{P}_{ij} V_j\}$:

$$ \mathrm{Skip}~ Q_iK_j^\top ~\mathrm{and}~ \widetilde{P}_{ij} V_j,~~\mathrm{If}~ M_{g}[i,j] = 0 $$

3.4 稀疏Warp在线Softmax

关键思想:我们可以在在线softmax过程中进一步识别注意力图中足够小的值。如果 $\widetilde{P}_{ij}$ 中的所有值都接近零,则 $\widetilde{P}_{ij} V_j$ 可以忽略并省略。

为了识别哪些 $\widetilde{P}_{ij} =\exp(S_{ij}- m_{i,j})$ 包含可以省略的足够小的值,我们注意到在FlashAttention的每个内循环中,$O_{ij}$ 将乘以 $\exp(m_{i, j-1}-m_{ij})$ 然后加上 $\widetilde P_{ij} V_j$:

$$ m_\text{local} = \mathrm{rowmax}(S_{ij}), ~~m_{ij} = \max\{m_{i, j-1}, m_\text{local}\} $$ $$ O_{ij}= \mathrm{diag}{\exp(m_{i, j-1}-m_{ij})} O_{i, j-1} + \widetilde P_{ij} V_j $$

如果 $\mathrm{rowmax}(S_{ij}) < m_{ij}$,则 $m_{ij} = m_{i,j-1}$,$O_{ij} = O_{i,j-1} + \widetilde P_{ij} V_j$。此外,如果 $\mathrm{rowmax}(S_{ij}) \ll m_{ij}$,则 $\widetilde P_{ij} = \exp(S_{ij} - m_{ij})$ 中的所有值都接近0,导致 $\widetilde P_{ij} V_j$ 中的所有值都接近0。当以下条件满足时,$\widetilde P_{ij} V_j$ 可以忽略:

$$ \max(m_\text{local} - m_{ij}) < \lambda $$

因此,在FlashAttention的内循环中,$S_{ij}$ 将被 $c_w$ 个GPU warp划分为 $\{S_{ij}[\frac{i_w*b_q}{c_w}:\frac{(i_w+1)*b_q}{c_w}, :]\}$。如果 $\max(m_\text{local}[I_w] - m_{ij}[I_w]) < \lambda$,则 $O_{ij}[I_w] \approx O_{i,j-1}[I_w]$,我们将跳过用于更新 $O_{ij}[I_w]$ 的 $\widetilde P_{ij}[I_w] V_j$ 的计算。

3.5 与SageAttention结合

为了进一步加速稀疏注意力实现,我们将方法集成到SageAttention中,该方法提出了量化来加速注意力。由于量化和稀疏操作正交,稀疏计算可直接应用于SageAttention。

算法如Algorithm 1所示。具体来说,首先需要在SageAttention内循环开始时添加一个判断(Line 10),以决定是否一次性跳过整个内循环。其次,在SageAttention内循环中更新 $O_{ij}$ 之前添加另一个判断,以决定是否跳过 $\widetilde P_{ij} V_j$ 的计算。此外,为最小化注意力图预测开销,我们使用CUDA实现预测,并采用一些核融合技术。

3.6 超参数确定

基于方法描述,我们的方法包含三个超参数:$\tau \in (0,1)$、$\theta \in (-1,1)$ 和 $\lambda < 0$。参数确定过程很简单:目标是找到一组超参数,既能最大化注意力稀疏度,又能将五个不同模型输入的注意力误差限制在一定范围内。使用相对L1距离作为误差度量:$L1 = \sum |O - O'| / \sum |O|$。首先设置两个L1误差阈值 $l_1$ 和 $l_2$。首先对 $\tau$ 和 $\theta$ 进行网格搜索,找到在确保 $L1 < l_1$ 的同时最大化稀疏度的最优对。然后对 $\lambda$ 进行另一次网格搜索,找到在保持 $L1 < l_2$ 的同时进一步最大化稀疏度的最优值。

图 5:在 $1\times 6 \times 6$ 空间中不同token排列方法的说明,块大小为4。

图 5:在 $1\times 6 \times 6$ 空间中不同token排列方法的说明,块大小为4。

:::

3.7 Hilbert曲线排列

关键思想:在保持准确性的同时提高稀疏度是增强稀疏注意力性能的关键挑战。在算法中,增加key和query块的自相似性可以减少固定块的数量。这允许更多选择性块参与TopCdf选择,从而提高稀疏度。由于注意力计算对token排列不变,问题归结为找到一种增强相邻token相似性的排列

图像和视频模型受益于强先验:相邻像素可能相似。为了更好地利用这一先验,我们提出了Hilbert曲线排列。对于3D视觉token $Q, K, V \in \mathbb{R}^{T\times H \times W \times d}$,我们使用Hilbert曲线填充3D空间,然后沿曲线将token展平为形状 $\mathbb{R}^{L \times d}, L = T \times H \times W$。Hilbert曲线有效地保持局部性,遍历整个3D空间而不跨行或列,从而增加相邻token的相似性和注意力的稀疏度。


4. 实验(Experiment)

4.1 设置

模型:在语言、图像和视频生成的多种代表性模型上验证SpargeAttention的有效性。具体来说,在Llama3.1(8B)上进行文本到文本的实验,在Cogvideo(2B)、Mochi和Open-Sora-Plan上进行文本到视频的实验,在Flux(.1-dev)和Stable-Diffusion3.5(large)上进行文本到图像的实验。

数据集:文本到文本模型在四个零样本任务上评估:WikiText(评估模型预测置信度)、Longbench和InfiniteBench的En.MC(全面评估长上下文理解能力)、Needle-in-a-Haystack任务(评估模型检索能力)。文本到视频模型使用Open-Sora提示集进行评估。文本到图像模型在COCO注释上进行评估。

端到端指标:对于Llama3.1,使用WikiText的困惑度(ppl.)、Longbench分数和Needle-in-a-Haystack任务的检索准确率。对于文本到视频模型,按照VIDIT-Q的评估方法,使用五个指标评估生成视频的质量:CLIPSIM和CLIP-Temp(CLIP-T)用于测量文本-视频对齐;VQA-a和VQA-t用于评估视频美学和技术质量;Flow-score(FScore)用于时间一致性。对于文本到图像模型,将生成的图像与COCO数据集中的图像在三个方面进行比较:FID用于保真度评估,Clipscore(CLIP)用于文本-图像对齐,ImageReward(IR)用于人类偏好。

速度和稀疏度指标:使用逆延迟 $1/t$ 来评估稀疏注意力方法的速度。速度指标是完全公平的,因为 $O(attn)$ 对于一组输入是固定的,速度由 $t$ 决定,这包括预测注意力图稀疏区域的时间。稀疏度定义为相对于完整注意力所需的 $Q_iK_j^\top$ 加上 $\widetilde{P}_{ij} V_j$ 的总数,被跳过的矩阵乘法比例。

实现和超参数:使用CUDA实现。需要为模型确定 $l_1, l_2$。对于Llama3.1使用 ($l_1=0.08, l_2=0.09$),对于Cogvideo和Mochi使用 ($l_1=0.05, l_2=0.06$),对于SD和Flux使用 ($l_1=0.07, l_2=0.08$),对于Open-Sora-Plan使用 ($l_1=0.03, l_2=0.035$)。

基线:目前适用于不同模型类型的稀疏注意力方法有限。我们选择块稀疏Minference和FlexPrefill作为基线。

4.2 质量和效率评估

端到端指标:表1展示了各种模型使用SpargeAttention与使用完整注意力和基线的端到端指标比较。可以观察到,我们的方法在各种模型上几乎没有端到端指标损失,并且在端到端准确性方面优于具有各种稀疏度的基线。图6-9展示了一些在CogvideoX、Flux、SD、Mochi和Open-Sora-Plan上的可见比较示例,表明SpargeAttention没有性能损失且优于基线。

注意力速度:表1显示我们的方法比完整注意力实现更快的速度,并且在注意力速度方面优于具有各种稀疏度的基线。图10展示了各种方法在不同稀疏度下的内核速度,突出了我们方法的高效性和相对于其他方法的显著优势。

图 6:CogvideoX上使用SpargeAttention的可见示例。

图 6:CogvideoX上使用SpargeAttention的可见示例。

:::

图 7:Flux和SD上的比较示例。

图 7:Flux和SD上的比较示例。

:::

图 8:Mochi上的比较示例。

图 8:Mochi上的比较示例。

:::

图 9:Open-Sora-Plan上的可见示例。

图 9:Open-Sora-Plan上的可见示例。

:::

图 10:不同稀疏度下的内核速度比较。输入张量的序列长度为22K,头维度为128。

图 10:不同稀疏度下的内核速度比较。输入张量的序列长度为22K,头维度为128。

:::

稀疏块预测的开销:表2比较了SpargeAttention中动态稀疏块预测的开销与注意力执行延迟。结果表明,预测开销很小,特别是对于较长的序列。

端到端加速:表3展示了在Cogvideo、Mochi和Llama3.1上使用SpargeAttention的端到端延迟。值得注意的是,SpargeAttention在Mochi上实现了1.83倍加速。

4.3 消融研究和关键洞察

Hilbert曲线排列的效果:我们通过比较三个指标来评估Hilbert曲线排列对Mochi的影响:query或key块的平均块相似度、定义的L1误差和稀疏度。表4表明HilbertCurve排列始终获得更高的块自相似度和稀疏度,而精度仅有微小差异。

自相似性判断的消融:表5展示了自相似性判断对Mochi的影响,发现自相似性判断可以保证端到端准确性。

$M_g$ 和 $M_{pv}$ 的稀疏度分析:表6显示了在仅使用 $M_g$、仅使用 $M_{pv}$ 和使用 $M_g$+$M_{pv}$ 时的稀疏度。

SpargeAttention增强LLM性能:从结果观察到SpargeAttention在长上下文任务中增强了LLM性能。这可能是因为稀疏注意力帮助LLM聚焦于更相关信息。

稀疏度随序列长度增加:表7显示了在Llama3.1上,稀疏度随序列长度增加。这表明序列越长,SpargeAttention可以实现的加速越高。

扩散模型稀疏度分析:我们对Cogvideo的所有层、头、时间步和样本使用SpargeAttention进行稀疏度详细分析。发现稀疏度随层和头变化,表明需要为每个层和头设置不同的超参数。还发现对于扩散模型,稀疏度随采样时间步增加。

图 11:Llama3.1上Needle-in-a-Haystack比较示例。

图 11:Llama3.1上Needle-in-a-Haystack比较示例。

:::

图 12:Llama3.1上Needle-in-a-Haystack比较示例。

图 12:Llama3.1上Needle-in-a-Haystack比较示例。

:::


5. 结论(Conclusion)

本文提出了SpargeAttention,一种通用的稀疏和量化注意力,对于任何输入都能高效准确地执行注意力。我们的方法使用两阶段在线过滤器:第一阶段,我们快速准确地预测注意力图,从而跳过注意力中的一些矩阵乘法。第二阶段,我们设计了一种不增加额外开销的在线softmax感知过滤器,进一步跳过一些矩阵乘法。实验表明,SpargeAttention在语言、图像和视频生成模型上实现了加速,且不牺牲端到端指标。


附录(Appendix)

A.1 排列消融的详细解释和结果

我们使用五个不同的提示词,并在Cogvideo和Mochi模型上使用预搜索的超参数 $l_1=0.05, l_2=0.06$ 进行实验。排列在位置嵌入后的注意力操作中分别对 $Q、K、V$ 执行。为保留输入序列的原始顺序,在注意力输出上执行逆排列;对于使用视觉-语言联合自注意力的模型(如Cogvideo),我们仅排列视觉token。

评估块自相似度时,我们为query选择块大小128,为key选择块大小64,与内核实现一致。精度指标(L1)使用FlashAttention2输出作为真值进行评估。

A.2 自相似性判断的消融研究

为研究自相似性判断对注意力性能的影响,我们进行消融研究,移除自相似性判断。在大多数情况下,高度局部化模式的存在导致非自相似块数量很少,因此在所有张量情况平均时,精度和稀疏度仅有微小差异。为获得更有意义和可解释的见解,我们专门分析精度差异具有统计意义的情况。

A.3 其他实验结果

图13-14和表8展示了在16-28K序列长度的Llama3.1上Needle-in-a-Haystack任务中的额外结果和比较示例。

A.4 扩散模型稀疏度分析

本节分析了Cogvideo在不同维度上的稀疏度模式:模型层、去噪时间步、输入样本和注意力头。

图 14:Cogvideo的逐层稀疏度。

图 14:Cogvideo的逐层稀疏度。

:::

图 15:Cogvideo的逐时间步稀疏度。

图 15:Cogvideo的逐时间步稀疏度。

:::

图 16:Cogvideo的逐样本稀疏度。

图 16:Cogvideo的逐样本稀疏度。

:::

图 17:Cogvideo的逐头稀疏度。

图 17:Cogvideo的逐头稀疏度。

:::


四、论文简评

创新点

  1. 通用性强:SpargeAttention是首个能够真正加速语言、图像和视频模型且不影响准确性的稀疏注意力方法,解决了现有方法通用性不足的问题。

  2. 无训练:无需对模型进行任何训练或微调,直接作为推理插件使用,极大提高了实用性。

  3. 两阶段过滤机制:提出创新的两阶段在线过滤器——第一阶段通过选择性token压缩预测稀疏块,第二阶段利用在线softmax特性进一步跳过计算,开销极小。

  4. 自相似性判断:通过引入自相似性判断,确保不会遗漏重要信息,保证端到端性能不损失。

  5. Hilbert曲线排列:针对视觉模型提出Hilbert曲线排列,有效增加块的自相似性,进一步提高稀疏度。

  6. 与量化结合:与SageAttention量化框架无缝集成,实现进一步的性能提升。

局限性

  1. 超参数调优:需要为不同模型调整 $l_1$、$l_2$ 等超参数,虽然过程不复杂但仍需一定的人工介入。

  2. 稀疏度与精度的权衡:虽然文中方法在保持精度方面表现优秀,但当稀疏度过高时仍可能影响性能。

  3. 预测开销:虽然预测开销相对较小,但对于极短序列(如8K以下),仍有一定的开销比例。

应用场景

  1. 大模型推理加速:适用于各类大模型的推理阶段,包括LLM、图像生成模型(如Stable Diffusion、Flux)和视频生成模型(如CogVideoX、Mochi、Open-Sora)。

  2. 长上下文处理:特别适合处理长序列任务,稀疏度随序列长度增加而提高。

  3. 资源受限部署:在GPU资源有限的情况下,通过稀疏计算实现更高效的模型部署。

  4. 生产环境部署:由于是无训练方法,可直接作为推理库集成到现有系统中。

可改进方向

  1. 自适应超参数:研究自动为不同层和头选择最优超参数的方法,减少人工调参。

  2. 更精确的预测:探索更准确的稀疏块预测算法,进一步提高稀疏度。

  3. 更多模型支持:将方法扩展到更多类型的模型,如多模态模型、语音模型等。

  4. 硬件协同优化:针对不同GPU架构进行更深入的硬件级优化。