KVキャッシュとPagedAttention：なぜLLMサーバーは遅くなるのか

LLMサーバーの動作が遅くなっています。

4枚のA100 GPUに70Bモデルをデプロイしました。午前8時の時点ではすべて順調です。しかし、昼食時になるとレイテンシが2倍に膨れ上がります。メモリを確認すると、その大部分が「テンソルバッファ（tensor buffers）」によって占有されていました。これらは、実際には過去の会話からキャッシュされた状態なのです。

これがKVキャッシュの問題です。これは、本番環境でのLLMサービングにおける最大のボトルネックとなっています。

KVキャッシュとは何か？

すべてのTransformerモデルは、トークンを一つずつ生成します。新しいトークンを作成するには、モデルはそれ以前のすべてのトークンからKey（キー）とValue（バリュー）のテンソルを必要とします。これらを毎回再計算するのは時間がかかりすぎます。そのため、エンジンはこれらを保存します。この保存領域がKVキャッシュです。

メモリの問題：

Llama 3.1 70Bモデルの場合、単一の4096トークンのシーケンスに約1.3 GBのメモリが必要です。

もし同時に256人のユーザーがいる場合、336 GBのメモリが必要になります。これは4枚のA100 GPUの容量を超えています。KVキャッシュは非常に急速に増大するため、モデルの重みそのものよりも多くのメモリを消費することがよくあります。

従来型のメモリ管理が失敗する理由：

  • 内部断片化（Internal fragmentation）：4096トークン分のスペースを割り当てたものの、実際には300トークンしか使用していない場合、そのスペースの93%が無駄になります。
  • 共有不可：同じシステムプロンプトを持つ2人のユーザーが、それぞれプロンプトのコピーを個別に保存してしまいます。
  • 全か無かの退避（All-or-nothing eviction）：メモリが不足すると、シーケンス全体をCPUに移動させなければなりません。これにより、GPUの処理が停止します。

PagedAttentionによる解決策：

PagedAttentionはオペレーティングシステム（OS）のように機能します。KVキャッシュを「ページ（page）」と呼ばれる小さな固定サイズのブロックに分割します。

これにより、主に3つの問題が解決されます：

  • オンデマンド割り当て：シーケンスは成長するにつれてページを占有していきます。未使用の容量に対してメモリを無駄にすることはありません。
  • 共有プレフィックスのサポート：複数のユーザーが、共通のシステムプロンプトに対して同じ物理ページを共有できます。これは「コピーオンライト（copy-on-write）」のロジックを使用して、膨大な量のメモリを節約します。
  • きめ細かな退避（Fine-grained eviction）：メモリがいっぱいになったとき、システムは巨大なシーケンスではなく、小さなページをCPUに移動させます。

結果：

PagedAttention（vLLMに採用されている技術）を使用することで、従来の手法と比較してスループットを2倍から4倍に向上させることができます。

使用すべきケース：

  • 高並列処理（High concurrency）。
  • 長さの異なるシーケンス。
  • 同じ開始部分を持つプロンプト。

使用を避けるべきケース：

  • 単一ユーザーによるローカル推論。
  • 非常に小さなモデル。
  • すべてのシーケンスが全く同じ長さであるタスク。

KV CacheとPagedAttention：その仕組みと重要性

大規模言語モデル（LLM）の推論は、計算量とメモリ消費の両面で非常に重いタスクです。特に、モデルが長いテキストを生成しようとすると、リソースの消費は爆発的に増加します。

この課題を解決するために不可欠な2つの技術が、KV CacheとPagedAttentionです。この記事では、これらが何であるか、そしてなぜ重要なのかを解説します。

KV Cacheとは？

LLMの核となるのは「アテンション・メカニズム（Attention Mechanism）」です。アテンションは、現在のトークンを生成するために、それ以前のすべてのトークンとの関係性を計算します。

推論プロセス（特に自己回帰的な生成）では、新しいトークンが1つずつ生成されます。各ステップで、以前のトークンのKey（K）とValue（V）のベクトルを再計算するのは非常に非効率です。

KV Cacheは、一度計算した過去のトークンのKeyとValueのベクトルをメモリに保存しておく手法です。これにより、新しいトークンを生成する際には、過去の計算結果を再利用できるため、計算量を大幅に削減できます。

KV Cacheの課題

KV Cacheは計算を速めますが、代わりにメモリ（VRAM）を大量に消費します。

問題は、KV Cacheのメモリ割り当て方法にあります。従来のシステムでは、各リクエストに対して、あらかじめ「最大シーケンス長（Max Sequence Length）」に基づいた連続的なメモリ領域を確保していました。

これには2つの大きな問題があります：

  1. 内部断片化 (Internal Fragmentation): 実際には短い文章しか生成していない場合でも、最大長分のメモリが予約されているため、未使用のメモリが無駄になります。
  2. 外部断片化 (External Fragmentation): メモリが連続したブロックとして要求されるため、合計の空き容量は足りていても、大きな連続領域が確保できず、新しいリクエストを受け入れられないことがあります。

PagedAttention：メモリ管理の革命

このメモリ管理の非効率性を解決するために登場したのが、PagedAttentionです。

PagedAttentionは、オペレーティングシステム（OS）の**仮想メモリ（Virtual Memory）とページング（Paging）**の概念から着想を得ています。

仕組み

従来の方式では、KV Cacheをメモリ上の「連続した巨大なブロック」として扱っていました。一方、PagedAttentionは、KV Cacheを小さな**「ブロック（Blocks）」**に分割して管理します。

  • 非連続な割り当て: KV Cacheの各ブロックは、物理メモリ上のどこに配置されても構いません。必ずしも連続している必要はありません。
  • ページテーブル: OSが仮想アドレスを物理アドレスに変換するように、PagedAttentionは「ページテーブル」を使用して、論理的なシーケンス内のトークンがどの物理ブロックにあるかを追跡します。

PagedAttentionのメリット

  1. 断片化の劇的な削減: メモリを小さなブロック単位で動的に割り当てるため、内部断片化が最小限に抑えられます。
  2. メモリ利用率の向上: 必要な分だけメモリを確保できるため、同じGPUメモリでより多くのリクエストを同時に処理（バッチサイズを拡大）できます。
  3. コンティニュアス・バッチング (Continuous Batching) の実現: リクエストが終了した瞬間にそのメモリを解放し、新しいリクエストに即座に割り当てることが容易になります。これにより、スループットが大幅に向上します。

まとめ

  • KV Cacheは、計算を高速化するために過去の情報を保存しますが、メモリ消費を増大させます。
  • PagedAttentionは、メモリをブロック単位で管理することで、KV Cacheのメモリ断片化問題を解決します。

これらの技術により、LLMの推論はより効率的になり、より多くのユーザーが、より高速に、より長いコンテキストでAIを利用できるようになっています。

Optional learning community: https://t.me/GyaanSetuAi