論文まとめ: vLLM(Paged Attention)
背景
vLLMの推論が早すぎてすごい。裏側でどのようなことが行われているのか知りたいため、論文を読んだ。その内容をまとめる。
論文メタ情報
タイトル
Efficient Memory Management for Large Language Model Serving with PagedAttention
URL
著者
概要
- 大規模言語モデル（LLM）の高スループット処理には、一度に十分な数のリクエストをバッチ処理する必要がある。
- しかし、既存のシステムでは、各リクエストの KV キャッシュメモリが膨大で、動的に増減するため、処理が困難である。非効率的に管理されると、断片化や冗長な重複によってメモリが大幅に浪費され、バッチサイズが制限される可能性がある。
- この問題に対処するため、オペレーティングシステムの古典的な仮想メモリとページング技術に着想を得たアテンションアルゴリズムである PagedAttention を提案する。
- KV キャッシュメモリの無駄をほぼゼロに抑え、リクエスト内およびリクエスト間で KV キャッシュを柔軟に共有することでメモリ使用量をさらに削減する LLM 処理システム vLLM を構築する。
- 評価の結果、vLLM は、FasterTransformer や Orca などの最先端システムと比較して、同じレイテンシで一般的な LLM のスループットを 2～4 倍向上させることが示された。
- この改善は、シーケンスが長いほど、モデルが大規模になるほど、そしてデコードアルゴリズムが複雑になるほど顕著である。
- vLLM のソースコードは https://github.com/vllm-project/vllm で公開されている。
詳細
背景
KV Cache は、Attention 機構の計算に用いる Key と Value の値をメモリに格納することで、メモリを消費する代わりに計算を高速にする方法である。また、LLM が高いスループットを出すためには、バッチ処理が必要である。
課題
KV Cache とバッチ処理に関する課題を以下の図を用いて説明する。最大トークン数がそれぞれ 2048, 512 と異なる 2 つのリクエストに対する KV Cache でのメモリ管理方法のイメージである。
意味はそれぞれ以下である。
| 要素 | 意味 | 必要かどうか |
|---|---|---|
| KV Cache states | KV Cache を行うために予約されたメモリ | 必要 |
| generated token | 生成されたトークン | 必要 |
| Reserved | リクエスト処理中において、将来生成するかもしれないトークン用に予約されたメモリ | 必要だが、過剰に確保していることが多い。 |
| Internal fragmentation | リクエスト終了後において、過剰に予約しすぎていたことがわかったメモリ | 不要 |
| External fragmentation | リクエストごとに異なるサイズのメモリを割り当てた結果、断片化してしまい、次のリクエストの処理に利用できないメモリ | 不要 |
バッチ処理を行う前提において、KV Cache の仕組みには、以下 3 つの課題が存在する。
- Internal fragmentationの問題。リクエストの KV キャッシュを連続した空間に格納するために、リクエストの最大長 (たとえば、2048 トークン) を持つ連続したメモリチャンクを事前に割り当てる。これは、リクエストの実際の長さが最大長よりもはるかに短くなる可能性があるため、深刻な内部フラグメンテーションを引き起こす可能性がある。
- Reserved（事前割り当て）が非効率である。リクエスト中、チャンク全体が予約されるため、他の短いリクエストは、現在未使用のチャンク部分を利用できない。
- External fragmentationの問題。事前割り当てサイズはリクエストごとに異なる可能性があるため、メモリが断片的に空くが他の処理に利用することができない状態、つまり外部メモリのフラグメンテーションが深刻になる可能性がある。
提案手法
Paged Attention はバッチ処理における KV Cache の無駄をなくし、より効率的にメモリを利用する手法である。
従来は、連続したメモリブロックに Key Value の値を格納する必要があった。しかし、Paged Attention を用いると、非連続なメモリブロックに Key Value の値を格納できるようになる。
以下の Figure5 は、「forth」という Query に対して、非連続なメモリブロックに Key と Value が格納されている様子である。
では、非連続なメモリブロックへの格納をどう実現しているかについて説明していく。以下の Figure6 は Paged Attention の仕組みを説明している。
まず、KV キャッシュを KV ブロックと呼ばれる固定サイズのブロックに分割する。更に、Logical KV blocks(論理 KV ブロック)と、Physical KV blocks(物理 KV ブロック)に分けて管理する。論理 KV ブロックは、論理上の KV ブロックの配置であり。Prompt を入力時と同じ順番になるように保持する。物理 KV ブロックは、実際の物理メモリ上の配置である。これらを、Block Table で紐付ける。この仕組みによりブロックを連続した空間に格納する必要がなくなる。
KV ブロックを利用することにより、以下のことが実現できる。
- 小さなブロックを使用し、必要に応じて動的に割り当てることで、Reserved、Internal fragmentation を軽減する。
- すべてのブロックが同じサイズであるため、サイズが可変長のリクエストを受け入れることが原因だった External fragmentation を解消する。
2つのリクエストを同時に処理する際のイメージを Figure7 に載せる。
2つのシーケンスの論理 KV ブロックは、GPUワーカーのブロックエンジンによって予約された空間内の異なる物理 KV ブロックにマッピングされる。両シーケンスの連続する論理 KV ブロックは、物理メモリ内で連続している必要はなく、物理 KV ブロックの空間は両方のシーケンスで有効に活用できる。
さらに、同じリクエストに関連する異なるシーケンス間、あるいは異なるリクエスト間でも、ブロック単位でのメモリ共有が可能になる。Figure8 に異なるリクエストにおけるブロック単位でのメモリ共有イメージを示す。
Figure 8 は、2 つの出力の並列デコードの例を示している。両方の出力は同じプロンプトを共有するため、プロンプトのコピー 1 つ分のみスペースを予約する。両方のシーケンスのプロンプトの論理ブロックは同じ物理ブロックにマップされる。1 つの物理ブロックを複数の論理ブロックにマップできるため、各物理ブロックに Ref count を導入する。この場合、物理ブロック 7 と 1 の Ref count はどちらも 2 である。
生成フェーズにおいて、vLLM は、複数のシーケンスによる変更が必要な物理ブロックに対して、ブロック粒度で Copy-on-write メカニズムを実装する。具体的には、図 8 では、Sample A1 が最後の論理ブロック (論理ブロック 1) に書き込む必要がある場合、vLLM は対応する物理ブロック (物理ブロック 1) の Ref count が 1 より大きい時は、新しい物理ブロック (物理ブロック 3) に情報をコピーし、Ref count を 1 に減らす。次に、Sample A2 が物理ブロック 1 に書き込む場合、Ref count が 1 であるため、A2 は新しい KV キャッシュを物理ブロック 1 に上書きする
実験
Figure1 より、vLLM は既存システムより、バッチサイズが増加してもメモリの使用量の増加割合が少ない。よって、スループットが改善される。
Figure2 より、従来のサービングシステムである Orca が、KV cache で 本来 cache したい要素である Token states に対し、KV Cache に用いているメモリのうち、約 20.4 ~ 38.2% 程度しか利用できていなかった点に対し、 vLLM は約 96.3% を利用できている。
Figure12 は単一の文章生成における Request Rate に対する latency を、 Figure14 は並列での文章生成における成における Request Rate に対する latency を示している。vLLM が最も小さいことがわかる。
終わり
- 本論文では、Attention の Key と Value を非連続なページメモリに格納できる新しいアテンションアルゴリズムである PagedAttention を提案し、PagedAttention によって効率的なメモリ管理を実現する高スループット LLM サービングシステム vLLM を紹介する。
- オペレーティングシステムに着想を得て、仮想メモリやコピーオンライトといった既存の技術をLLMサービングにおけるKVキャッシュの効率的な管理と様々なデコードアルゴリズムの処理に適用する方法を示した。
- 実験では、vLLMは最先端システムと比較して 2～4 倍のスループット向上を実現することが示された。
思ったこと
- 単一の処理についても早くなるが、バッチ処理を前提にすることでより高速化が可能な手法だと理解した。ここを理解していないと、vLLM をバッチ処理なし利用してしまい、より高速化可能であることに気付けない、などあるため、読んでおいてよかった。
- ただ、↑の話は、Continuous batching を用いることで解決できるかもしれない。
Discussion