数日前から、個人的な学習として、Multi-Head Attention に関連する記事を書いています。今回は Multi-Head Attention における KV Cache の容量と Multi-Query Attention について調べてみます。
KV Cache の容量
一般的な Multi-Head Attention の KV Cache の容量は次式で表されます。
Size = 2 \times B \times L \times T \times H \times D \times S
- 2: Key, Value の2種類
- B: Batch Size
- L: Transformer の層数
- T: 保存するトークン数（コンテキスト長）
- H: Head 数
- D: Head あたりの次元数 (head_dim)
- S: 1要素のサイズ (FP16の場合2Byte)
なお、一般的な Transformer では、
H \times D = HiddenSize
となります。HiddenSize が一定であれば、Head 数を増やしても KV Cache の総容量は変わりません。
Llama 3 8B の場合
Llama 3 8B の場合の例として、以下のパラメータを想定します。
- Batch = 1
- Layer = 32
- Context = 4096
- hidden_size = 4096
- FP16
この時、1層あたりの Key, Value の容量は以下のようになります。
\begin{align}
Size(Key+Value) &= 2 \times 4096 \times 4096 \times 2 \\
&= 67,108,864 Byte \\
&≈ 64 MiB
\end{align}
32層の場合は、2GiB ほどになります。
Head 数を変えた場合でも、H × D = 4096 は固定のため、KV Cache の容量は変わりません。
-
Head = 32->head_dim = 128 -
Head = 64->head_dim = 64
Multi-Query Attention
Multi-Query Attention (MQA)1 は、複数の Query Head が単一の共有 Key Head と Value Head を使い回すことで、推論時の KV Cache 容量を削減するとともに、GPUメモリ帯域の使用量も削減する Attention 機構です。
Query Head が32個、KV Head が1個の場合、1 x head_dim となります。これにより KV Cache の容量は $ \frac{1}{32} $ になります。
実際の LLM では、MQA ではなく、その発展形である Grouped-Query Attention (GQA) が広く採用されています。次回は、GQA の KV Cache 削減効果について見て見たいと思います。