跳转至

Efficient Attention

现阶段做高效Attention的机制主要有两个，一个是低比特矩阵计算的Attention，另一种是稀疏的注意力机制

其实还存在另一种高效方案就是Linear Attention，包括是Kimi最近在做的一系列工作，但是这种方案还未能做到完全的工业场景的落地

在Context Length的增长中，Attention计算会变成一个巨大的瓶颈，在长文本，视频之类的领域尤其明显

1. 低比特计算

这一系列的代表工作主要是SageAttention，在FP32的Full Attention的基线中引入了各项低比特机制

  • SageAttention1 INT8 + FP16 inference
  • SageAttention2 INT4/INT8 + FP8 inference
  • SageAttention3 FP4 inference FP8 Training

SageAttention，已经集成在Diffusers，TensorRT等库中，很多视频生成产品都用了

在应用的时候直接替换FlashAttention，一行就可以实现加速

SageAttention也是基于FlashAttention之上的

在低比特量化的时候，由于两个单位的上下界不同，因此一般在低比特量化中需要做一些缩放，量化完成的矩阵还会带一个小标量，就是一个缩放因子，可以看下图中的 \(s_q\)

也是建立在FlashAttention的基础上，对于分块好的矩阵先进行量化，在 \(QK^T\) 这个计算中使用INT8 在最后和 V 矩阵的相乘中使用FP16

但是量化本身发生在token维度上，一个序列的token的attention值其实差异波动是很大的，在这种情况下，量化的效果会比较糟糕，因为他是基于同一个缩放因子去进行缩放的

在这里引入了一个新的方案叫做Smooth K，其用于平均化序列中的异常值，其实就是减去每列的均值，并且这个操作不影响softmax

在SageAttention1中其实没有动最后和V矩阵相乘这一块，但是在SageAttention2中也量化这一部分，随之也引入了SmoothQ

2. 稀疏注意力

这部分的主要工作是Sparge Attention等一系列的工作，并且他们都是建立在Saged Attention之上的