PagedAttention

PagedAttentionは、大規模言語モデル（LLM）のサービング向けのアテンションアルゴリズム兼メモリマネージャーです。各リクエストのKVキャッシュを小さな固定サイズのブロックに分け、GPUメモリ内の任意の場所に配置します。リクエストごとのブロックテーブルでブロックを追跡する仕組みは、OSが仮想メモリのページを割り当てる方法に似ています。UC Berkeleyと提携大学のWoosuk Kwon氏、Zhuohan Li氏らは、vLLM サービングエンジンとともに、この手法を発表しました。発表先はSOSP 2023で、論文のタイトルはEfficient Memory Management for Large Language Model Serving with PagedAttentionです。連続したキャッシュ割り当てによる断片化をなくすことで、1台のGPUで同時に処理できるシーケンス数が大幅に増えます。論文では、同じレイテンシで、FasterTransformerおよびOrcaよりスループットが2～4倍高いと報告しています。

解決するメモリの問題#

テキスト生成中、トランスフォーマーは過去の各トークンのキーと値のベクトルを保持し、再計算を避けます。このキャッシュはレイヤーごとにトークン1つ分ずつ増加し、130億パラメーターのLLaMAモデルでは、単一のシーケンスだけで最大1.7 GBを占めることがあります。リクエストを受け取った時点では出力長が分からないため、以前のサービングシステムは、最大長を想定した連続領域をリクエストごとに予約していました。予約領域の大部分は使われず、領域間の隙間も小さすぎて再利用できませんでした。vLLMチームの測定によると、こうしたシステムではKVキャッシュメモリの60～80%が無駄になり、計算リソースが枯渇するより先にバッチサイズが制限されていました。

ページングの仕組み#

PagedAttentionはキャッシュをブロックに分割し、各ブロックに固定数のトークンのキーと値を格納します。vLLMではデフォルトで16トークンです。リクエストは、直前のブロックが埋まった時点でのみ、空きブロックの中から新しいブロックを割り当てられます。ブロックテーブルには順序が記録されます。アテンション機構のカーネルはこのテーブルに従って、キーと値をブロック単位で取得するため、シーケンスのキャッシュを連続した領域に置く必要はありません。無駄になるのは各シーケンスの最後のブロックの未使用部分だけで、vLLMの測定では4%未満です。

以下のPyTorchの例では、共有プール内の分散したブロックに1つのリクエストのキーを保存し、ブロックテーブルを介して読み戻します。

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

ブロックテーブル[41, 7, 19]とkeys intact: Trueが出力されます。次に、ページングを使うリクエストの未使用スロット数として8が出力されます。2,048トークン用に予約した領域では2,008個です。

シーケンス間でのブロック共有#

ブロックテーブルを使うと、複数のシーケンスから同じ物理ブロックを参照できます。並列サンプリング、ビームサーチ、同じシステムプロンプトから始まるリクエストでは、共有プレフィックスのコピーを1つだけ再利用し、各ブロックの参照カウントで管理します。ほかのシーケンスも使用しているブロックに書き込む必要が生じた場合は、エンジンが先にそのブロックをコピーします。これはOSから着想を得たコピーオンライト処理です。

vLLMの著者らは、この共有によって並列サンプリングとビームサーチのメモリ使用量が最大55%削減され、スループットが最大2.2倍になると報告しています。ブロック全体をハッシュ化すると、この考え方を複数のリクエスト間に拡張できます。これがvLLMのプロンプトキャッシュの仕組みです。

関連手法#

PagedAttentionは、キャッシュ済みのキーと値をどこに配置するかを決めます。保存量や読み込み速度を変える手法と組み合わせることもできます。

• Flash Attention： アテンション計算をタイル状に分割し、スコア行列全体がGPUのメインメモリに書き込まれるのを防ぎます。現在のエンジンは、ページングされたブロックから読み込むFlashAttention形式のカーネルを実行します。
• 連続バッチ処理： 別のリクエストが完了するたびに新しいリクエストを受け入れます。ページングでは完了したシーケンスのブロックが直ちに空きプールに戻されるため、新しいリクエスト用の領域を確保できます。
• Grouped-query attention： キーと値のヘッドをクエリヘッド間で共有し、各ブロックに格納するデータ量を減らします。
• FP8キャッシュの量子化： キーと値を8ビット精度で保存します。16ビット形式と比べて、キャッシュメモリをおよそ半分にできます。

Ultralytics LLMインターフェースはリクエストの引数を変更せずに転送します。そのため、n=3を指定したChat Completions呼び出しを、base_url経由でローカルのvLLMサーバーに送信すると、プロンプトのキャッシュブロックを共有する3つのサンプルが返されます。新しいメモリを使うのは、それぞれのサンプルで分岐した後続部分だけです。