複数のローカルAIリクエストが同時に実行される場合、Paged AttentionはKVキャッシュをどのように管理するのか？

エヴァ・ウォン

IceWhale author

エヴァ・ウォン は テクニカルライター であり ZimaSpaceの常駐ティンカーでもあります。 生涯のオタクであり、 ホームラボとオープンソースソフトウェアに情熱を持っています。彼女は複雑な技術的概念をわかりやすく、 実践的なガイドに翻訳することを専門としています。エヴァはセルフホスティングは楽しくあるべきで、怖がるものではないと信じています。彼女のチュートリアルを通じて、コミュニティが ハードウェアのセットアップを解明する手助けをしています。初めてのNAS構築からDockerコンテナの習得まで。

How Does Paged Attention Manage KV Cache Under Concurrent Local AI Requests? - Zima Store Online

Paged Attentionは、各リクエストで増加するトークン状態を固定サイズのブロックに割り当てます。これらのブロックは個別に確保・再利用できるため、同時実行されるKVキャッシュを効率的に管理できます。

これはホームAIサーバーで重要です。アクティブなチャット、ドキュメント処理、エージェントは、同じコンテキスト長になることも、同じタイミングで終了することもほとんどありません。連続した領域を確保する方式では、このように長さが異なるシーケンスの周囲で、限られたアクセラレータメモリを無駄にしてしまいます。ページングによってメモリレイアウトが変わり、スケジューラーは利用可能な容量をより密に使えるようになります。

同時実行リクエストでは、異なる速度で増加するKVキャッシュが作られる

アクティブな自己回帰リクエストでは、処理済みトークンのキーとバリューの状態が必要です。あるチャットは100トークンで終了する一方、別のチャットは数千トークンまで続くことがあります。そのため、バッチ開始時点で必要なメモリ量を正確に把握することはできません。

PagedAttentionの論文では、動的に増加するKVキャッシュメモリがサービングのボトルネックとなり、割り当てが非効率だとバッチ処理を制限すると説明されています。

複数のユーザーやバックグラウンドエージェントが動作するホームサーバーでは、この変動によって、大きな空きメモリ容量がサイズの異なる多数の使用中割り当てへと分かれていきます。

固定サイズの連続領域を予約すると、長さの異なるシーケンスの周囲でメモリが無駄になる

単純なランタイムでは、最大長を基準にリクエストごとに大きな連続領域を1つ予約することがあります。短いリクエストでは予約領域の内部に未使用スペースが残り、長いリクエストでは再割り当てや過剰なメモリ確保が必要になる場合があります。

vAttentionは、ページングの元々の動機を次のように要約しています。事前にKV領域を予約すると内部断片化が発生するのです。無駄になったメモリによって、同時にアクティブにできるリクエスト数が減少します。

実際には、使用中のトークン状態の合計が利用可能なVRAMより少なく見えるにもかかわらず、メモリ不足の制限に達したり、最大バッチサイズが小さくなったりします。

PagedAttentionはKV状態を固定サイズのブロックに分割する

PagedAttentionは、物理的に連続した1つのリクエストバッファーを、論理的なブロック列に置き換えます。トークンが増えると、ランタイムは共有プールから追加の物理ブロックを割り当てます。

vLLMの説明によると、KVブロックは物理的に連続していないメモリ上に配置できます。ブロックテーブルが、リクエストの論理的なトークン順序をそれらの物理的な場所に対応付けます。

モデルは引き続き正しいシーケンスに対してアテンションを実行します。この間接参照によって変わるのはメモリの配置であり、会話の意味上の順序ではありません。

-15% OFF

ZimaBoard 2 - あなたの大きなアイデアのためのミニホームサーバー

Single board computer zimaboard2

リクエストが終了すると、解放されたブロックは共有プールに戻る

あるリクエストが完了すると、そのリクエストだけが使用していたブロックはすぐに解放され、別のアクティブなリクエストや新たに受け付けたリクエストで再利用できます。サーバーは、巨大なアリーナ全体が空になるまで待つ必要がありません。

NVIDIAの推論スタックでは、KVキャッシュ容量をサービングリソースとして扱い、モデルの重みやランタイムのワークスペースとのバランスを取る必要があります。

このブロック単位の再利用は、リクエストの長さが不規則になりやすい家庭環境で有効です。ある人が短いスマートホームに関する質問をする一方で、別の人が長いドキュメントアシスタントを実行し、バックグラウンドエージェントがファイルを要約することもあります。

ページングは各トークンを小さくするのではなく、断片化を減らして同時実行数を増やす

PagedAttentionは、1トークン分のキーとバリューのテンソルを根本的に小さくするものではありません。主なメモリ上のメリットは、シーケンスの実際の増加量に近い形で割り当て、可能な場合に共有を有効にすることから生まれます。

PyTorchの分散処理ガイダンスでは、モデル並列化とサービングメモリ管理を区別しています。ページングはテンソル並列化ではなく、それ自体でモデルの重みを複数のGPUに分割するものでもありません。

モデルの重みがすでにデバイスメモリのほぼすべてを占有している場合、ページングによって大規模な同時実行に必要な容量を生み出すことはできません。ページングが改善するのは、残っているKV領域の利用効率です。

ブロック管理には固有のオーバーヘッドとスケジューリング上の制限がある

ランタイムは、ブロックテーブル、参照カウント、割り当て判断、そしてページングされたレイアウトに対応するアテンションカーネルを管理する必要があります。バッチが極端に小さい場合や、対応していないハードウェアでは、メリットが小さくなる可能性があります。

ZimaSpaceの家族での同時利用とKV負荷に関する記事では、ユーザーが体感する限界を説明しています。メモリ効率が重要なのは、待ち行列が長くなる前に、いくつのアクティブなコンテキストを共存させられるかが決まるためです。

1つの合成バッチだけでなく、複数のリクエスト長と到着パターンをテストしてください。ZimaSpaceのスライディングコンテキストウィンドウとKVメモリに関するガイドは、比較の参考になります。ウィンドウ方式は各シーケンスが保持する状態量を変えるのに対し、ページングはその状態の割り当て方法を変えます。アクティブなシーケンス数、使用中のKVブロック数、キュー待ち時間、プリエンプションを記録してください。

Hugging FaceのTGIスケジューラーも、同時サービングにおけるプリフィルとデコードのスケジューリングを区別しています。ページングされたメモリはスケジューラーによる処理の受け入れを助けますが、公平性とレイテンシーは、処理がどのようにキューへ入れられるかにも左右されます。