LLM スループットは 2 ～ 4 倍に増加し、モデルが大きくなるほど効果も高まります。カリフォルニア大学バークレー校、スタンフォード大学などのオープンソースの効率的なメモリ管理メカニズム PagedAttention

大規模言語モデル (LLM) のパフォーマンスは驚異的ですが、ユーザーからのリクエストを受け取るたびに大量のビデオメモリとコンピューティングリソースを消費します。リクエスト数が予想を超えると、ダウンタイム、キューイング、高レイテンシなど、ChatGPT が最初にリリースされたときと同じジレンマに直面する可能性が非常に高くなります。

高スループットの LLM サービスを構築するには、モデルは 1 回のバッチでできるだけ多くのリクエストを処理する必要があります。しかし、既存のシステムのほとんどは、リクエストを処理するたびに大量のキー値 (KV) キャッシュを適用します。管理効率が高くない場合、断片化と冗長レプリケーションで大量のメモリが浪費され、バッチ サイズの増加が制限されます。

最近、カリフォルニア大学バークレー校、スタンフォード大学、カリフォルニア大学サンディエゴ校の研究者らは、オペレーティングシステムの古典的な仮想メモリとページング技術に基づいた新しいアテンションアルゴリズムPagedAttentionを提案し、LLMサービスシステムvLLMを作成した。

論文リンク: https://arxiv.org/pdf/2309.06180.pdf

オープンソースリンク: https://github.com/vllm-project/vllm

vLLM は KV キャッシュの無駄をほぼゼロにし、「リクエスト内」および「リクエスト間」で KV キャッシュを柔軟に共有できるため、メモリ使用量がさらに削減されます。

評価結果によると、vLLM は一般的に使用される LLM のスループットを 2 ～ 4 倍向上でき、レイテンシの点では最先端のシステム (FasterTransformer や Orca など) に匹敵し、シーケンスが長く、モデルが大きく、デコード アルゴリズムが複雑なほど改善が顕著になります。

ページングされた注意

アテンション メカニズムのメモリ管理問題を解決するために、研究者は新しいアテンション アルゴリズム PagedAttention を開発し、集中型スケジューラを使用して分散 GPU ワーカー スレッドの実行を調整する LLM サービス エンジン vLLM を構築しました。

1. アルゴリズム

PagedAttention は、オペレーティング システムのページング アルゴリズムにヒントを得て、シーケンス内の KV キャッシュを KV ブロックに分割します。各ブロックには、固定数のトークンのキー (K) と値 (V) のベクトルが含まれ、アテンション計算がブロック レベルの操作に変換されます。

注意計算中に、PagedAttention カーネルはそれぞれ異なる KV ブロックを識別して取得します。たとえば、次の例では、キー ベクトルと値ベクトルが 3 つのブロックに分散されており、3 つのブロックは物理メモリ内で不連続です。次に、クエリ ベクトルをブロック内のキー ベクトルに乗じて部分注意スコアを取得し、次にブロック内の値ベクトルに乗じて最終的な注意出力を取得します。

この設計により、KV ブロックを非連続の物理メモリに格納できるようになり、vLLM でのページング メモリ管理がより柔軟になります。

2. KVキャッシュマネージャー

オペレーティング システムは、メモリを複数の固定サイズのページに分割し、ユーザー プログラムの論理ページを物理ページにマップします。連続した論理ページは、連続していない物理メモリ ページに対応できるため、ユーザーがメモリにアクセスすると、連続しているように見えます。

さらに、物理メモリ空間を事前に完全に予約する必要がないため、オペレーティング システムは需要に応じて物理ページを動的に割り当てることができます。

vLLM は、PageAttention によって分割された KV ブロックを通じて、仮想メモリ メカニズムを使用して KV キャッシュを一連の論理 KV ブロックとして表現し、新しいトークンと KV キャッシュを生成するときに左から右にそれらを埋めていきます。最後の KV ブロックの埋められていない位置は、後続の生成操作のために予約されます。

KV ブロック マネージャーは、各リクエストの論理 KV ブロックと物理 KV ブロック間のマッピングであるブロック テーブルを維持する役割も担います。

論理 KV ブロックと物理 KV ブロックを分離することで、vLLM はすべての場所に KV キャッシュ メモリを事前に予約することなく動的に拡張できるため、既存のシステムでのメモリの無駄の大部分が排除されます。

3. デコード

次の例は、vLLM が単一の入力シーケンスのデコード処理中に PagedAttention を実行し、メモリを管理する方法を示しています。

① オペレーティングシステムの仮想メモリと同様に、vLLM は最初に生成されるシーケンスの最大長分のメモリを予約する必要はなく、即時の計算中に生成される KV キャッシュに対応するために必要な KV ブロックのみを予約します。

プロンプト ワードには 7 つのトークンが含まれているため、vLLM は最初の 2 つの論理 KV ブロック (0 と 1) を 2 つの物理 KV ブロック (7 と 1) にマッピングします。事前入力ステップでは、vLLM は自己注意アルゴリズムを使用してプロンプトと最初の出力トークンの KV キャッシュを生成します。次に、最初の 4 つのトークンの KV キャッシュが論理ブロック 0 に格納され、次の 3 つのトークンが論理ブロック 1 に格納されます。残りのスロットは、後続の自己回帰生成用に予約されます。

②最初の自己回帰復号ステップでは、vLLMは物理ブロック7と1にPagedAttentionアルゴリズムを使用して新しいトークンを生成する。

最後の論理ブロックにはまだ 1 つのスロットが使用可能なため、新しく生成された KV キャッシュがこのスロットに格納され、ブロック テーブルの #filled レコードが更新されます。

③ 2回目のデコードステップでは、最後の論理ブロックがいっぱいになると、vLLMは新しく生成されたKVキャッシュを新しい論理ブロックに保存し、それに新しい物理ブロック（物理ブロック3）を割り当てて、ブロックテーブルに保存されるようにマッピングします。

LLM の計算プロセス中に、vLLM は PagedAttention カーネルを使用して、以前に論理 KV ブロックの形式で保存された KV キャッシュにアクセスし、新しく生成された KV キャッシュを物理 KV ブロックに保存します。

1 つの KV ブロック (ブロック サイズ > 1) に複数のトークンを格納すると、PagedAttention カーネルはより多くの場所で KV キャッシュを並列処理できるようになり、ハードウェアの使用率が向上し、レイテンシが短縮されますが、ブロック サイズが大きくなるとメモリの断片化も増加します。

トークンとその KV キャッシュがどんどん生成されると、vLLM は新しい物理ブロックを論理ブロックに動的に割り当て、すべてのブロックを左から右に埋め、以前のブロックがすべていっぱいになった場合にのみ新しい物理ブロックを割り当てます。つまり、リクエストのすべてのメモリ浪費は 1 つのブロックに制限され、すべてのメモリを効果的に利用できるため、より多くのリクエストをメモリに入れてバッチ処理できるようになり、スループットが向上します。リクエストが生成されると、その KV ブロックを解放して、他のリクエストの KV キャッシュを保存できます。

4. ユニバーサルデコード

このアルゴリズムは、貪欲なデコードとサンプリング、および単一のユーザープロンプト入力から単一の出力シーケンスを生成するなどの基本的なシナリオのサポートに加えて、並列サンプリング、ビーム検索、共有プレフィックスなどのより複雑なデコードシナリオもサポートできます。

5. スケジューリングとプリエンプション

要求トラフィックがシステム容量を超えると、vLLM は公平性を確保し、リソース不足を防ぐために、先着順 (FCFS) スケジューリング戦略を使用して要求サブセットに優先順位を付ける必要があります。

ただし、LLM への入力ヒントは長さが大きく異なる場合があり、出力の長さは入力ヒントとモデルに応じて事前に不明です。リクエストの数とその出力が増加すると、vLLM は新しく生成された KV キャッシュを保存するための GPU の物理ブロックを使い果たす可能性があります。

スワッピングは、解放されたページをディスク上のスワップ領域にコピーするために、ほとんどの仮想メモリ アルゴリズムで使用される従来の手法です。

GPU ブロック アロケータに加えて、vLLM には CPU RAM にスワップされた物理ブロックを管理する CPU ブロック アロケータも含まれています。vLLM で新しいトークン用の空き物理ブロックがなくなると、一連のシーケンスが選択され、KV キャッシュが解放されて CPU に転送されます。

この設計では、CPU RAM にスワップされるブロックの数は GPU RAM 内の物理ブロックの合計数を超えることはないため、CPU RAM 上のスワップ領域は KV キャッシュに割り当てられた GPU メモリによって制限されます。

再計算: プリエンプトされたシーケンスが再スケジュールされると、KV キャッシュは単純に再計算され、そのレイテンシは元のレイテンシよりも大幅に低くなります。これは、デコード中に生成されたトークンを元のユーザー プロンプトと新しいプロンプトとして連結し、すべての位置の KV キャッシュを 1 回のプロンプト フェーズ反復で生成できるためです。

スワッピングと再計算のパフォーマンスは、CPU、RAM、GPU メモリ間の帯域幅と、GPU の計算能力によって異なります。

6. 分散実行

vLLM は、SPMD (Single Program Multiple Data) 実行スケジュールに従う Megatron-LM スタイルのテンソル モデル並列戦略をサポートします。この戦略では、線形レイヤーが分割されてブロック単位の行列乗算が実行され、GPU は allreduce 操作を通じて中間結果を常に同期します。

具体的には、アテンション オペレータはアテンション ヘッド次元で分割され、各 SPMD プロセスはマルチヘッド アテンション内のアテンション ヘッドのサブセットを担当します。ただし、各モデル シャードは引き続き同じ入力トークン セットを処理するため、同じ場所に KV キャッシュが必要です。

異なる GPU ワーカーは、マネージャーと論理ブロックから物理ブロックへのマッピングを共有し、各入力要求に対してスケジューラによって提供される物理ブロックを使用してモデルを実行します。各 GPU ワーカー スレッドは同じ物理ブロック ID を持ちますが、1 つのワーカー スレッドには、対応するアテンション ヘッドの KV キャッシュの一部のみが格納されます。

各ステップでは、スケジューラはまずバッチ内の各リクエストの入力トークン ID を含むメッセージと、各リクエストのブロック テーブルを準備します。

次に、スケジューラは制御メッセージを GPU ワーカーにブロードキャストし、GPU ワーカーは入力トークン ID を使用してモデルを実行します。アテンション レイヤーでは、制御メッセージ内のブロック テーブルに従って KV キャッシュが読み取られます。実行中、中間結果はスケジューラの調整なしに all-reduce 通信プリミティブと同期されます。

最後に、GPU ワーカーはその反復のサンプリング トークンをスケジューラに送り返します。

評価結果

基本的なサンプリング

ShareGPT データセットでは、リクエスト レートが増加すると、レイテンシは最初はゆっくりと増加しますが、その後突然急増します。これは、リクエスト レートがサービス システムの容量を超えるとキューの長さが無限に増加するためと考えられます。

vLLM は、PagedAttention がメモリ使用量を効果的に管理し、Orca よりも多くのリクエストをバッチで処理できるため、同様のレイテンシを維持しながら、Orca よりも 1.7 ～ 2.7 倍、Orca(Max) よりも 2.7 ～ 8 倍高いリクエスト レートを維持できます。

OPT-13B モデルの場合、vLLM は Orca よりも 2.2 倍、Orca (Max) よりも 4.3 倍多くのリクエストを同時に処理します。

FasterTransformer と比較すると、vLLM は最大 22 倍のリクエスト レートを実現します。これは、vLLM がきめ細かいスケジューリング メカニズムを活用しておらず、Orca (Max) のようなメモリ管理が非効率的であるためと考えられます。

複数のシーケンス

並列サンプリングでは、リクエスト内のすべての並列シーケンスがヒントの KV キャッシュを共有でき、サンプリングされたシーケンスの数が増えるにつれて、vLLM は Orca ベースラインよりも大きな改善を実現します。

ビーム検索では共有コンテンツが多くなるため、vLLM ではパフォーマンス上の利点が大きくなります。

OPT-13B および Alpaca データセットでは、Orca (Oracle) に対する vLLM の改善は、基本サンプリングでは 1.3 倍、幅 6 のビーム検索では 2.3 倍に増加します。

メモリ節約は、共有によって節約されたブロックの数を共有されていないブロックの合計数で割ることによって計算されました。結果は、並列サンプリングによってメモリが 6.1% ～ 9.8% 節約され、ビーム検索によってメモリが 37.6% ～ 55.2% 節約されたことを示しました。

ShareGPT データセットを使用した同じ実験では、並列サンプリングによってメモリが 16.2% ～ 30.5% 節約され、ビーム検索によってメモリが 44.3% ～ 66.3% 節約されることがわかります。