SageAttention2 Technical Report:
Accurate 4 Bit Attention for Plug-and-play Inference Acceleration

自己注意の計算は以下のように定式化できる：$S=QK^{\top}/\sqrt{d},~{}P=\sigma(S),~{}O=PV$、ここで$\sigma(S)_{ij}=\exp(S_{ij})/\sum_{k}\exp(S_{ik})$はソフトマックス演算である。 行列$Q$、$K$、および$V$はそれぞれ$N\times d$の次元を持ち、一方で行列$S$、$P$は$N\times N$である。$d$は通常小さく、例えば64や128であるが、$N$は数千あるいは数百万にもなり得る。したがって、$N\times N$行列$(S,P)$は$(Q,K,V)$よりもはるかに大きく、素朴な実装では$(S,P)$の読み書きに膨大なグローバルメモリIOが必要となる。FlashAttention (Dao, 2023)は、$Q$、$K$、および$V$をトークン次元からブロックサイズ$b_{q}$、$b_{kv}$、$b_{kv}$のブロック$\{Q_{i}\},\{K_{i}\},\{V_{i}\}$にタイル化することを提案している。そして、$(S,P)$のメモリIOを避けるために、オンラインソフトマックス(Milakov & Gimelshein, 2018)を使用して$O$の各ブロック、つまり$O_{i}$を段階的に計算する：

まず、$\{K_{i}\},\{V_{i}\}$の各ブロックに対して、以下の方程式を反復的に計算する：

ここで、$m_{i}^{j}$と$l_{i}^{j}$は$b_{q}\times 1$ベクトルであり、それぞれ$-\infty$と$0$で初期化される。$\tilde{\sigma}()$はオンラインソフトマックス演算子である：$m^{j}_{i}=\max\{m^{j-1}_{i},\mathrm{rowmax}(S^{j}_{i})\},~{}\widetilde{P}_{j}^{i}=\exp(S^{j}_{i}-m^{j}_{i})$。

最後に、出力$O_{i}$は$O_{i}=\mathrm{diag}(l_{i}^{j})^{-1}O_{i}^{j}$によって計算できる。

行列乗算 $C=AB$ は、以下のように量子化によって加速できる：

$\psi$ は量子化器であり、高精度（例：FP32またはFP16）の行列 $A$ を低精度形式 $\hat{A}$（例：INT4またはFP8）にスケール $\delta_{A}$ で変換し、$\psi^{-1}$ は高精度に戻す逆量子化器である。我々は $\psi^{-1}_{\delta_{A}}(\hat{A})\approx A$ を持つべきである。実際の行列乗算 $\hat{A}\hat{B}$ は低精度で行われる。現代のGPUでは、低精度の行列乗算は通常、高精度のものよりも数倍高速である。 多くの量子化器は、数値形式と粒度に依存する。例えば、共通のスケール因子を共有する要素の数などである。 例えば、INT4のテンソル単位の量子化器は、まずテンソル全体の絶対値の最大値としてスケールを計算し、要素をINT4の最大表現可能範囲[-7, +7]にスケーリングし、その後四捨五入してINT4にキャストする：$\hat{A}=\lceil A/\delta_{A}\rfloor,\delta_{A}=\max(|{A}|)/7$。 同様に、トークン単位の量子化器はテンソルの各トークンにスケール因子を割り当てる：$\hat{A}[i,:]=\lceil A[i,:]/\delta_{A}[i]\rfloor,\delta_{A}[i]=\max(|{A[i,:]}|)/7$。 また、チャンネル単位の量子化器はテンソルの各チャンネルにスケール因子を割り当てる。つまり、チャンネル次元に沿って：$A[:,i]=\lceil A[:,i]/\delta_{A}[i]\rfloor,\delta_{A}[i]=\max(|A[:,i]|)/{7}$。 逆量子化プロセスは要素ごとのスケーリングである：$\psi_{\delta_{A}\delta_{B}}^{-1}(\hat{A}\hat{B})=\hat{A}\hat{B}*\delta_{A}*\delta_{B}$。

FlashAttention2 (Dao, 2023) のブロックタイリングに基づき、SageAttention (Zhang et al., 2024) はブロック単位の粒度で $Q,K$ をINT8に量子化する。さらに、量子化の精度を保つために、SageAttentionは最初に $K$ を平滑化することを提案している：

ここで、$Q_{i},K_{j}$ はFlashAttentionでタイル化されたブロックであり、$\text{mean}(K)$ は $K$ のチャンネル次元の平均値である。SageAttentionは $\widetilde{P}V$ をFP16として保持し、$\widetilde{P}V$ にはFP16アキュムレータを使用したFP16行列乗算を使用する。しかし、FP16アキュムレータを使用したFP16行列乗算は、RTX4090およびRTX3090 GPUでのみ高速化効果がある。

第2節で紹介したFlashAttentionと量子化に基づき、我々が開発した量子化注意機構アプローチについて説明する。

$\phi_{Q}$、$\phi_{K}$、$\phi_{V}$は、量子化された$Q$、$K$、$V$を得るための3つの変換であり、これらについては後続のセクションで議論する。 簡略化のため、すべての上付き文字と下付き文字は省略するが、注意機構で使用される行列は依然としてタイルであり、計算は第2.1節で説明したFlashAttentionとして組織化されている。式2、3に示されている元の全精度版と比較して、SageAttention2は$Q,K,P,V$に量子化器を追加し、積に逆量子化器を追加することで、$QK$と$\widetilde{P}V$の両方の行列乗算を加速している。

SageAttentionはブロック単位の量子化を使用しており、GPUのストリーミングプロセッサごとに$Q$と$K$の各ブロックを量子化する。このような量子化戦略は、トークン単位の量子化に近い精度性能を達成し、量子化スケールベクトル$\delta_{Q}$と$\delta_{K}$のドット積のオーバーヘッドを回避できる。しかし、$Q$と$K$をINT4に量子化するには、より正確な量子化粒度が必要である。我々はワープ単位の量子化を提案する。これはブロック単位の量子化器よりも精密で粒度の細かい量子化アプローチであり、ベクトルのドット積による追加のオーバーヘッドもない。

具体的には、SageAttentionにおける$Q$の各ブロックは$c_{w}$個のセグメントに分割され、それぞれがGPUストリーミングプロセッサ（SM）内の$c_{w}$個のGPUワープの1つによって処理される。その後、各ワープが割り当てられたセグメントに対してMatmulを実行する。ワープ単位のINT4量子化は、$Q$の$b_{q}/{wc}$トークンごとにスケール因子を割り当てる：

ワープ単位の量子化は、SageAttentionで使用されているブロック単位の量子化よりも優れた精度性能を提供する。これについては3.5節で議論する。

INT4量子化の代表的な範囲は著しく制限されており、すなわち$2^{4}-1=15$値のみです。この制限は、注意機構において$Q,K$をINT4に量子化する際に性能を大幅に低下させます。例えば、$Q,K$をINT4で量子化すると、WikiTextにおけるLlama3.1のパープレキシティが90%以上増加し、CogvideoXが生成する動画の品質が約3倍低下します（表1参照）。我々は実際のモデルにおける$Q,K$のデータ分布を分析しました。例えば、Llama3.1とCogvideoXでは$Q,K$においてチャンネル単位の顕著な外れ値が見られることがわかりました（図4参照）。チャンネルごとの量子化はこのような外れ値による量子化誤差を軽減できますが、この方法は$Q,K$には適用できません。なぜなら、量子化は$QK^{\top}$の外れ値がもたらす量子化誤差を排除する方法を提案します：

ここで、$\mathrm{mean}(K)$と$\gamma(K)$はSageAttention (Zhang et al., 2024)で議論されており、$K$の変換は注意スコア$\widetilde{P}$を変更しません。$\mathrm{mean}(Q)=\frac{1}{N}\sum_{t=1}^{N}Q[t,:]$は形状$1\times d$のベクトルです。$Q$の変換については、$\overrightarrow{Q_{m}}$と$K^{\top}$の間でGEMV（一般行列ベクトル乗算）を実行します。つまり、$\overrightarrow{Q_{m}}K^{\top}$を$\Delta S$として計算します。この$\Delta S$は注意の計算中に$S$に加算されます。 最終的に、全精度の$Q,K$から量子化された$\hat{Q},\hat{K}$への変換は式6のように表現できます。ここで、$\psi_{Q}$、$\psi_{K}$は$Q$と$K$のための2つの量子化器です。言い換えれば、全精度の$Q,K$はチャンネル次元の平均値を減算してから量子化されます。

図5は、CogvideoXから平滑化Qの有無による量子化された$Q$の分布の例を示しています。平滑化Qを使用することで、INT4の範囲がより均一かつ完全に利用されていることがわかります。 表1は、Llama3.1 (Dubey et al., 2024)とCogvideoX (Yang et al., 2024)における平滑化Q+Kの有無による異なる量子化方法のエンドツーエンドの指標を示しています。結果は平滑化Q+Kが精度に大きな利点をもたらすことを示しています。また、表10と表10は、効果の順序が平滑化Q+K $>$ 平滑化Q $>$ 平滑化K $>$ その他のベースラインであることを示しています。

SageAttentionの欠点である、$PV$に対するFP16累算器を持つFP16 MatmulがRTX4090のようなGPUでのみ効果的であることを避けるため、我々は$P$と$V$をFP8に量子化してFP8テンソルコアの普遍的な加速を活用するアプローチを採用した。しかし、我々はAda アーキテクチャ上のmma(f32f8f8f32)命令の累算器が実際にはFP22、具体的には1符号ビット、8指数ビット、13仮数ビットであることを発見した。具体的には、mma(f32f8f8f32)命令$C=AB+D$において、$A,B$はFP8データ型のテンソルで、$C,D$はFP32データ型のテンソルである。我々は$A,B$をゼロに初期化し、$D$を変化させて累算器のデータ型をテストした。表7に示すように、$D$が1符号ビット、8指数ビット、13仮数ビットで初期化されると、$C$の値はmma(f16f16f32)命令の結果と一致する。しかし、$D$が13ビット以上の仮数ビットで初期化されると、$C$の誤差はmma(f32f16f16f32)とmma(f32f8f8f32)の結果の差に対応する。 したがって、FP8に量子化された行列$\widetilde{P}$と$V$の行列乗算は、FP32累算器を使用する場合と比較して、ある程度の精度損失を被る。この精度損失を可能な限り軽減するために、我々は$V$を滑らかにすることを提案する：

図6に示すように、この戦略は以下の理由により$\widetilde{P}V$の値に対するFP22の精度を向上させる：$\widetilde{P}$の各行は0から1の値範囲にわたり、$V$の各列は一貫してチャンネル方向の外れ値を持ち、それらは排他的に正または負で、通常8から9の範囲にある。結果として、$\widetilde{P}V$の値はかなり大きくなる可能性がある。しかし、浮動小数点数の表現範囲は一様ではなく、ゼロ付近でより密である。したがって、チャンネル次元に沿って平均を$V$から減算することで、$\widetilde{P}V$の値はゼロに近くなり、より高い表現精度が得られる。我々は、このような説明と戦略がコミュニティによって報告された多くの問題¹¹1https://github.com/triton-lang/triton/issues/4476, https://github.com/triton-lang/triton/issues/5065を解決できると考えている。 さらに、アテンション計算の正確性を維持するためには、$\overrightarrow{V_{m}}$を$O$の最終計算に加えるだけで十分である：$O=O+\overrightarrow{V_{m}}$。これは、$\widetilde{P}$行列の各行の合計が1に等しいため、$\widetilde{P}\overrightarrow{V_{m}}=\overrightarrow{V_{m}}$となるからである。 言い換えれば、この方法は$V$を2つの部分に分解する：$\overrightarrow{V_{m}}$と$V$。$V$については、各列の値をゼロの周りに中心化し、これにより量子化された$\widetilde{P}$行列の行と量子化された$V$行列の列との内積結果がゼロに近くなる。これによりFP22の表現がより正確になる。一方、$\overrightarrow{V_{m}}$はFP16で保持され、最後に$O$に加算されるため、$\overrightarrow{V_{m}}$部分の精度損失を引き起こさない。

表6は、CogvideoXからサンプリングされた実際のテンソルに対する、$V$の平滑化の有無によるアテンション精度を示している。これは、$Q,K$をINT4に、$\widetilde{P},V$をFP8に量子化する際に、Vを平滑化することでSageAttention2の精度を向上させることができることを示している。

セクション3.5の議論に基づき、我々は$\psi_{Q}(Q)$と$\psi_{K}(K)$にINT4またはINT8量子化を使用する選択に基づいて2つのアテンションカーネル（表8参照）を実装する。 これらのカーネルの速度は（SageAttn2-4b$>$SageAttn2-8b）の順であるが、精度の順序は逆である。 SageAttn2-4bは多くのタスクで十分に機能する可能性があり、例えばLlama3.1のLambdaベンチマークでのエンドツーエンドの指標損失はごくわずかであるが、SageAttn2-4bは一部の困難なシナリオでは不十分である。したがって、我々はまずSageAttn2-4bとSageAttn2-8bを使用して様々なモデルの精度性能の詳細な分析を行う。その後、効果的で使いやすい適応的手法を提案する。