Attentionの記事で、Cross-AttentionのKVキャッシュは一度計算すれば増えないと書いた。その「増える／増えない」の実体を、自分の手で測ってみたことがなかった。この記事で考えてほしいのは、KVキャッシュを小さくする方法には性質の違う2つの対策があるということだ。1つはキャッシュする量そのものを減らす方法（アーキテクチャ側）、もう1つは持っているキャッシュの使い方を効率化する方法（サービング側）。手元の環境はコンシューマー向けハイエンド構成の一例（VRAM 32GB級のGPU）で、大きなモデルは動かせない。だからNumPyだけでキャッシュありなしの生成を再現し、速度差を自分の目で確認するところから始めた。
この記事の対象読者
- Self-Attentionの基本（Query・Key・Value）を理解している人
- LLMのサービング（vLLM等）でバッチサイズが伸びない、VRAMがすぐ尽きると感じたことがある人
- MQA・GQAという言葉は知っているが、KVキャッシュのどの部分を削っているのか説明できない人
この記事で得られること
KVキャッシュを削る2つの対策──ヘッド数を減らしてキャッシュそのものを小さくする方法と、既存のキャッシュの持ち方を効率化する方法──のうち、自分の状況にどちらが合うかを判断できるようになります。
- なぜKVキャッシュがないと生成が遅くなるかを、実測したコードで説明できるようになります
- MHA・MQA・GQAがKVキャッシュのサイズにどう影響するかを自分で計算できるようになります
- PagedAttentionがKVキャッシュのメモリ管理をどう変えたかを説明できるようになります
この記事で扱わないこと
- vLLM等のサービングエンジンの実際のインストール・デプロイ手順
- Flash Attentionなど、計算そのものを速くする手法の実装詳細
- Attention層の「数」を削るアプローチ（Jambaの記事を参照してください）
1. なぜキャッシュがないと生成は遅くなるのか
このセクションで分かること：KVキャッシュがない場合、生成のたびに何が無駄に繰り返されているか、そしてその無駄が実際にどれだけの速度差を生むか。
LLMは1トークンずつ自己回帰的に文章を生成する。Self-Attention層は、新しいトークンを予測するたびに、それまでの全トークンに対してKey・Valueベクトルを計算する必要がある。ここで愚直な実装をすると、すでに計算したはずの過去のトークン分のKey・Valueまで、毎回ゼロから計算し直すことになる。KVキャッシュは、この過去分の計算結果を保存しておき、新しいトークンの分だけ追記する仕組みだ。
実際にどれだけ無駄が生じるのか、NumPyで比較する。
import time
import numpy as np
np.random.seed(0)
def naive_generate(seq_len, d_model, n_layers):
"""毎ステップ、これまでの全トークン分のK・Vを再計算する（キャッシュなし）"""
Ws = [(np.random.randn(d_model, d_model), np.random.randn(d_model, d_model),
np.random.randn(d_model, d_model)) for _ in range(n_layers)]
for t in range(1, seq_len + 1):
seq = np.random.randn(t, d_model)
for (W_q, W_k, W_v) in Ws:
K = seq @ W_k # tトークン分のKを毎回作り直す
V = seq @ W_v # tトークン分のVを毎回作り直す
Q = seq[-1:] @ W_q
scores = Q @ K.T
_ = scores @ V
def cached_generate(seq_len, d_model, n_layers):
"""KVキャッシュを使い、新規トークン分だけK・Vを計算して追記する"""
Ws = [(np.random.randn(d_model, d_model), np.random.randn(d_model, d_model),
np.random.randn(d_model, d_model)) for _ in range(n_layers)]
K_cache = [np.zeros((0, d_model)) for _ in range(n_layers)]
V_cache = [np.zeros((0, d_model)) for _ in range(n_layers)]
for t in range(1, seq_len + 1):
x = np.random.randn(1, d_model)
for l, (W_q, W_k, W_v) in enumerate(Ws):
k_new, v_new = x @ W_k, x @ W_v # 新規トークン1つ分だけ計算
K_cache[l] = np.vstack([K_cache[l], k_new])
V_cache[l] = np.vstack([V_cache[l], v_new])
Q = x @ W_q
scores = Q @ K_cache[l].T
_ = scores @ V_cache[l]
d_model, n_layers = 256, 4
for seq_len in [100, 200, 400]:
t0 = time.perf_counter(); naive_generate(seq_len, d_model, n_layers); t_naive = time.perf_counter() - t0
t0 = time.perf_counter(); cached_generate(seq_len, d_model, n_layers); t_cached = time.perf_counter() - t0
print(f"seq_len={seq_len}: naive={t_naive:.3f}s cached={t_cached:.3f}s speedup={t_naive/t_cached:.1f}x")
計測環境はCPU実行・NumPyのみ（GPU不使用）、d_model=256、n_layers=4。実行結果は次のとおり。
| 生成トークン数 | キャッシュなし | キャッシュあり | 速度差 |
|---|---|---|---|
| 100 | 0.208秒 | 0.061秒 | 3.4倍 |
| 200 | 0.558秒 | 0.125秒 | 4.5倍 |
| 400 | 2.076秒 | 0.346秒 | 6.0倍 |
生成トークン数が伸びるほど、速度差そのものが3.4倍→4.5倍→6.0倍と拡大している。これは偶然ではない。キャッシュなしの実装は、t番目のトークンを生成するたびにt個分のKey・Value射影をやり直すので、無駄な計算量は生成が進むほど二乗のペースで積み上がる。キャッシュありの実装は、新規トークン1つ分の射影を追記するだけなので、無駄な計算量は増えない。
Cross-Attentionの場合はさらに事情が単純で、Key・ValueはEncoder側の出力から一度だけ計算すればよく、Decoderが何トークン生成してもKey・Valueの計算自体は増えない（詳しくはAttentionの記事を参照）。増え続けるのはSelf-Attention側のキャッシュだけ、という点は覚えておいて損はない。
2. KVキャッシュのサイズはヘッド数で決まる ── MHA・MQA・GQA
前節で、キャッシュがないと計算量が無駄に積み上がると確認した。ではキャッシュを使うとして、そのキャッシュ自体のサイズは何で決まるのか。これを理解しないと、VRAMが足りないときに何を削ればいいのか判断できない。このセクションで分かること：Key・Valueの「ヘッド数」を変えるだけで、キャッシュサイズが最大32分の1まで縮むという実測。
KVキャッシュのサイズは、以下の式で決まる。
\text{KVキャッシュ（バイト）} = 2 \times L \times H_{kv} \times D_{head} \times S \times B \times \text{bytes}
$L$はAttention層数、$S$は文脈長、$B$はバッチサイズ、bytesは精度（16bitなら2）を表す。Jambaの記事では$L$（Attention層の数）を減らすレバーを扱った。この記事で見るのは、もう1つのレバー、$H_{kv}$（Key・Value用のヘッド数）だ。
| 記号 | 読み | たとえでの役割 | 型 |
|---|---|---|---|
| $H$ | エイチ | Query側のヘッド数（分業する担当者の人数） | 整数 |
| $H_{kv}$ | エイチ・ケーブイ | Key・Value側のヘッド数（実際に原本を持つ担当者の人数） | 整数 |
通常のMulti-Head Attention（MHA）では $H_{kv} = H$ で、担当者全員が自分専用の原本（Key・Value）を持つ。2019年、Shazeerは担当者全員に原本を持たせる必要はなく、1人の担当者が持つ原本を全員で共有すればよいと提案した。これがMulti-Query Attention（MQA、$H_{kv}=1$）で、論文は「推論時にKeyとValueという大きなテンソルを繰り返し読み込むメモリ帯域コストが、逐次的な生成を遅くしている」という問題意識を明確に述べている。2023年のGQA（Grouped-Query Attention）は、その中間として担当者を数グループに分け、グループ内でだけ原本を共有する方式（$1 < H_{kv} < H$）を提案した。
hidden_size=4096、num_attention_heads=32、文脈長256K、全32層がAttentionという条件で、$H_{kv}$だけを変えて計算する。
def kv_cache_bytes(num_attn_layers, num_kv_heads, head_dim, seq_len, batch=1, dtype_bytes=2):
return 2 * num_attn_layers * num_kv_heads * head_dim * seq_len * batch * dtype_bytes
def to_gb(b):
return b / (1024 ** 3)
hidden_size, num_attention_heads = 4096, 32
head_dim = hidden_size // num_attention_heads # 128
seq_len, num_layers = 262144, 32
configs = {
"MHA（num_kv_heads=32）": 32,
"GQA（num_kv_heads=8、Mixtral等）": 8,
"GQA（num_kv_heads=4）": 4,
"MQA（num_kv_heads=1）": 1,
}
for name, h_kv in configs.items():
b = kv_cache_bytes(num_layers, h_kv, head_dim, seq_len)
print(f"{name}: {to_gb(b):.1f} GB")
| 方式 | $H_{kv}$ | KVキャッシュ（256K文脈、全32層Attention） |
|---|---|---|
| MHA | 32 | 128.0 GB |
| GQA | 8 | 32.0 GB |
| GQA | 4 | 16.0 GB |
| MQA | 1 | 4.0 GB |
興味深いのは、全32層をAttentionにしたままMQA（$H_{kv}=1$）を採用しても、JambaがAttention層の数を4層まで減らして達成したのと同じ4GBに到達する点だ。同じ「KVキャッシュを32分の1にする」という結果でも、削る対象が層の数かヘッドの数かという、まったく違うレバーになっている。両方を同時に使えば、さらに掛け算で削れる。
ここまでのまとめ
キャッシュがない実装は、生成が進むほど無駄な再計算が積み上がり、実測では最大6倍の速度差になった。キャッシュを使う前提で次に効くのがヘッド数で、MHAからMQAにするとKVキャッシュは最大32分の1（この記事の条件では128GB→4GB）まで縮む。GQAはMHAとMQAの中間を選べる設定で、品質と省メモリのバランスを取りたいときに使われる。
3. 溜まったキャッシュの「持ち方」を変える ── PagedAttention
前節で、キャッシュそのものを小さくする方法（MQA・GQA）を見た。だが複数のリクエストを同時にさばくサービングの現場では、キャッシュを小さくするだけでは足りない問題がある。このセクションで分かること：KVキャッシュの「中身」ではなく「置き方」を変えることで、スループットが2〜4倍になった仕組み。
初期のサービングシステムは、リクエストごとに「起こりうる最大文脈長」ぶんの連続したメモリ領域をあらかじめ確保していた。実際の生成が短く終わっても、確保した分は他のリクエストに使い回せず、断片化やムダな予約によってGPUメモリが圧迫される。
2023年、KwonらはOSの仮想メモリ・ページングの発想を借りて、KVキャッシュを固定サイズの「ブロック」単位に分割し、物理的に連続していないメモリ領域にも配置できるようにした。これがPagedAttentionで、これを実装したサービングエンジンがvLLMだ。
論文の実測では、vLLMは既存システム（FasterTransformer、Orca）と比べて同程度のレイテンシで2〜4倍のスループットを達成している。しかもこの手法はKVキャッシュ自体のサイズを削るわけではない。MQA・GQAが「原本の数を減らす」対策なら、PagedAttentionは「持っている原本を無駄なく棚に並べる」対策で、両者は競合せず組み合わせられる。実際、GQAを採用したモデルをvLLMのようなPagedAttention対応エンジンで配信するのが、現在の実務ではほぼ標準的な組み合わせになっている。
| 状況 | 優先すべき対策 |
|---|---|
| 自分でモデルを選べる・設計できる、推論コストを根本的に下げたい | GQA・MQAを採用したモデルを選ぶ（アーキテクチャ側） |
| 既存モデルで複数リクエストを同時にさばきたい、スループットを上げたい | PagedAttention対応のサービングエンジンを使う（サービング側） |
| 両方の制約がある | 両方を組み合わせる（GQA採用モデル＋PagedAttention対応エンジン） |
トラブルシューティング
| 症状 | 原因 | 対処 |
|---|---|---|
| バッチサイズを上げるとすぐOOMになる | KVキャッシュが文脈長×バッチサイズに比例して線形に増える | GQA・MQA採用モデルへの切り替え、またはPagedAttention対応エンジンの利用を検討する |
| 長文脈タスクで推論が急に遅くなる | KVキャッシュのメモリ帯域読み込みがボトルネックになっている | $H_{kv}$を減らせるモデルか確認する。減らせないなら量子化やページング対応のエンジンを検討する |
| 同じシステムプロンプトを使う複数リクエストで無駄が多い | 各リクエストが同じ内容のKVキャッシュを個別に保持している | PagedAttention対応エンジンのprefix caching（共有プレフィックスのキャッシュ共有）機能を確認する |
| Cross-AttentionのKVキャッシュサイズを計算に入れ忘れる | Self-Attention側のキャッシュだけを見積もっている | Encoder-Decoderモデルでは、Encoder出力ぶんのKey・Valueも別途VRAMを消費することを忘れない |
用語集
- KVキャッシュ｜Attention層が過去に計算したKey・Valueベクトルを保持しておく記憶領域。新しいトークンぶんだけ追記すれば、過去分の再計算を省ける
- MHA（Multi-Head Attention）｜Query・Key・Valueそれぞれに同じ数のヘッドを持つ通常のAttention
- MQA（Multi-Query Attention）｜全ヘッドでKey・Valueを1組だけ共有する方式。KVキャッシュを最も削減できるが品質への影響もある
- GQA（Grouped-Query Attention）｜ヘッドをいくつかのグループに分け、グループ内でKey・Valueを共有する方式。MHAとMQAの中間
- PagedAttention｜KVキャッシュをOSのページングのようにブロック単位で非連続に管理する手法。キャッシュのサイズではなく置き方を最適化する
- スループット｜単位時間あたりに処理できるリクエスト数やトークン数。PagedAttentionはこの数値を主に改善する
学習ロードマップ
- Self-Attention：Query・Key・Valueの基本計算を理解する
- Attention（別記事）：Self-AttentionとCross-Attentionでキャッシュの増え方がどう違うかを理解する
- Jamba：Attention層の「数」を減らすことでKVキャッシュを削る実例を見る
- vLLM公式ドキュメント：PagedAttentionの実装を実際に動かしてみる
まとめ
KVキャッシュを「あって当然のもの」として使っているうちは、その正体を意識する機会がなかった。だが実際にキャッシュなしの実装を書いて速度を測ってみると、無駄な再計算が生成の長さとともに二乗のペースで積み上がっていく様子が数字としてはっきり見えた。そしてキャッシュを使う前提に立ったとき、次に効くレバーが「ヘッドの数を減らす」ことと「持っているキャッシュの置き方を変える」ことという、性質の異なる2つの対策だと分かった。どちらか一方が正解ということはなく、Jambaのように層の数を削る、GQAのようにヘッドの数を削る、PagedAttentionのように置き方を変える──これらは互いに競合しない、掛け算できる対策なのだと理解できたのが、この記事を書いていちばんの収穫だった。
参考文献
- Shazeer, N. "Fast Transformer Decoding: One Write-Head is All You Need" arXiv:1911.02150（2019）（邦題：高速なTransformerデコーディング ── 書き込みヘッドは1つで十分）── Multi-Query Attentionの初出。 https://arxiv.org/abs/1911.02150
- Ainslie, J. et al. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" arXiv:2305.13245（2023）（邦題：GQA ── Multi-Headチェックポイントから汎用Multi-Queryモデルを学習する）── Grouped-Query Attentionの初出。 https://arxiv.org/abs/2305.13245
- Kwon, W. et al. "Efficient Memory Management for Large Language Model Serving with PagedAttention" arXiv:2309.06180（2023、SOSP）（邦題：PagedAttentionによる大規模言語モデルサービングの効率的メモリ管理）── PagedAttentionとvLLMの出典、2〜4倍のスループット改善の実測値。 https://arxiv.org/abs/2309.06180
- Vaswani, A. et al. "Attention Is All You Need" arXiv:1706.03762（2017）（邦題：Attentionこそが必要なすべて）── Self-Attentionの基礎計算の出典。 https://arxiv.org/abs/1706.03762