第8章 — 次世代のKVキャッシュ管理

LLM Primer VI: AIシステムのスケーリング を章ごとに紹介していくウォークスルー、第8回。オペレーティングシステムのページングの洞察を推論エンジンに持ち込む章 — そしてKVキャッシュを、予約されたバイトのスラブから、共有可能でエビクション可能でプレフィックスキャッシュ可能なリソースへ変える。

なぜこの章があるのか

第7章は連続バッチングに借金を残しました。系列は反復ごとに出入りし、ナイーブなレイアウトは各スロットに最大サイズのスラブを与え、その大半が無駄になり、バッチングのスループット勝利は部分的に返される。借金は内部断片化 — まさに1960年代にオペレーティングシステムがページングで解決した失敗モード。第8章はその解をLLMサービングに適用する話です。KVキャッシュを小さな物理ブロックに分割し、それらをページテーブルで論理トークン位置と分離し、エビクションとキャッシュポリシーで系列間でブロックを回収または共有する。PagedAttentionが基礎的な動きで、H2OとInfiniGenがエビクションポリシー、プレフィックスキャッシュこそが本番クラスタが十数台のGPUで数百万のエージェント的リクエストを捌けるようにする技法です。

ひとことで言うと: KVキャッシュをページ化仮想メモリのように扱う — 小さなブロック、ページテーブル、参照カウント共有 — そうすれば連続バッチングのメモリ問題は問題でなくなる。

8.1 PagedAttentionはKVキャッシュのための仮想メモリ

vLLMのPagedAttention論文(2023)は、OSの洞察をエンジンに直接持ち上げます。KVキャッシュは固定サイズのブロック — 通常16トークン — に分割され、平坦な物理プールに保持される。系列はブロックテーブルで表現される。論理位置を物理ブロックIDにマップするポインタの配列です。アテンションカーネルはブロックテーブルを追加入力として取り、連続したスライスではなく間接参照でキーとバリューを収集する。HopperではL2がランダムアクセスパターンを十分によく吸収し、カーネルはスラブベース版の数パーセント以内で走ります。勝利は大きい。内部断片化は60〜80パーセントから約4パーセント(系列あたり部分末尾ブロック1つ)に落ち、利用可能な同時実行数を2〜4倍に上げます。参照カウント付きのブロック共有は複雑なサンプリングをほぼ無料にし — 2,000トークンプロンプトのbest-of-8が16,800から2,800トークンブロックに落ち — プレフィックスキャッシュが上に構築される基盤です。

8.2 H2OとInfiniGenは重要でないトークンをエビクトする

PagedAttentionは断片化を解きますが、どんなレイアウトでもVRAMを超過するコンテキストは解けません。200,000トークンのLlama-3-70Bコンテキストは重みと並んで60 GBのKVを必要とする。H2O(「Heavy Hitter Oracle」)は、デコード時のアテンション重みが少数の元位置 — 直近トークン、最初にあるアテンションシンクトークン、そして疎に散らばる内容関連のヒット — に集中し、大半の履歴位置は本質的にゼロ重みを受け取ることを観察します。エンジンは位置ごとの累積アテンションスコアを追跡し、系列のKV予算が上限に近づくと、最近と最初のトークンの保証保持ウィンドウを除いて最低スコアの位置をエビクトする。節約は大きい。コストは恒久性 — 後のクエリがエビクトされた位置を必要としても、エンジンは復元できない。InfiniGenは動的で復元可能な選択でトレードを洗練します。トークンをそのまま落とすのではなく、そのKVをCPUメモリにオフロードし、そこにアテンションが再び集中し始めたら戻す。適切なエビクションポリシーは、ワークロードが自身の長い履歴をどれだけ再クエリするかに依存します。エージェント的ワークロードは恒久エビクションを罰し、InfiniGen型の復元を報います。

8.3 プレフィックスキャッシュはPagedAttentionが解き放つ最大の効果レバー

実トラフィックでは、大半のプロンプトの最初の数千トークンはリクエスト間で同一です。チャットサービスは会話ごとに同じシステムプロンプトを再利用する。RAGサービスは同じ検索通過を検索した全ユーザーのプロンプトに貼り付ける。エージェントはステップごとに同じツール説明と推論スキャフォールドを注入する。PagedAttentionは共有を機構化します。プロンプトをブロックサイズのチャンクでハッシュ化する。ハッシュが計算済みブロックのグローバルキャッシュにあれば、新規リクエストのブロックテーブルをそのキャッシュされたブロックに指させ、そのプレフィックスのプリフィルを完全にスキップする。なければプリフィルが走り、結果のブロックが登録される。本番ヒット率は劇的です — チャットサービスのシステムプロンプトで99パーセント超、RAG検索依存のプレフィックスで30〜60パーセント、エージェント的スキャフォールドで1.0近く。SGLangのRadixAttentionは、ブロック境界に揃わない任意長の共有プレフィックスをキー付けする基数木でこの考え方をさらに押し進めます。プレフィックスキャッシュは、予算超過のサービングクラスタを予算内に変える単一の技法として最も頻繁に効きます。

覚えておきたいこと: ワークロードに共有システムプロンプト、共有ツールスキャフォールド、あるいは安定した検索コーパスがあるなら、プレフィックスキャッシュは投資対効果が最も高いスイッチである可能性が高い。エンジンで有効化されていると仮定する前にヒット率を計測する。その数字が、プリフィルのうちどれだけを二重に払っているかを教えてくれる。

第8章が敷いたもの

ページング、エビクション、プレフィックスキャッシュはトークンあたりのKVフットプリントを縮め、エンジンの高同時実行下のメモリ振る舞いを扱いやすくしました。いずれもデコードの根本的な逐次依存 — 系列あたり反復あたり1出力トークン、アクティブスロット数によらず — には対処しません。第9章はその制約に投機的デコードで挑みます。安価なドラフトで数トークン先を予測し、単一の高価な順伝播で推測を検証する技法群 — ユーザーが感じるレイテンシに最も重要な系列について、1ステップ1トークンの床を破ります。

次回 — 第9章: 投機的デコード。 ドラフト、検証、そして投機が割に合う条件の算術。

全体像を押さえたい方へ: 書籍版には動くページ化KVデータ構造コード、プレフィックスキャッシュ用のvLLMとSGLangの設定フラグ、H2Oエビクションのチューニングノブ、InfiniGenのオフロードと復元機構、そしてこの4つを本番KVスタックに合成するリファレンスアーキテクチャが収録されています。Amazonで『LLM Primer VI』を見る →