GQA/MQAとは？KV Cacheを減らすAttention改良

MHAからGQA、MQAへ進むほどKV head共有が強くなりKV Cacheが小さくなるアイキャッチ

GQA（Grouped-Query Attention、Query headをgroup化してK/Vを共有するAttention）とMQA（Multi-Query Attention、すべてのQuery headで1組のK/Vを共有するAttention）は、Query headの多様性を残しながらKV head数を減らす設計です。KV Cacheの容量とdecode時に読み出すK/Vのbyte数を抑えることが主な狙いです。

「Attention headを減らす」とだけ覚えると正確ではありません。複数のQuery headは維持し、減らすのはKey/Value headです。またKV Cache部分が4分の1になっても、model weightやactivationを含む総VRAMが4分の1になるわけではありません。

MHAからGQA、MQAへ進むほどKV head共有が強くなりKV Cacheが小さくなるアイキャッチ

3文要約

  1. MHAはQueryと同数のK/V headを持ち、MQAはK/Vを1組へ共有し、GQAは複数のQuery head groupごとに1組のK/Vを共有します。
  2. 他の条件が同じなら、KV Cache payloadとdecode時のK/V読出し量はKV head数に比例するため、MHAに対する比率は\(H_{kv}/H_q\)で概算できます。
  3. K/V共有を強くすると容量と帯域を減らせる一方、品質、kernel対応、checkpoint変換がtrade-offになるため、model configと実測結果を確認します。

代表的な論文

MQAの代表資料は、Noam ShazeerのFast Transformer Decoding: One Write-Head is All You Needです。incremental decode（生成時に1 tokenずつ順番に計算する処理）では、過去のK/V tensorを各stepで読み出すmemory bandwidthがbottleneckになり得ると分析し、複数のQuery headでK/Vを共有するMQAを提案しました。

GQAの代表資料は、Joshua AinslieらのGQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointsです。K/Vを1組まで集約するMQAと、headごとにK/Vを持つMHAの間として、複数groupのK/V headを使うGQAを導入しました。同論文は既存のMHA checkpointからK/V projectionをmean poolingし、追加pretrainingで適応させるuptrainingも扱っています。

MHAを基準にhead構成を整理する

MHA（Multi-Head Attention、複数headで異なる表現へ注目するAttention）では、各headがそれぞれQuery、Key、Valueを持ちます。Attention Is All You Needで導入された基本形です。

batch sizeを\(B\)、系列長を\(N\)、Query head数を\(H_q\)、KV head数を\(H_{kv}\)、1 headの次元を\(d_h\)とします。tensor shapeは次のように書けます。

\[Q\in\mathbb{R}^{B\times H_q\times N\times d_h}\]
\[K,V\in\mathbb{R}^{B\times H_{kv}\times N\times d_h}\]

MHAでは\(H_q=H_{kv}\)です。たとえばQuery headが8個なら、Key headも8個、Value headも8個あります。各head \(i\)は自分に対応する\(Q_i\)、\(K_i\)、\(V_i\)でAttentionを計算します。

\[O_i=\operatorname{softmax}\left(\frac{Q_iK_i^{\mathsf{T}}}{\sqrt{d_h}}\right)V_i\]

各headの出力を連結してoutput projectionへ渡すため、複数の視点を保てます。一方、decode中のKV CacheにもheadごとのK/Vを保存するため、KV head数が多いほどcacheは大きくなります。Self-Attentionの式を先に確認したい場合は、Self-AttentionとScaled Dot-Product Attentionの解説を参照してください。

MQAはK/Vを1組へ共有する

MQAではQuery headを複数残したまま、Key headとValue headをそれぞれ1個にします。

\[H_{kv}=1\]

各Query headは独自の\(Q_i\)を持ちますが、全headが共通の\(K_0\)と\(V_0\)を参照します。

\[O_i=\operatorname{softmax}\left(\frac{Q_iK_0^{\mathsf{T}}}{\sqrt{d_h}}\right)V_0\]

この変更により、過去tokenごとに保存するK/Vは1組で済みます。Query headが32個なら、MHAの32組に対してMQAは1組です。他条件が同じK/V payloadだけを比べれば、容量と読出し量はMHAの1/32になります。

ただし、Query headが1個になるわけではありません。Query projection、headごとのscore、softmax weight、出力は残ります。MQAは「single-head Attentionへ戻す方式」ではなく、「Multi-Query、single-KV」の設計です。

GQAはMHAとMQAの中間を選べる

MQAはK/V共有を最も強くするため、KV Cacheと帯域を大きく減らせます。一方、K/V表現の多様性も1組へ集約します。GQAはQuery headを複数groupへ分け、groupごとに1組のK/V headを持たせます。

Query head数を\(H_q\)、KV head数を\(H_{kv}\)とすると、一般的なGQAは次の範囲です。

\[1 < H_{kv} < H_q\]

1個のKV headを共有するQuery head数、つまりgroup size \(s\)は次の式です。

\[s=\frac{H_q}{H_{kv}}\]
方式 Query head数 KV head数 共有範囲 KV Cache payload
MHA \(H_q\) \(H_q\) headごとに専用 基準
GQA \(H_q\) \(1より多くH_q未満\) Query groupごとに共有 中間
MQA \(H_q\) 1 全Query headで共有 最小

GQAのgroup数が1ならMQA、group数がQuery head数と等しければMHAです。つまりMQAとMHAはGQAの両端として表せます。

8個のQuery headに対しMHAは8組、GQAは2組、MQAは1組のK/V headを使う比較

「group数」と「group size」を分ける

用語を逆に読むと設定を間違えます。

  • group数: KV head数\(H_{kv}\)と同じ
  • group size: 1個のKV headを共有するQuery head数\(H_q/H_{kv}\)

たとえば\(H_q=32\)、\(H_{kv}=8\)なら、group数は8、group sizeは4です。Query head indexを0から数える場合、対応は次のようになります。

KV head 共有するQuery head
0 0、1、2、3
1 4、5、6、7
2 8、9、10、11
… …
7 28、29、30、31

連続したgroupへ均等に割り当てるなら、Query head \(i\)が使うKV head index \(g(i)\)は次の式です。

\[g(i)=\left\lfloor\frac{i}{s}\right\rfloor,\qquad s=\frac{H_q}{H_{kv}}\]

この構成では\(H_q\)が\(H_{kv}\)で割り切れる必要があります。現行PyTorchのGQA機能も、この割り切りを制約にしています。

32個のQuery headを4個ずつ8 groupへ分け各groupを1組のK/V headへ対応させる図

同じK/Vを使っても出力は同じにならない

GQAやMQAでは複数のQuery headが同じK/Vを参照します。それでも、各headの出力が同じになるわけではありません。

同じgroupに属するQuery head \(i\)と\(j\)を考えます。両者は同じ\(K_g\)、\(V_g\)を使いますが、\(Q_i\)と\(Q_j\)は別のprojectionから得られます。

\[A_i=\operatorname{softmax}\left(\frac{Q_iK_g^{\mathsf{T}}}{\sqrt{d_h}}\right)\]
\[A_j=\operatorname{softmax}\left(\frac{Q_jK_g^{\mathsf{T}}}{\sqrt{d_h}}\right)\]

\(Q_i\neq Q_j\)なら、通常はscoreとattention weight \(A_i,A_j\)も異なります。その結果、同じ\(V_g\)に対するweighted sumも異なります。

\[O_i=A_iV_g,\qquad O_j=A_jV_g\]

共有するのはK/Vという参照対象です。「どの過去tokenへ、どの強さで注目するか」はQueryごとに変わります。

2つのQuery headが同じK/Vを参照しても異なるscoreとattention weightから別の出力を得る図

KV Cache容量はKV head数に比例する

DecoderのKV Cacheが保存する生のK/V payloadを概算します。layer数を\(L\)、batch sizeを\(B\)、保存token数を\(N\)、KV head数を\(H_{kv}\)、head dimensionを\(d_h\)、1要素のbyte数を\(b\)とすると、容量は次の式です。

\[M_{KV}=2\times L\times B\times N\times H_{kv}\times d_h\times b\]

先頭の2はKeyとValueの2 tensorを表します。

Queryは次のtokenを計算するたびに新しく作るため、通常のKV Cacheへ過去Queryを保存しません。GQA/MQAが減らすのは式中の\(H_{kv}\)です。

MHAに対するKV Cache比は、他条件が同じなら次の値になります。

\[\frac{M_{GQA/MQA}}{M_{MHA}}=\frac{H_{kv}}{H_q}\]

decodeの各stepでは、現在のQueryから過去すべてのK/Vを参照します。そのためK/V読出しbyte数も、同じ実装条件のHkvに比例する部分があります。これはKV Cacheを保存して再計算を避ける仕組みとは別の層です。基礎はKV Cacheの記事で説明しています。

layer、batch、token、KV head、head dimensionからKV Cache容量を計算するtensor図

32層・8192 tokenで容量を計算する

次の条件で、1 requestのK/V payloadを比較します。

  • layer数\(L=32\)
  • batch size\(B=1\)
  • 保存token数\(N=8,192\)
  • Query head数\(H_q=32\)
  • head dimension\(d_h=128\)
  • dtypeはBF16、1要素\(b=2\) bytes

MHAは4 GiB

MHAでは\(H_{kv}=32\)です。

\[M_{MHA}=2\times32\times1\times8192\times32\times128\times2\]
\[M_{MHA}=4,294,967,296\ \mathrm{bytes}=4\ \mathrm{GiB}\]

GQAは1 GiB

\(H_{kv}=8\)のGQAでは、group sizeは4です。

\[M_{GQA}=2\times32\times1\times8192\times8\times128\times2\]
\[M_{GQA}=1,073,741,824\ \mathrm{bytes}=1\ \mathrm{GiB}\]

MHAに対してKV Cache payloadは1/4です。

MQAは128 MiB

MQAでは\(H_{kv}=1\)です。

\[M_{MQA}=2\times32\times1\times8192\times1\times128\times2\]
\[M_{MQA}=134,217,728\ \mathrm{bytes}=128\ \mathrm{MiB}\]

MHAに対して1/32です。

方式 \(H_q\) \(H_{kv}\) group size K/V payload
MHA 32 32 1 4 GiB
GQA 32 8 4 1 GiB
MQA 32 1 32 128 MiB

同じ32層8192 token条件でMHA 4 GiB、GQA 1 GiB、MQA 128 MiBとなる比較

これはK/V要素を一度ずつ数えた生payloadです。model weight、activation、Attention workspace、allocatorの予約・断片化、padding、cache metadataは含みません。したがってGQAでK/V payloadが1/4になっても、総VRAMが1/4になるとは限りません。

K/V projection parameterも減る

GQA/MQAでは、cacheだけでなくK/Vを生成するprojection matrixも狭くなります。hidden sizeを\(d_{model}\)とすると、biasを除くshapeは次のとおりです。

\[W_Q\in\mathbb{R}^{d_{model}\times(H_qd_h)}\]
\[W_K,W_V\in\mathbb{R}^{d_{model}\times(H_{kv}d_h)}\]

\(H_q\)を維持して\(H_{kv}\)だけを減らすため、\(W_Q\)とoutput projectionは同条件なら維持され、\(W_K\)と\(W_V\)が小さくなります。

\(d_{model}=4096\)、\(H_q=32\)、\(d_h=128\)の1 Attention layerで、K/V projection parameterを合計すると次の値です。

方式 \(H_{kv}\) K/V projection parameter
MHA 32 33,554,432
GQA 8 8,388,608
MQA 1 1,048,576

これはmodel全parameterの比較ではありません。MLP、embedding、Q projection、output projectionなどは残るため、model file全体が同じ比率で小さくなるわけではありません。

Query projectionの幅は維持しK/V projectionだけがMHAからGQA、MQAへ狭くなる図

品質・速度・容量のtrade-off

KV head数を減らすと、cache容量と読出し量を減らせます。一方、複数Query headが利用できるK/V表現の種類も減ります。

比較軸 MHA GQA MQA
K/V表現の数 多い 中間 1組
KV Cache 大きい group数に応じて削減 最小
decodeのK/V読出し 多い 中間 少ない
品質の余地 K/V headごとに独立 共有と多様性の折衷 共有が最も強い
kernel実装 広く対応 group mapping対応が必要 共有K/V対応が必要

MQA原論文は、WMT14翻訳とBillion Word language modelingの実験で、baselineに対する品質差を小さく抑えつつincremental decodeを高速化できたと報告しています。ただし、当時のmodel、TPU、系列長、batch、実装における結果です。

GQA論文は、uptrained GQAがMHAに近い品質とMQAに近い速度を得たと報告しています。これも論文のmodelと評価条件での結果であり、「GQAなら常に品質低下なし」「MQAと同じ速度」とは一般化できません。

実際の選択では、KV head数だけでなく次を同条件で測ります。

  • task別のloss、perplexity、生成品質
  • prefill latencyとdecode latency
  • inter-token latencyとthroughput
  • request数、context長ごとのpeak memory
  • runtimeが選んだAttention kernel

MHA checkpointは設定変更だけでGQAにならない

既存MHA modelのnum_key_value_headsだけを書き換えても、K/V projectionのshapeが合いません。さらに、異なるK/V headを共有表現へまとめるため、modelが新しい構造へ適応する必要があります。

GQA論文のuptrainingは、概念的に次の2段階です。

  1. 同じgroupに入るMHAのK/V projection headをmean poolingし、少数のK/V headへ変換する
  2. 変換したcheckpointを元のpretraining recipeの一部で追加学習する

MQAへ変換する場合は、全K/V headを1組へmean poolingします。GQAならgroup内だけを平均します。論文は元のpretraining computeの5%を用いたrecipeを評価しましたが、この5%はすべてのmodelに必要十分な値ではありません。

MHAのK/V projectionをgroup内で平均してGQA checkpointへ変換し追加学習する流れ

modelを新規学習する場合は、初めからGQA/MQAのprojection shapeで学習できます。既存checkpointを変換する場合は、weight変換、追加学習、品質評価を一つの工程として扱います。

model configで方式を確認する

Hugging Face Transformersの現行LlamaConfig documentでは、num_attention_headsとnum_key_value_headsから方式を区別できます。

config条件 方式
num_key_value_heads == num_attention_heads MHA
1 < num_key_value_heads < num_attention_heads GQA
num_key_value_heads == 1 MQA

configを確認する最小例です。

from transformers import AutoConfig

model_id = "your-model-id"
config = AutoConfig.from_pretrained(model_id)

query_heads = config.num_attention_heads
kv_heads = getattr(config, "num_key_value_heads", None) or query_heads

# headの対応を曖昧にしないため、group sizeを出す前に割り切りを検証します。
if query_heads % kv_heads != 0:
    raise ValueError(
        "num_attention_heads must be divisible by num_key_value_heads"
    )

if kv_heads == query_heads:
    attention_type = "MHA"
elif kv_heads == 1:
    attention_type = "MQA"
else:
    attention_type = "GQA"

group_size = query_heads // kv_heads
print(
    f"type={attention_type}, Hq={query_heads}, "
    f"Hkv={kv_heads}, group_size={group_size}"
)

古いconfigや別architectureではfield名・意味が異なる可能性があるため、対象modelの公式documentと実装も確認してください。

PyTorch SDPAでshapeを確認する

現行PyTorchのscaled_dot_product_attentionはenable_gqa=Trueを提供しています。QueryとK/Vのhead数を別shapeで渡します。

import torch
import torch.nn.functional as F

batch_size = 2
query_heads = 32
kv_heads = 8
sequence_length = 1024
head_dimension = 128

query = torch.randn(
    batch_size,
    query_heads,
    sequence_length,
    head_dimension,
    device="cuda",
    dtype=torch.float16,
)
key = torch.randn(
    batch_size,
    kv_heads,
    sequence_length,
    head_dimension,
    device="cuda",
    dtype=torch.float16,
)
value = torch.randn_like(key)

# K/Vを手動複製せず、GQA対応backendへhead mappingを伝えるために有効化します。
output = F.scaled_dot_product_attention(
    query,
    key,
    value,
    dropout_p=0.0,
    is_causal=True,
    enable_gqa=True,
)

公式documentが示す主なhead条件は次のとおりです。

\[H_q\bmod H_{kv}=0\]
\[H_k=H_v\]

enable_gqaはexperimental機能で、backend、device、Nested Tensorなどの対応条件はversionにより変わります。reference実装の説明にK/Vのrepeat_interleaveが現れても、cacheを物理的にHq個へ恒久複製することがGQAの目的ではありません。実際のfused kernelが使われたか、fallbackや一時的な展開がないかをprofilerとmemory計測で確認します。

KV Cache・FlashAttention・量子化との違い

技術 変える対象 主に減らすもの
KV Cache 過去K/Vを保存して再利用 decode時のK/V再計算
GQA/MQA Query headに対するKV head数 KV Cache payloadとK/V読出し
FlashAttention Attention内部の計算順序とdata movement N×N中間値とHBM IO
量子化 weight、activation、KV Cache等のbit幅 要素あたりのbyte数

これらは排他的ではありません。GQA modelでKV Cacheを使い、対応kernelでFlashAttentionを実行し、weightやcacheを量子化する構成もあり得ます。

FlashAttentionの記事では、KV head数ではなくN×N中間値をHBMへ保存しない計算順序を扱っています。量子化の記事では、head数ではなく値のbit幅を減らします。

導入・model選択のチェックリスト

  1. num_attention_headsとnum_key_value_headsを確認する
  2. \(H_q/H_{kv}\)からgroup sizeを計算し、割り切れることを確認する
  3. KV Cacheのdtype、layer数、context長を含めてpayloadを概算する
  4. 総VRAMとKV Cache部分の削減率を分ける
  5. training、prefill、decodeを分けてlatencyとmemoryを測る
  6. runtimeとAttention kernelが対象GQA/MQA shapeへ対応するか確認する
  7. fallbackやK/Vの物理複製が起きていないかprofilerで確認する
  8. perplexityだけでなく、実運用taskの生成品質を比較する
  9. MHA checkpoint変換ではweight変換と追加学習を計画する

KV head数の少なさだけで方式を選ばず、品質条件と同時request数、context長、hardware、runtimeを同じ測定条件へ揃えます。

よくある質問

GQAではQuery headも減りますか

通常は減らしません。Query head数\(H_q\)を維持し、Key/Value head数\(H_{kv}\)を減らします。そのため各Query headは別のQとattention weightを持てます。

GQAのgroup数はどう決まりますか

group数はKV head数と同じです。\(H_q=32\)、\(H_{kv}=8\)なら8 groupで、1 groupのQuery head数は4です。最適値はmodelの学習・品質・推論条件に依存します。

GQAにするとAttention計算も4分の1になりますか

KV CacheとK/V読出し、K/V projectionには\(H_{kv}\)削減の効果があります。一方、Query head数と各Queryのscore・weighted sumは残るため、Attention全体のFLOPsやlatencyが同じ比率で減るとは限りません。

既存MHA modelのconfigだけ変更できますか

できません。K/V projection weightのshapeが変わります。GQA論文のようにgroup内headをmean poolingしてcheckpointを変換し、追加学習で適応させる方法がありますが、必要な学習量と品質はmodelごとに評価します。

MQAとGQAはKV Cacheの代わりですか

いいえ。KV Cacheを使う前提で、cacheに保存するK/V head数を減らします。過去K/Vを保存しない場合、decode時の再計算を避けるというKV Cacheの利点は得られません。

まとめ

GQAとMQAは、複数のQuery headを残しながらK/V headを共有し、KV Cacheとdecode時のK/V読出しを減らすAttention改良です。MHAは\(H_{kv}=H_q\)、MQAは\(H_{kv}=1\)、GQAはその中間に位置します。

理解の要点は、group数\(H_{kv}\)とgroup size\(H_q/H_{kv}\)を分けることです。KV Cacheの削減比は\(H_{kv}/H_q\)で概算できますが、総VRAM、Attention全体のFLOPs、latency、品質は同じ比率では変わりません。

model configではQuery/KV head数を確認し、runtimeでは対応kernelとfallbackを確認します。既存MHA checkpointを変換する場合は、projectionのmean poolingだけで終わらせず、追加学習とtask別品質評価まで含めて判断します。

参考文献

スポンサーリンク

コメント

タイトルとURLをコピーしました