第15章. KV cache / batching / paged attention

この章の目的

この章では、LLMの推論速度を向上させるための主要な技術であるKV cache、batching、そしてpaged attentionについて、その基本的な概念と仕組みを理解することを目的とします。これらの技術は、LLMを実用的な速度で動作させる上で不可欠であり、推論ランタイム（第9章参照）の性能を大きく左右します。

この章で覚えるべきこと

KV cacheがLLMの推論においてどのようにメモリと計算を節約するか
batchingがスループット向上にどのように貢献するか
paged attentionがKV cacheのメモリ効率をどのように改善するか
これらの技術が互いにどのように関連し、推論性能を最適化するか

導入

大規模言語モデル（LLM）は、その強力な生成能力の代償として、膨大な計算リソースとメモリを必要とします。特に、ユーザーからのプロンプトを受け取り、それに対する応答を生成する「推論」の段階では、いかに高速かつ効率的に処理を行うかが実用性向上の鍵となります。

これまでの章では、モデルのサイズを小さくする量子化（第10章）や、特定のハードウェアに最適化された推論ランタイム（第9章）について学びました。本章で扱うKV cache、batching、paged attentionは、これらとは異なるアプローチで推論のボトルネックを解消し、LLMの応答速度と処理能力を飛躍的に向上させるための技術です。これらを理解することで、特定の推論ランタイムが高速である理由や、どのような状況でどの技術が有効なのかを判断できるようになります。

基本概念

KV cache (Key-Value Cache)

ひとことで言うと:

LLMが過去に計算したAttentionの「Key」と「Value」を保存しておき、再利用することで計算量を減らし、推論を高速化する技術です。

何のカテゴリか:

推論最適化技術、メモリ最適化技術

何に使うのか:

LLMのテキスト生成（デコーディング）フェーズにおいて、Attention機構の計算負荷を軽減し、推論速度を向上させます。

代表例:

Transformerベースの全てのLLM推論で利用されます。vLLMやTGIなどの高速推論ランタイムでは、KV cacheの効率的な管理が性能の要となります。

よく混同される用語:

Attention機構そのもの: KV cacheはAttention機構の一部であるKeyとValueをキャッシュするものであり、Attention機構そのものではありません。
通常のCPU/GPUキャッシュ: 一般的なハードウェアキャッシュとは異なり、LLMの特定の計算結果をソフトウェア的に管理するものです。

初心者向け注意点:

KV cacheはメモリを消費します。特に長いテキストを生成する場合、大量のKeyとValueがキャッシュされるため、GPUメモリの消費量が大きくなる点に注意が必要です。

詳細説明:

TransformerモデルのAttention機構（第1章参照）は、入力シーケンスの各トークンに対して、他の全てのトークンとの関連度を計算します。この計算には、Query (Q)、Key (K)、Value (V) という3つのベクトルが使われます。

テキスト生成（デコーディング）の際、LLMは一度に1トークンずつ生成していくのが一般的です。例えば、「こんにちは、元気ですか？」という応答を生成する場合、

「こんにちは」を生成する。
「こんにちは」を考慮して「、」を生成する。
「こんにちは、」を考慮して「元気」を生成する。 ... というように、生成済みのトークン列全体を考慮して次のトークンを生成します。

ここで問題となるのが、各ステップでAttention計算を行う際に、既に計算済みのトークンに対するKeyとValueの計算を毎回繰り返してしまうことです。これは非常に非効率です。

KV cacheは、この無駄を省くために導入されました。一度計算されたKeyとValueのベクトルをGPUメモリ上に保存（キャッシュ）しておき、次のトークンを生成する際には、新しく入力されたトークン（Query）と、キャッシュされた過去のKey/Value、そして新しいトークン自身のKey/Valueを使ってAttention計算を行います。これにより、過去のトークンに対するKey/Valueの再計算が不要となり、計算量が大幅に削減されます。

KV cacheがない場合: 各生成ステップで、入力シーケンス全体のQ, K, Vを計算し、Attentionを適用。シーケンス長が長くなるほど計算量が二次関数的に増加。

KV cacheがある場合: 各生成ステップで、新しいトークンのQ, K, Vのみを計算。過去のK, Vはキャッシュから取得。Attentionは新しいQと、キャッシュされたK/V + 新しいK/Vに対して適用。計算量が大幅に削減される。

graph TD
    A[プロンプト入力] --> B{トークン1生成}
    B -- Q1, K1, V1 --> C[Attention計算]
    C -- K1, V1をキャッシュ --> D{トークン2生成}
    D -- Q2, K2, V2 --> E["Attention計算 (Q2とK1,V1,K2,V2)"]
    E -- K2, V2をキャッシュ --> F{トークン3生成}
    F -- Q3, K3, V3 --> G["Attention計算 (Q3とK1,V1,K2,V2,K3,V3)"]
    G -- K3, V3をキャッシュ --> H[...]

図15.1: KV cacheの動作原理

KV cacheの導入による計算量の削減は、特にシーケンス長 $L$ が長くなるにつれて顕著になります。Attention計算の計算量は通常 $O(L^2)$ ですが、KV cacheを使用することで、各生成ステップでの計算量は $O(L)$ に削減されます。

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ 式15.1: Attention機構の基本式。ここで $Q, K, V$ はそれぞれQuery, Key, Value行列、$d_k$ はKeyベクトルの次元です。KV cacheは、この $K$ と $V$ を過去のステップから再利用することで、計算負荷を軽減します。

Batching (バッチ処理)

ひとことで言うと:

複数の独立したリクエスト（プロンプト）をまとめて同時に処理することで、GPUの利用効率を高め、全体のスループットを向上させる技術です。

何のカテゴリか:

推論最適化技術、スループット向上技術

何に使うのか:

複数のユーザーからのリクエストや、大量のテキストを並行して処理する必要がある場合に、LLMの推論サーバーの処理能力を最大化します。

代表例:

Web APIとしてLLMを提供するサービス（OpenAI APIなど）や、vLLM、TGIといった推論ランタイムで広く利用されています。

よく混同される用語:

KV cache: KV cacheは単一のリクエスト内の計算を最適化するのに対し、batchingは複数のリクエストをまとめて処理します。両者は併用されます。
並列処理: batchingは一種の並列処理ですが、特に「複数の独立した入力データをまとめて一つの大きな入力として処理する」という点で特徴的です。

初心者向け注意点:

batchingはスループットを向上させますが、個々のリクエストのレイテンシ（応答時間）は、バッチサイズが大きくなるとわずかに増加する可能性があります。これは、バッチ内の最も長いシーケンスに合わせて処理時間が決まるためです。

詳細説明:

GPUは、大量の並列計算を得意とするハードウェアです。しかし、単一の小さな計算タスクだけでは、GPUの計算リソースを十分に活用しきれないことがあります。LLMの推論においても、1つのプロンプトだけを処理する場合、GPUのコアの多くがアイドル状態になることがあります。

Batchingは、この問題を解決するために、複数のユーザーからのプロンプトや、複数のテキスト生成タスクを一つの「バッチ」としてまとめ、GPUに一度に投入する手法です。これにより、GPUはより多くの計算を並行して実行できるようになり、全体としての処理能力（スループット）が向上します。

例えば、3人のユーザーが同時にLLMに質問を投げたとします。

Batchingなしの場合: ユーザーAの質問を処理 → ユーザーBの質問を処理 → ユーザーCの質問を処理（各処理の間でGPUがアイドル状態になる可能性あり）
Batchingありの場合: ユーザーA, B, Cの質問をまとめて一つのバッチとしてGPUに投入 → 3つの質問を同時に処理（GPUの利用率が向上し、全体として早く処理が完了）

graph TD
    A[ユーザー1リクエスト] --> B{バッチキュー}
    C[ユーザー2リクエスト] --> B
    D[ユーザー3リクエスト] --> B
    B -- バッチサイズNに達したら --> E[LLM推論エンジン]
    E -- 並列処理 --> F[ユーザー1応答]
    E -- 並列処理 --> G[ユーザー2応答]
    E -- 並列処理 --> H[ユーザー3応答]

図15.2: Batchingの動作原理

Batchingには、静的バッチング（固定サイズのバッチで処理）と動的バッチング（リアルタイムで到着するリクエストに応じてバッチサイズを調整）があります。特にLLMの推論では、リクエストごとにプロンプト長や生成する応答長が異なるため、動的バッチングが非常に重要になります。

Paged Attention (ページド・アテンション)

ひとことで言うと:

KV cacheのメモリ管理を、オペレーティングシステムの仮想メモリ管理（ページング）に似た方法で行うことで、メモリの断片化を解消し、より効率的にKV cacheを利用する技術です。

何のカテゴリか:

推論最適化技術、メモリ管理技術、KV cache最適化技術

何に使うのか:

特に動的バッチングと組み合わせることで、可変長のシーケンスを持つ複数のリクエストを効率的に処理し、GPUメモリを最大限に活用してスループットを向上させます。

代表例:

vLLMがこの技術を初めて導入し、その高いスループット性能の主要因となっています。

よく混同される用語:

KV cache: Paged attentionはKV cacheのメモリ管理方法であり、KV cacheそのものではありません。KV cacheをより賢く使うための技術です。
仮想メモリ/ページング: オペレーティングシステムの概念をLLMのKV cache管理に応用したものであり、OSの仮想メモリそのものではありません。

初心者向け注意点:

Paged attentionは、KV cacheのメモリ効率を劇的に改善しますが、その恩恵は主に複数のリクエストを同時に処理する（バッチ処理する）場合に顕著です。単一のリクエストを処理するだけでは、その効果は限定的です。

詳細説明:

前述の通り、KV cacheはLLMの推論を高速化しますが、メモリを大量に消費します。特に、複数のリクエストをバッチ処理する場合、各リクエストのプロンプト長や生成する応答長はバラバラです。

従来のKV cache管理では、各リクエストに対してKV cache用の連続したメモリ領域を事前に確保していました。この方法にはいくつかの問題があります。

メモリの断片化: リクエストの完了やキャンセルによって解放されたメモリ領域が、次に到着するリクエストのサイズに合わない場合、その領域が再利用されず、メモリが断片化して無駄が生じます。
メモリの過剰確保: 最大のシーケンス長に合わせてメモリを確保する必要があるため、短いシーケンスのリクエストでは多くのメモリが無駄になります。
連続メモリの要件: GPUメモリは連続した領域を確保するのが難しく、特に大規模なバッチ処理では制約となります。

Paged attentionは、これらの問題を解決するために、オペレーティングシステムの仮想メモリ管理（ページング）の概念をKV cacheに応用します。

KV cacheのメモリを固定サイズの「ブロック（ページ）」に分割します。
各リクエストのKV cacheは、連続したメモリ領域ではなく、これらのブロックを不連続に割り当てて構成されます。
物理メモリ上のブロックと、論理的なKV cacheシーケンスの対応付けは、「ブロックテーブル」と呼ばれるマッピングによって管理されます。

graph TD
    subgraph "リクエスト1のKV cache"
        R1_B1[ブロック1]
        R1_B2[ブロック2]
        R1_B3[ブロック3]
    end

    subgraph "リクエスト2のKV cache"
        R2_B1[ブロック4]
        R2_B2[ブロック5]
    end

    subgraph "GPU物理メモリ"
        P_B1[物理ブロックA]
        P_B2[物理ブロックB]
        P_B3[物理ブロックC]
        P_B4[物理ブロックD]
        P_B5[物理ブロックE]
    end

    R1_B1 --> P_B1
    R1_B2 --> P_B3
    R1_B3 --> P_B5

    R2_B1 --> P_B2
    R2_B2 --> P_B4

    subgraph "ブロックテーブル (マッピング)"
        T1[リクエスト1: [A
        C
        E]]
        T2[リクエスト2: [B
        D]]
    end

    R1_B1 -- 論理 --> T1
    R2_B1 -- 論理 --> T2

図15.3: Paged Attentionの概念図

この方式により、以下のメリットが得られます。

メモリの断片化解消: 空いているブロックを柔軟に割り当てられるため、メモリの断片化が起こりにくくなります。
メモリの有効活用: 必要な分だけブロックを割り当てられるため、メモリの過剰確保が減り、GPUメモリをより多くのリクエストで共有できます。
Copy-on-Writeの実現: 複数のリクエストが同じプロンプトを共有する場合（例: 共通のシステムプロンプト）、その部分のKV cacheブロックを共有し、書き込みが必要になったときに初めてコピーを作成する「Copy-on-Write」が可能になります。これにより、さらにメモリ効率が向上します。

Paged attentionは、特に動的バッチングと組み合わせることで、GPUメモリを最大限に活用し、LLM推論のスループットを大幅に向上させる画期的な技術です。

KV cacheのメモリ消費量 $M_{KV}$ は、通常、以下の式で近似できます。

$$ M_{KV} = 2 \times N_{layers} \times N_{heads} \times d_{head} \times L \times \text{dtype_size} $$ 式15.2: KV cacheのメモリ消費量。$N_{layers}$ はTransformer層の数、$N_{heads}$ はAttentionヘッドの数、$d_{head}$ は各ヘッドの次元、$L$ はシーケンス長、$\text{dtype_size}$ はデータ型（例: fp16なら2バイト）のサイズです。Paged Attentionは、この $L$ の部分を動的に管理し、メモリの無駄を削減します。

具体例

KV cacheの例

LLMが「日本の首都は」というプロンプトに対して「東京」と生成するプロセスを考えます。

プロンプト処理フェーズ: 「日本」「の」「首都」「は」の各トークンが入力されます。
- 各トークンに対してQ, K, Vベクトルが計算され、Attentionが適用されます。
- このとき計算されたK, VベクトルはKV cacheに保存されます。
- 最初の出力トークンとして「東」が生成されます。
生成フェーズ (1トークン目):
- 新しく生成された「東」のQ, K, Vベクトルが計算されます。
- Attention計算では、「東」のQベクトルと、KV cacheに保存されている「日本」「の」「首都」「は」のK, Vベクトル、そして「東」自身のK, Vベクトルが使われます。
- 計算後、「東」のK, VベクトルもKV cacheに追加されます。
- 次の出力トークンとして「京」が生成されます。
生成フェーズ (2トークン目):
- 新しく生成された「京」のQ, K, Vベクトルが計算されます。
- Attention計算では、「京」のQベクトルと、KV cacheに保存されている「日本」「の」「首都」「は」「東」のK, Vベクトル、そして「京」自身のK, Vベクトルが使われます。
- 計算後、「京」のK, VベクトルもKV cacheに追加されます。
- 応答が完了します。

このように、KV cacheがあることで、生成フェーズにおいて過去のトークンに対するK, Vベクトルの再計算が不要になり、計算量が大幅に削減されます。

Batchingの例

WebサービスでLLM APIを提供している状況を想定します。

ユーザーA: 「今日の天気は？」
ユーザーB: 「おすすめのレストランを教えて」
ユーザーC: 「PythonでFizzBuzzを書いて」

これらのリクエストがほぼ同時に到着した場合、推論ランタイムは以下のように処理します。

リクエストのキューイング: 各リクエストは推論ランタイムのキューに格納されます。
バッチの形成: キューから複数のリクエストが選択され、一つのバッチとしてまとめられます。例えば、上記3つのリクエストが1つのバッチになります。
パディング: 各リクエストのプロンプト長が異なるため、バッチ内の最長プロンプトに合わせて短いプロンプトにはパディングトークンが追加され、同じ長さになります。
並列推論: 形成されたバッチがLLMに一度に投入され、GPU上で並列に処理されます。
結果の分離: LLMからの出力は、元のリクエストごとに分離され、各ユーザーに返されます。

これにより、GPUは常に高い稼働率を維持し、個々のリクエストを順番に処理するよりもはるかに高い全体スループットを実現します。

Paged Attentionの例

Batchingの例で、さらに生成される応答の長さもバラバラだとします。

ユーザーA: 「今日の天気は？」 → 「晴れです。」（短い）
ユーザーB: 「おすすめのレストランを教えて」 → 「新宿にあるイタリアンレストラン『ラ・パスタ』がおすすめです。特にランチのコースは...」（中くらい）
ユーザーC: 「PythonでFizzBuzzを書いて」 → 「はい、承知いたしました。以下にPythonでFizzBuzzを実装したコードを示します。python ... 」（長い）

従来のKV cache管理では、各リクエストに対して最大生成長を考慮した連続メモリを確保する必要があり、短い応答のリクエストではメモリが無駄になります。

Paged attentionでは、

各リクエストのKV cacheは、固定サイズのブロック（例: 256トークン分のK, Vベクトルを格納できるブロック）に分割されます。
リクエストAには1ブロック、リクエストBには3ブロック、リクエストCには5ブロックといったように、必要な分だけ不連続な物理メモリブロックが割り当てられます。
これらのブロックの物理的な位置は、各リクエストの「ブロックテーブル」に記録されます。
リクエストAが完了し、そのKV cacheブロックが解放されると、そのブロックはすぐに他のリクエストに再利用可能になります。メモリの断片化が起こりにくく、GPUメモリをより効率的に共有できます。

これにより、同じGPUメモリ量で、より多くのリクエストを同時に処理したり、より長い応答を生成したりすることが可能になります。

よく混同される用語との比較

特徴	KV cache	Batching	Paged Attention
目的	単一リクエスト内のAttention計算を高速化	複数リクエストのスループットを向上	KV cacheのメモリ効率を向上
対象	TransformerのAttention計算	複数の独立した推論リクエスト	KV cacheのメモリ管理
仕組み	過去のK, Vベクトルをメモリに保存し再利用	複数のリクエストをまとめて並列処理	KV cacheを固定サイズのブロックで管理
効果	推論レイテンシの削減、計算量削減	全体スループットの向上、GPU利用率向上	メモリ断片化の解消、メモリ利用効率向上
単独効果	単一リクエストでも効果あり	複数リクエストがないと効果なし	主に複数リクエストで効果あり
関連性	BatchingやPaged Attentionと併用される	KV cacheやPaged Attentionと併用される	KV cacheの最適化であり、Batchingと相乗効果

表15.1: KV cache, Batching, Paged Attentionの比較

graph TD
    subgraph "推論最適化技術"
            A
            KV
            Cache
            C
            Batching
            E
            Paged
            Attention
    end

    A -- 依存 --> G[Transformer推論]
    C -- 適用 --> G
    E -- KV Cacheの管理 --> A

    D -- 恩恵を受ける --> F
    F -- 恩恵を受ける --> D

    G -- 最終目標 --> H[高速・効率的なLLM推論]

図15.4: LLM推論最適化技術の関係性

実務での位置づけ

これらの技術は、LLMを実用的なアプリケーションに組み込む上で不可欠な要素です。

KV cache: ほとんど全てのLLM推論ランタイムで標準的に実装されており、単一のプロンプトに対する応答速度（レイテンシ）に直接影響します。これがなければ、LLMの生成速度は非常に遅くなります。
Batching: LLMをAPIサービスとして提供したり、大量のテキストを処理するバックエンドシステムで利用する際に、サーバーのスループットを最大化するために必須です。これにより、より多くのユーザーを同時に捌けるようになります。
Paged Attention: 特にvLLMのような高性能な推論ランタイムの核となる技術であり、動的バッチングと組み合わせることで、限られたGPUメモリで最大限のスループットを引き出すことを可能にします。これにより、より安価なGPUでより多くのリクエストを処理できるようになり、運用コストの削減にも貢献します。

これらの技術は、それぞれが独立して機能するだけでなく、互いに組み合わさることで相乗効果を発揮します。例えば、vLLMはKV cacheをPaged Attentionで効率的に管理しながら、動的バッチングによって高いスループットを実現しています。推論ランタイム（第9章参照）を選ぶ際には、これらの技術がどのように実装され、最適化されているかを理解することが重要です。

まとめ

3行まとめ

KV cacheは、LLMが過去に計算したAttentionのKeyとValueを保存し再利用することで、推論時の計算量を減らし、応答速度を向上させます。
Batchingは、複数のリクエストをまとめてGPUに投入することで、GPUの利用効率を高め、LLMサーバーの全体スループットを向上させます。
Paged Attentionは、KV cacheのメモリ管理をOSのページングのように行うことで、メモリの断片化を解消し、特にバッチ処理時のKV cacheのメモリ効率を劇的に改善します。

混同しやすい用語

KV cacheとAttention機構そのもの
Batchingと単なる並列処理
Paged AttentionとKV cacheそのもの、またはOSの仮想メモリ

次に読むべき章

本章で推論高速化の基礎的な概念を理解しました。これらの技術がどのように実装されているか、より具体的な推論ランタイムの内部構造に興味がある場合は、第9章「vLLM / TGI / TensorRT-LLM」を再読すると、より深く理解できるでしょう。また、モデルのサイズを小さくする量子化（第10章）と組み合わせることで、さらに効率的な推論が可能になります。