KV cache / batching / paged attention

生成は2つのフェーズでできている

• prefill: プロンプト全体を一度に読み込み、最初の1トークンを出すまで。数千トークンを並列に処理できるので、GPUの計算能力をフルに使えます。ここにかかる時間が TTFT（Time To First Token、最初の一語までの時間）です。
• decode: 1トークンずつ続きを書く段階。毎回1トークン分の計算しかしないのに、モデルの重み全体をメモリから読み直す必要があるので、計算器はほとんど暇をしています。1トークンごとの間隔は TPOT や ITL（inter-token latency）と呼ばれます。

この章の工夫の多くは「decodeの暇な時間をどう埋めるか」と「prefillとdecodeをどう混ぜるか」の2つに集約されます。

KVキャッシュ：毎回読み直さないためのメモ帳

Transformerの注意機構（attention）では、新しいトークンを書くたびに、それまでの全トークンを振り返ります。そのとき各トークンから作られるのがKey（K）とValue（V）というベクトルです。Kは「このトークンは何についての情報か」という見出し、Vは「中身」にあたります。

過去のトークンのKとVは、あとから変わることがありません。だから一度計算したら保存しておき、次のトークンでは新しい1トークン分だけを計算して追記します。これがKVキャッシュで、ほぼすべての推論エンジンが使っています。

flowchart TD
  A["プロンプト"] --> B["prefill
全トークンのKとVを計算"]
  B --> C["KVキャッシュに保存"]
  C --> D["decode
新しい1トークンだけ計算"]
  D --> E["KとVを追記"]
  E --> C
  D --> F["次のトークンを出力"]

KVキャッシュはどれくらいメモリを食う？

KVキャッシュの大きさは、次の掛け算で決まります。

2（KとV）× 層数 × KVヘッド数 × ヘッド次元 × トークン数 × 同時リクエスト数 × 1要素のバイト数

Llama 3.1 8B（32層、KVヘッド8個、ヘッド次元128）を16bitで持つと、1トークンあたり 2 × 32 × 8 × 128 × 2バイト = 128KiB。たいしたことはなさそうですが、トークン数と人数を掛けると一気に膨らみます。

ウィジェットの初期状態（Llama 3.1 8B、16Kトークン、1リクエスト、FP16）では約2GiBと表示されます。文脈のスライダーを128Kまで動かすと約16GiB。8Bモデルの重み（16bitで約16GB）と同じくらいの大きさが、たった1人の会話のために必要になるわけです。

次に、同時リクエスト数を8に増やしてみてください。16Kトークンのままでも約16GiBになります。プリセットを Llama 3.1 70B（80層）に切り替えると、1トークンあたり320KiBになり、128K文脈1人分で約40GiBです。サーバーで「何人を同時に捌けるか」は、重みではなくKVキャッシュで決まることが多いのです。キャッシュの精度を「FP8 / INT8」にすると、どの場合もちょうど半分になります。これが第10章で触れたKVキャッシュの量子化です。

KVを小さくするモデル側の工夫：MQA・GQA・MLA

KVキャッシュを減らすいちばん根本的な方法は、モデルの設計そのものを変えることです。

ウィジェットのプリセット「比較用: GQAなし」を選ぶと、同じ形でKVヘッドが32個の場合が計算できます。1トークンあたり512KiBで、GQAの4倍です。GQAの論文は、MQAほど品質を落とさずに、MQAに近い速さが得られると報告しています。

MLAはさらに大胆です。DeepSeek-V2 の論文では、従来の DeepSeek 67B と比べてKVキャッシュを93.3%削減したと述べています。DeepSeek-V3（総パラメータ671B）の設定値（61層、潜在次元512、位置情報用64次元）から計算すると、16bitで1トークンあたり約69KiB。8Bの Llama 3.1 の半分ほどで済む計算です。MoEの仕組みについては第26章で扱います。

バッチング：暇なGPUにまとめて仕事を渡す

decodeのGPUは暇をしている、という話を思い出してください。重みを1回読み出すなら、そのついでに何人分ものトークンを同時に計算したほうがお得です。これがバッチング（batching）です。

素朴なバッチング（static batching）では、何件かのリクエストをまとめて開始し、全員が書き終わるまで次を入れません。すると、短い返事で終わった人の席が空いたまま、長文を書く人を待つことになります。

continuous batching（連続バッチング、iteration-level scheduling とも）は、1トークン生成するたびに「終わった人を降ろし、待っている人を乗せる」方式です。2022年のOrca（OSDI '22）が提案し、今では vLLM、SGLang、TensorRT-LLM などの標準機能になっています。llama.cpp のサーバーでも既定で有効です。

PagedAttention：メモリを「ページ」で貸し出す

continuous batching で人の出入りが激しくなると、今度はKVキャッシュの置き場所が問題になります。以前のシステムは、リクエストごとに「最大の長さ分」の連続した領域を先に確保していました。実際には短く終わる会話が多いので、大半が使われないまま無駄になります。

vLLMの論文（2023年）は、既存システムではKVキャッシュ用メモリのうち実際のトークンの保存に使われていたのは20.4〜38.2%にすぎないと測定しました。そこで提案されたのがPagedAttentionです。

KVキャッシュを一定のトークン数ごとのブロックに分け、必要になるたびに空いているブロックを割り当てます。ブロックは物理的に連続していなくてよく、注意機構の計算はブロックテーブルを見ながら行います。論文の測定では、vLLMはメモリの96.3%を実際のトークンの保存に使えていました。

ブロック単位なので、同じプロンプトから複数の回答を作るときは共通部分のブロックを共有し、書き換えが必要になった時点でだけコピーする（copy-on-write）こともできます。こうした工夫で、vLLMは同程度の遅延のまま、従来システムより2〜4倍のスループットを出したと報告されています。vLLMそのものは第9章で紹介しています。

プレフィックスキャッシュとchunked prefill

プレフィックスキャッシュ（prefix caching）は、プロンプトの先頭が同じリクエスト同士でKVキャッシュを使い回す仕組みです。長いシステムプロンプト、毎回同じ資料、チャットの過去のやり取りなどは、前回作ったKVキャッシュをそのまま使えばprefillを省けます。vLLMは Automatic Prefix Caching としてブロックのハッシュで一致を探し、SGLang は RadixAttention という木構造で管理します。商用APIの「プロンプトキャッシュ」割引も、考え方は同じです。

chunked prefillは、長いプロンプトのprefillを小さな塊（チャンク）に分け、他の人のdecodeと同じバッチに混ぜて処理する方法です。巨大なprefillが割り込むと、生成中の人の文字の流れが止まってしまいます。分割して少しずつ混ぜれば、計算が重いprefillとメモリが重いdecodeがちょうどよく補い合います。Sarathi-Serve（2024年）がこの方式を提案し、vLLMのV1エンジンでは可能な場合に既定で有効になっています。

FlashAttention：注意計算そのものを速く

最後に、計算の中身の話です。注意機構の素直な実装は、トークン数の2乗の大きさの中間結果をGPUの大きなメモリ（HBM）に書いては読み直すため、メモリの往復で時間を食います。FlashAttention（2022年）は、計算を小さなタイルに分けてGPUチップ上の高速なSRAMの中で済ませ、HBMとの往復を減らしました。近似ではなく、結果は通常の注意計算と同じです。

FlashAttentionはKVキャッシュの量を減らすものではなく、同じ計算をより速く行う技術です。PagedAttentionやGQAとは役割が違い、組み合わせて使われます。