Multi-Query / Grouped-Query Attention（MQA/GQA）— KVキャッシュを削ってLLM推論を軽くする

長い文章をLLMに書かせていると、最初の応答が出るまでは待たされるのに、そこから先はスラスラと文字が流れてきます。ところが入力の文脈がどんどん長くなると、今度はその「スラスラ」自体がだんだん重くなっていきます。文脈が10万トークンにもなると、GPUのメモリがぎりぎりになり、複数ユーザーを同時にさばくとさらに苦しくなります。この苦しさの正体は、実は計算量ではなく メモリ帯域 です。

その帯域を食いつぶす最大の犯人が KVキャッシュ です。標準的な Multi-Head Attention（MHA） は、ヘッドの数だけ独立したKey・Valueを持ちます。ヘッドが多いほど表現力は上がりますが、そのぶんキャッシュも膨らみ、1トークン生成するたびに大量のKey・Valueをメモリから読み出さなければなりません。Multi-Query Attention（MQA） と Grouped-Query Attention（GQA） は、この「ヘッドごとに別々のK,Vを持つ」という前提そのものを崩し、K,Vを共有することでキャッシュを劇的に減らす手法です。

MQA/GQAを理解すると、次の2つの問いに明確に答えられるようになります。

• なぜLlama 2 70BやMistralはMHAではなくGQAを採用したのか、それでどれだけメモリと速度が改善するのか
• K,Vを共有すると品質はどうなるのか、そのトレードオフをどこで折り合わせるのか

本記事の内容

• なぜ生成はメモリ帯域律速になるのか（KVキャッシュと演算強度の話）
• MQAの仕組み — 全ヘッドでK,Vを1組だけ共有する
• GQAの仕組み — グループ単位でK,Vを共有し、MHAとMQAの中間を取る
• KVキャッシュ量の式とMHA/GQA/MQAでの比較
• 学習済みMHAをGQAに作り替える「uptraining」
• PyTorchでの実装と、テンソル形状・キャッシュ量・デコード時間の実測

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

MQA/GQAは、この3つの続編です。MHAでヘッドを分ける仕組み、KVキャッシュで過去のK,Vを貯める仕組みを前提に、「そのK,Vをどこまで共有できるか」を掘り下げます。

なぜ生成はメモリ帯域律速になるのか

MQA/GQAの狙いを腹落ちさせるには、まず「LLMの生成が何で遅くなっているのか」を正確に押さえる必要があります。素朴には「計算量が多いから遅い」と思いがちですが、自己回帰生成のボトルネックはむしろ メモリからデータを運ぶ速度 にあります。

生成は「行列×ベクトル」で演算強度が低い

学習やプロンプトの一括処理（プリフィル）では、多数のトークンをまとめて処理します。これは「行列×行列」の演算で、1バイトのデータを読むごとに多くの浮動小数点演算をこなせます。GPUのテンソルコアが最も得意とする形です。

一方、1トークンずつ吐き出す生成フェーズ（デコード）では、いま処理しているのは たった1個のトークン です。この1トークンのQueryを、キャッシュに貯まった大量のK,Vと突き合わせます。演算の中身は「1本のベクトル × 大きな行列」、つまり 行列×ベクトル です。

この違いを「演算強度」（1バイトのデータ読み出しあたり何回の演算をするか）という指標で見ると、両者は真逆の位置にいます。生成は演算強度が低く、GPUの演算器はスカスカに空いたまま、K,Vをメモリから読み出す時間だけが伸びていきます。これが メモリ帯域律速 です。

この図は横軸に演算強度、縦軸に到達スループットを取ったルーフライン図です。左側の斜めの領域（赤）は「帯域律速」で、演算強度が低いほどスループットが頭打ちになります。右側の平らな領域（青）は「演算律速」で、演算器の性能が上限になります。自己回帰生成は左端の赤い点にいて、GPUの演算ピークにまるで届いていません。逆に学習・プリフィルは右上の青い点にいて演算律速です。つまり生成を速くしたいなら、演算を減らすのではなく メモリから運ぶバイト数を減らす のが正解だと読み取れます。

運ぶバイト数を支配するのがKVキャッシュ

では生成中に何をいちばん多く運んでいるのでしょうか。それがKVキャッシュです。1トークン生成するたびに、全層・全ヘッドのK,Vをメモリから読み出してQueryと突き合わせます。文脈が長いほど、バッチが大きいほど、この読み出し量はふくらみます。

KVキャッシュの記事 で見たとおり、KVキャッシュの大きさは次の式で決まります。

$$ \begin{equation} \text{Memory}_{\text{KV}} = 2 \times L \times H_{kv} \times d_k \times T \times B \times \text{bytes} \end{equation} $$

ここで $L$ は層数、$H_{kv}$ はKVヘッド数、$d_k$ は1ヘッドあたりの次元、$T$ は系列長、$B$ はバッチサイズ、$\text{bytes}$ はデータ型のバイト数（fp16なら2）です。$2$ はKeyとValueの2種類ぶんです。

この式をよく見ると、生成の重さを直接に左右する項のうち、$L$（層数）や $d_k$（次元）はモデルの表現力に直結していて簡単には削れません。$T$ や $B$ はユーザーの使い方で決まる量です。ところが $H_{kv}$（KVヘッド数）だけは事情が違います。「ヘッドごとに別々のK,Vを持たなければならない」という前提を緩めれば、ここを大きく削れる のです。これがMQA/GQAの着眼点です。

次の節から、その $H_{kv}$ を減らす具体的な仕組みを見ていきます。

MHAの復習 — ヘッドごとに独立したK,Vを持つ

出発点として、標準のMulti-Head Attentionを $H_{kv}$ の視点で振り返っておきましょう。詳しくは Multi-Head Attention の記事にありますが、要点は「Query・Key・Valueをそれぞれ $h$ 個のヘッドに分ける」ことです。

入力 $\bm{x}$ から、ヘッド $i$（$i = 1, \dots, h$）ごとに独立した射影行列でQ,K,Vを作ります。

$$ \begin{equation} \bm{q}_i = \bm{x}\bm{W}^Q_i, \quad \bm{k}_i = \bm{x}\bm{W}^K_i, \quad \bm{v}_i = \bm{x}\bm{W}^V_i \quad (i = 1, \dots, h) \end{equation} $$

各ヘッドは自分専用のK,Vを使ってAttentionを計算し、それを連結して出力します。

$$ \begin{equation} \text{head}_i = \text{softmax}\!\left(\frac{\bm{q}_i \bm{k}_i^\top}{\sqrt{d_k}}\right)\bm{v}_i, \quad \text{MHA}(\bm{x}) = \text{concat}(\text{head}_1, \dots, \text{head}_h)\,\bm{W}^O \end{equation} $$

ここで大事なのは、ヘッドの数だけK,Vの組が存在するという点です。ヘッド $i$ のQueryはヘッド $i$ のKey・Valueとしか照合しません。1対1対応です。したがって $H_{kv} = h$ となり、KVキャッシュは「Queryヘッド数」と同じだけ膨らみます。

これは表現力の面では理にかなっています。各ヘッドが別々の部分空間に注目できるからです。しかし生成時には、$h$ 個ぶんのK,Vをすべてメモリから読み出す必要があり、まさにここが帯域を食います。

そこで自然な疑問が生まれます。「Queryはヘッドごとに違ってもよいが、K,Vまでヘッドごとに完全に別々である必要は本当にあるのか？」 この問いに「いや、K,Vは1組で足りる」と答えたのがMQAです。

MQA — 全ヘッドでK,Vを1組だけ共有する

直感 — 質問は人それぞれ、資料は共通でよい

図書館の比喩で考えてみましょう。Attentionでは、Queryが「探したい情報」、Key・Valueが「本棚に並ぶ資料とその中身」でした（MHAの記事 参照）。

MHAは、ヘッドごとに 別々の本棚 を用意するようなものです。文法担当ヘッドには文法用の本棚、意味担当ヘッドには意味用の本棚、というふうに。これは贅沢ですが場所を取ります。

Multi-Query Attention（Shazeer, 2019）の発想はこうです。本棚（K,V）は1つの共通のものを全員で使い、質問（Query）だけ人それぞれ変える。文法担当も意味担当も、同じ本棚を見に行きます。ただし「何を探すか」という問いは各自で違うので、同じ本棚からでも違う情報を引き出せます。

数式 — Queryだけ $h$ 個、K,Vは1個

MQAでは、Queryはヘッドごとに作りますが、Key・Valueは全ヘッドで共有する1組だけを作ります。

$$ \begin{equation} \bm{q}_i = \bm{x}\bm{W}^Q_i \quad (i = 1, \dots, h), \qquad \bm{k} = \bm{x}\bm{W}^K, \quad \bm{v} = \bm{x}\bm{W}^V \end{equation} $$

$\bm{W}^K, \bm{W}^V$ の添字 $i$ が消えていることに注目してください。全ヘッドで同じ $\bm{k}, \bm{v}$ を使います。各ヘッドのAttentionは次のように、共有K,Vに対して計算します。

$$ \begin{equation} \text{head}_i = \text{softmax}\!\left(\frac{\bm{q}_i \bm{k}^\top}{\sqrt{d_k}}\right)\bm{v} \quad (i = 1, \dots, h) \end{equation} $$

Queryは $h$ 種類あるので、$h$ 個の異なる注意分布が得られます。共有された同じ本棚から、ヘッドごとに違う読み方をするわけです。

効果 — KVキャッシュが $1/h$ に

KVキャッシュの式で $H_{kv} = h$ が $H_{kv} = 1$ になるので、キャッシュ量はきれいに $1/h$ に縮みます。ヘッド数が32のモデルなら 1/32 です。生成時に読み出すK,Vのバイト数も同じ比率で減るので、帯域律速のデコードがそのぶん速くなります。

代償は表現力です。全ヘッドが同じK,Vを見るため、「ヘッドごとに違う情報を蓄える」という自由度が失われます。実際、MQAはMHAよりわずかに品質が落ちること、学習が不安定になることがあると報告されています。この「品質の落ち込み」を和らげるために提案されたのがGQAです。

その前に、MQAの共有構造を図で確認しておきましょう。

この図は左からMHA・MQA・GQAのK,V共有パターンを示しています。上段の青い箱がQueryヘッド（8個で共通）、下段の緑の箱がKVヘッドです。左のMHAはQueryヘッド8個に対しKVヘッドも8個で1対1、キャッシュは基準の×1です。中央のMQAはKVヘッドがたった1個で、8本のQueryすべてがそこへ集まり、キャッシュは×1/8です。右のGQAはKVヘッドが2個で、Queryを4本ずつのグループに分けて共有し、キャッシュは×1/4です。矢印の集まり方を見れば、共有が強いほどKVヘッドが減る＝メモリが減ることが直感的に読み取れます。

GQA — グループ単位でK,Vを共有する

直感 — 全員で1つは極端、班ごとに1つで折り合う

MHAは「1人1本棚」、MQAは「全員で1本棚」でした。GQA（Grouped-Query Attention、Ainslie et al., 2023）は、その中間の「班ごとに1本棚」を取ります。

8人のヘッドを2つの班（グループ）に分け、各班4人で1つの本棚を共有する。班が違えば本棚も違うので、MQAよりは多様な情報を持てます。同時に、本棚の数は8個から2個に減るので、MHAよりはずっとメモリが軽くなります。品質と効率のちょうどいい妥協点を探る、というのがGQAの思想です。

数式 — $g$ 個のグループでK,Vを共有

$h$ 個のQueryヘッドを $g$ 個のグループに分けます（$g$ はKVヘッド数でもあります）。グループ $j$ には $h/g$ 個のQueryヘッドが属します。グループごとに1組のK,Vを作り、そのグループ内のQueryヘッドが共有します。

$$ \begin{equation} \bm{k}^{(j)} = \bm{x}\bm{W}^{K}_j, \quad \bm{v}^{(j)} = \bm{x}\bm{W}^{V}_j \quad (j = 1, \dots, g) \end{equation} $$

Queryヘッド $i$ が属するグループを $g(i)$ と書くと、そのヘッドのAttentionは自分のグループのK,Vを使います。

$$ \begin{equation} \text{head}_i = \text{softmax}\!\left(\frac{\bm{q}_i \,{\bm{k}^{(g(i))}}^\top}{\sqrt{d_k}}\right)\bm{v}^{(g(i))} \end{equation} $$

両端としてのMHAとMQA

この定義の美しいところは、$g$ を動かすだけでMHAもMQAも表現できる点です。

• $g = h$ のとき：各グループにQueryヘッドが1個ずつ入る。つまりK,Vもヘッドごとに独立 → MHAそのもの。
• $g = 1$ のとき：全ヘッドが1グループに入り、K,Vは1組だけ共有 → MQAそのもの。
• $1 < g < h$ のとき：その中間。GQAの本領。

したがってMHA・GQA・MQAは別々の3手法というより、「グループ数 $g$」という1本のつまみを回してつながった連続スペクトル だと理解するのがいちばん正確です。

この図は、Queryヘッド数を8で固定したまま、グループ数 $g$ を8→4→2→1と減らしていく様子を示しています。$g=8$ は右端のMHA、$g=1$ は左端のMQAに対応し、その間がGQAです。$g$ を減らすほどKVヘッドが減り、キャッシュメモリが $g/8$ に縮む一方で、共有が強まって表現力は下がります。左に行くほど表現力寄り、右に行くほどメモリ効率寄りという、トレードオフの向きが読み取れます。

効果 — キャッシュは $g/h$ に

KVヘッド数が $g$ なので、キャッシュ量はMHA比で $g/h$ になります。Llama 2 70Bは $h=64$ に対して $g=8$（グループサイズ8）を採用し、KVキャッシュを 1/8 に抑えています。Mistral 7Bも $h=32$、$g=8$ でKVキャッシュを1/4にしています。$g = h/8 \sim h/4$ 程度に取ると、MHAとの品質差はごくわずかで、MQAより明確に安定するというのが実験的な知見です。

ここまでで3手法の仕組みが揃いました。次は、この共有がKVキャッシュ量にどれだけ効くかを、具体的な数字で確かめます。

KVキャッシュ量を数字で比較する

抽象的な $g/h$ だけではピンと来ないので、Llama 2 70Bに近い設定で実際のバイト数を計算してみましょう。層数 $L=32$、1ヘッド次元 $d_k=128$、系列長 $T=4096$、バッチ $B=1$、fp16（2バイト）とします。QueryヘッドはMHA基準で32個とします。

キャッシュ量の式 $\text{Memory}_{\text{KV}} = 2 \times L \times H_{kv} \times d_k \times T \times B \times \text{bytes}$ に、$H_{kv}$ だけ変えて代入します。

• MHA（$H_{kv}=32$）：$2 \times 32 \times 32 \times 128 \times 4096 \times 1 \times 2 = 2.147$ GB
• GQA（$H_{kv}=8$）：$0.537$ GB（MHAの1/4）
• MQA（$H_{kv}=1$）：$0.067$ GB（MHAの1/32）

MHAでは1系列あたり約2.1GBだったキャッシュが、GQAで0.5GB、MQAでは67MB弱まで縮みます。これはあくまで1系列ぶんです。推論サーバーは何十もの並列リクエストをさばくので、この差がそのまま「同じGPUに何ユーザー乗せられるか」に直結します。

この棒グラフは上記の3設定のキャッシュ量を並べたものです。MHAの2.15GBに対し、GQAは0.54GB、MQAは0.07GBと、共有を強めるほど棒が劇的に低くなります。GQAでも4倍のメモリ節約が得られ、MQAに至ってはほぼ底を這っています。生成時に読み出すバイト数もこの比率なので、帯域律速のデコードがそのぶん軽くなると読み取れます。

系列長・バッチを変えるとどうなるか

キャッシュ量は系列長 $T$ とバッチ $B$ に線形なので、長文・大バッチではこの差が効いてきます。バッチを8にして系列長を伸ばした場合を見てみましょう。

この折れ線は、バッチ8のときの系列長（横軸）に対するキャッシュ量（縦軸）です。破線はRTX 3090（24GB）とA100（80GB）のメモリ上限を示します。MHA（赤）は系列長が伸びるとすぐに24GB・80GBの壁にぶつかりますが、GQA（緑）はずっと余裕があり、MQA（青）はほとんど張り付いています。同じGPUで扱える文脈長やバッチサイズが、共有の度合いで大きく変わることが分かります。

KVヘッド数と系列長を同時に振ると、削減効果の全体像がさらに見えます。

このヒートマップは、縦軸にKVヘッド数（下がMQAの1、上がMHAの32）、横軸に系列長を取り、セルの色でキャッシュ量（GB）を示しています。同じ系列長でも、下（MQA）から上（MHA）へ行くほど色が明るくなり、メモリが増えていきます。逆に言えば、上から下へKVヘッドを減らすだけで、系列長を変えずにメモリを大きく節約できると読み取れます。行方向の変化（KVヘッド削減）が、列方向の変化（系列長短縮）と同じくらい効いている点がポイントです。

数字でメモリ削減が確認できました。では、この削減が実際の「速さ」にどう効くのかを次に見ます。

デコード時間はKVヘッド数に比例する

生成が帯域律速だという話に立ち返ると、1ステップのデコード時間は「そのステップで読み出すK,Vのバイト数」にほぼ比例するはずです。KVヘッドを減らせば読むバイトが減るので、時間も減ります。これを簡単な実測で確かめます。

系列長 $T=4096$、$d_k=128$ の設定で、KVヘッド数を1から32まで変えながら「1トークンのQueryを全K,Vと突き合わせる」処理の時間を測ると、次のようになりました。

KVヘッドが2倍になるとほぼ時間も2倍で、きれいな比例関係が出ています。

この図は上の表を折れ線にしたものです。KVヘッド数（横軸）に対してKV読み出し時間（縦軸）が直線的に増えています。MQA（KV=1）はMHA（KV=32）に比べて約25倍速く、GQA（KV=8）でもMHAの約1/4の時間で済みます。生成が帯域律速だからこそ、読み出すK,Vのバイト数を減らすことがそのまま速度に効く、という因果がはっきり見えます。

速度とメモリの利点は明快ですが、共有すれば品質が下がるのも事実です。そのトレードオフを整理します。

品質とのトレードオフ

MHAは表現力最大・メモリ最大、MQAはメモリ最小・表現力最小、GQAはその間です。実務で問われるのは「どこまで共有しても品質が保てるか」です。

Ainslieらの実験では、MQAはMHAに比べて品質がやや落ちるのに対し、$g = h/8$ 程度のGQAはMHAとほぼ同等の品質を、MQAとほぼ同等の速度で達成できると報告されています。つまりGQAは「MQAの速さを、MHAの品質を大きく損なわずに手に入れる」おいしい妥協点です。

この図は横軸にKVヘッド数（左ほどメモリ効率が高い）、縦軸にモデル品質（相対・定性）を取った概念図です。右端のMHAから左へKVヘッドを減らしていくと、しばらくは品質がほぼ横ばいのままメモリだけが減り、GQA（Llama 2の$g=8$相当）付近までは品質がよく保たれます。ところが左端のMQA（KV=1）に近づくと品質がやや落ち込みます。「$g$ をある程度まで減らすのはほぼ無料、そこから先は品質を少し払う」という曲線の形が、GQAが好まれる理由をよく表しています。

なぜ共有すると品質が落ちるのかは、各ヘッドが「見に行く先」を見ると分かります。

この図は、ランダム初期化した小さなモデルで、最終トークンの各Queryヘッドが過去トークンにどう注意を配るか（注意重み）をヒートマップにしたものです。左のMHAでは4つのQueryヘッドがそれぞれ別々のK,Vを見るため、行ごとにパターンが大きく違います。中央のGQAでは2ヘッドずつが同じK,Vを共有するため、ペアで似たパターンが現れます。右のMQAでは全ヘッドが同じK,Vを見るので、Queryの違いで多少ずれるものの、パターンが揃ってきます。共有が強いほど「ヘッドの多様性」が減る様子が読み取れ、これが品質低下の源だと理解できます。

品質を回復するには、共有K,Vをきちんと学習し直すのが理想です。ただしゼロから学習するのは高価です。そこで既存のMHAモデルを賢く作り替える方法が使われます。

uptraining — 学習済みMHAをGQAに作り替える

GQAの実務上の魅力は、すでに大金をかけて学習したMHAモデルを、そのまま活かしてGQA化できる ことにあります。この変換をuptrainingと呼びます。

手順はシンプルです。GQAの各グループには、元のMHAで複数のヘッドが対応します。そこで、そのグループに属する複数のヘッドのK射影・V射影を 平均プーリング して、共有K,Vの初期値にします。

$$ \begin{equation} \bar{\bm{W}}^K_j = \frac{1}{|G_j|}\sum_{i \in G_j} \bm{W}^K_i, \quad \bar{\bm{W}}^V_j = \frac{1}{|G_j|}\sum_{i \in G_j} \bm{W}^V_i \end{equation} $$

ここで $G_j$ はグループ $j$ に属するヘッドの集合、$|G_j| = h/g$ です。平均を初期値にするのがポイントで、いきなり1つのヘッドを選ぶより、グループ内の情報をまんべんなく引き継げます。

初期化しただけでは品質が少し落ちるので、そのあと 元の学習データのごく一部（数%）で追加学習 します。ゼロから学び直すより桁違いに安価で、これでMHAに近い品質のGQAモデルが手に入ります。

この図はuptrainingの流れを示しています。左の学習済みMHAには、1グループに対応する4つのK射影 $\bm{W}^K_1 \dots \bm{W}^K_4$ があります。これらを平均プーリングして、右のGQAの共有K射影 $\bar{\bm{W}}^K$ の初期値にします。その後、全データのわずか数%でuptrainingすると、初期化のずれが吸収されて品質が回復します。学習済みの資産を捨てずにGQA化できる、という実務上の強みが読み取れます。

これで仕組みは一通り揃いました。最後に、MHA/GQA/MQAを実際にPyTorchで実装し、形状とキャッシュ量を自分の手で確かめます。

PyTorchでの実装と検証

MHA・GQA・MQAは、KVヘッド数 $H_{kv}$ を変えるだけの1つの実装で書けます。共有の肝は、少ないK,Vヘッドをグループ内のQueryヘッド数だけ複製して形を合わせる repeat_interleave の一行です。

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


def attention(q, k, v):
    # q: (B, Hq, T, dk), k,v: (B, Hkv, T, dk)  ※Hkv <= Hq
    B, Hq, T, dk = q.shape
    Hkv = k.shape[1]
    rep = Hq // Hkv
    # KVヘッドをグループ内Queryヘッド数だけ複製して形を合わせる
    k = k.repeat_interleave(rep, dim=1)  # (B, Hq, T, dk)
    v = v.repeat_interleave(rep, dim=1)
    scores = (q @ k.transpose(-1, -2)) / (dk ** 0.5)  # (B, Hq, T, T)
    attn = F.softmax(scores, dim=-1)
    return attn @ v  # (B, Hq, T, dk)


class GroupedAttention(nn.Module):
    def __init__(self, d_model, n_q_heads, n_kv_heads):
        super().__init__()
        assert n_q_heads % n_kv_heads == 0
        self.hq, self.hkv = n_q_heads, n_kv_heads
        self.dk = d_model // n_q_heads
        self.Wq = nn.Linear(d_model, n_q_heads * self.dk, bias=False)
        self.Wk = nn.Linear(d_model, n_kv_heads * self.dk, bias=False)  # KVは少ない
        self.Wv = nn.Linear(d_model, n_kv_heads * self.dk, bias=False)
        self.Wo = nn.Linear(n_q_heads * self.dk, d_model, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.Wq(x).view(B, T, self.hq, self.dk).transpose(1, 2)
        k = self.Wk(x).view(B, T, self.hkv, self.dk).transpose(1, 2)
        v = self.Wv(x).view(B, T, self.hkv, self.dk).transpose(1, 2)
        print(f"  Q={tuple(q.shape)}  K={tuple(k.shape)}  V={tuple(v.shape)}")
        o = attention(q, k, v)                      # (B, Hq, T, dk)
        o = o.transpose(1, 2).reshape(B, T, self.hq * self.dk)
        return self.Wo(o)

Wk・Wv の出力次元が n_kv_heads * dk になっている点が要です。ここが n_q_heads より小さいぶんだけ、K,Vの射影パラメータもキャッシュも減ります。では同じ入力に対して、KVヘッド数だけ変えて3手法を動かし、テンソル形状を見てみます。

B, T, d_model, Hq = 2, 5, 64, 8
x = torch.randn(B, T, d_model)
print("入力 x:", tuple(x.shape))

print("[MHA] Hq=8, Hkv=8")
o1 = GroupedAttention(d_model, Hq, 8)(x); print("  出力:", tuple(o1.shape))
print("[GQA] Hq=8, Hkv=2")
o2 = GroupedAttention(d_model, Hq, 2)(x); print("  出力:", tuple(o2.shape))
print("[MQA] Hq=8, Hkv=1")
o3 = GroupedAttention(d_model, Hq, 1)(x); print("  出力:", tuple(o3.shape))

実行結果:

入力 x: (2, 5, 64)
[MHA] Hq=8, Hkv=8
  Q=(2, 8, 5, 8)  K=(2, 8, 5, 8)  V=(2, 8, 5, 8)
  出力: (2, 5, 64)
[GQA] Hq=8, Hkv=2
  Q=(2, 8, 5, 8)  K=(2, 2, 5, 8)  V=(2, 2, 5, 8)
  出力: (2, 5, 64)
[MQA] Hq=8, Hkv=1
  Q=(2, 8, 5, 8)  K=(2, 1, 5, 8)  V=(2, 1, 5, 8)
  出力: (2, 5, 64)

出力形状はどれも入力と同じ (2, 5, 64) で、3手法とも問題なく計算できています。注目すべきはKとVの2番目の軸（ヘッド数）です。MHAは 8、GQAは 2、MQAは 1 と、Queryの 8 はそのままにK,Vだけが減っています。キャッシュに貯めるのはこのK,Vなので、この軸の縮小がそのままメモリ削減になる わけです。Queryヘッド数は変わらないので、出力の表現力（ヘッドの本数）は保たれる点も確認できます。

キャッシュ量を実際に計算する

続いて、Llama 2 70B相当の設定でキャッシュのバイト数を計算し、先ほどの図の数字を再現します。

def kv_bytes(L, hkv, dk, T, B, byt=2):
    return 2 * L * hkv * dk * T * B * byt

L, dk, T, Bb = 32, 128, 4096, 1
for name, hkv in [("MHA", 32), ("GQA", 8), ("MQA", 1)]:
    b = kv_bytes(L, hkv, dk, T, Bb)
    print(f"{name}(KV={hkv:2d}): {b/1e9:.3f} GB  (MHA比 {hkv/32:.3f})")

実行結果:

MHA(KV=32): 2.147 GB  (MHA比 1.000)
GQA(KV= 8): 0.537 GB  (MHA比 0.250)
MQA(KV= 1): 0.067 GB  (MHA比 0.031)

計算結果は棒グラフと完全に一致します。GQAでMHAの1/4（0.537GB）、MQAで1/32（0.067GB）です。KVヘッド数を式に入れるだけで、これだけの差が出ることを自分の手で確認できました。

実装・形状・キャッシュ量がすべて理論どおりに揃いました。最後に要点をまとめます。

まとめ

本記事では、MQAとGQAがKVキャッシュを削ってLLM推論を軽くする仕組みを解説しました。

• 生成は帯域律速: 自己回帰生成は行列×ベクトルで演算強度が低く、GPUの演算器ではなくメモリ帯域がボトルネック。速くするには演算でなく「運ぶバイト数」を減らす
• 運ぶバイトの主役はKVキャッシュ: 大きさは $2 L H_{kv} d_k T B \times \text{bytes}$。このうちKVヘッド数 $H_{kv}$ は、K,V共有で削れる唯一の項
• MQA: 全ヘッドでK,Vを1組だけ共有。キャッシュを $1/h$ に激減させるが、表現力がやや落ちる
• GQA: $g$ グループでK,Vを共有。$g=h$ でMHA、$g=1$ でMQAとなる連続スペクトルの中間点。$g=h/8$ 程度でMHA並みの品質とMQA並みの速度を両立し、Llama 2 70BやMistralが採用
• 実測の一致: PyTorchでK,Vのヘッド軸だけが縮むことを形状で確認。キャッシュ量はMHA 2.15GB → GQA 0.54GB → MQA 0.07GB、デコード時間もKVヘッド数に比例
• uptraining: 学習済みMHAの複数ヘッドの射影を平均プーリングしてGQAを初期化し、少量再学習するだけで品質を回復できる

MQA/GQAは、KVキャッシュのメモリと帯域を根本から削る「アーキテクチャ側の対策」です。これに、キャッシュの管理を効率化するPagedAttentionや、K,Vをint8/int4に落とすKV量子化などの「実装側の対策」を組み合わせることで、長文・大バッチのLLM推論が現実的な速度で回るようになります。

次のステップとして、以下の記事も参考にしてください。