PagedAttention
PagedAttention
PagedAttentionは、LLMのKVキャッシュを固定サイズのブロックに分割するメモリ管理アルゴリズムです。推論時のGPUメモリの無駄を削減し、vLLMを支える技術として活用されています。
PagedAttentionとは?
PagedAttentionは、大規模言語モデル(LLM)のKey-Value(KV)キャッシュを固定サイズのブロック(「ページ」)に分割し、非連続的なメモリ割り当てを可能にすることで、推論時のGPUメモリの無駄を劇的に削減するメモリ管理アルゴリズムです。Kwonらによって2023年の論文で発表されたPagedAttentionは、UC Berkeleyで開発された高性能LLM推論エンジンであるvLLMの基盤となっています。
従来のLLMサービングシステムは、KVキャッシュを単一の連続したメモリチャンクとして事前割り当てするため、シーケンスの長さが変動したり、バッチサイズが変化したりする際に大きな無駄が生じていました。PagedAttentionは、オペレーティングシステムの仮想メモリページングにヒントを得て、シーケンスが必要なメモリのみを使用できるようにし、ほぼ最適な利用率を実現します。このアーキテクチャ上の革新により、LLMが推論時にメモリを処理する方法が変革され、より高いスループット、より大きなバッチサイズ、より効率的なハードウェア使用が可能になりました。
このアルゴリズムの影響は学術的な関心を超えて広がっています。LMSYS Chatbot Arena、Databricks、Dropbox、AWS、AMD、NVIDIAなどの主要な本番システムが、LLM推論ワークロードにvLLMとPagedAttentionを採用しています。オープンソースのvLLMプロジェクトは、20,000以上のGitHubスターを獲得し、効率的なLLMサービングにおけるイノベーションを推進し続けています。
コアコンセプト
KVキャッシュの基礎
Key-Valueキャッシュは、トランスフォーマーベースのLLM推論における重要なコンポーネントです。各デコーディングステップで、モデルは各トークンに対してキーとバリューのテンソルを生成し、アテンションメカニズムのためのコンテキスト関係をエンコードします。
目的と機能
- 事前計算されたアテンションキーとバリューを保存し、以前に処理されたトークンの冗長な計算を排除
- 各ステップでコンテキストの再計算を回避することで推論を劇的に高速化
- 効率的な自己回帰テキスト生成を実現
規模と影響
- LLaMA-13Bのようなモデルでは、単一シーケンスのKVキャッシュが1.7GBに達することがある
- より長いシーケンスでのバッチ処理では、総キャッシュサイズが40GB以上になることもある
- キャッシュは大規模かつ動的であり、そのサイズはアクティブなシーケンス、その長さ、バッチサイズに依存
- 非効率的なKVキャッシュ管理は、同時リクエスト処理を著しく制限し、コストを増加させ、ハードウェア利用率を低下させる
メモリフラグメンテーションの課題
フラグメンテーションは、GPUベースのLLM推論における最も重要なボトルネックの1つであり、硬直的なメモリ割り当てスキームから生じます。
内部フラグメンテーション
事前割り当てされたメモリブロックが実際のシーケンス要件を超える場合に発生し、割り当て内に未使用のスペースが残ります。
- シナリオ例: シーケンスが200トークンしか使用しないのに2048トークン分のスペースを事前割り当てすると、残りの容量が無駄になる
- 深刻度: 従来のシステムでは、内部フラグメンテーションによりGPUメモリの60-80%が無駄になる可能性がある
- パフォーマンスへの影響: 同時リクエストを制限し、実効スループットを低下させる
外部フラグメンテーション
さまざまな長さのシーケンスが開始および終了するにつれて発展し、空きメモリが新しい割り当てに適さない小さな非連続的なチャンクに分散します。
- 結果: 十分な総空きメモリがあっても、連続したスペースがないため新しいリクエストの割り当てができない
- 結果: 利用可能な容量があるにもかかわらず、早期のメモリ不足エラーが発生
PagedAttentionは、OSの仮想メモリページングと同様に、オンデマンドのブロック割り当てと解放を通じて、両方のタイプのフラグメンテーションを排除します。
仮想メモリとページングアーキテクチャ
PagedAttentionは、LLM推論のために古典的なオペレーティングシステムの仮想メモリ概念を適応させています。
コア原則
- 仮想メモリ: 論理メモリアドレスを物理的な場所から分離し、非連続的なストレージを可能にする
- ページング: メモリを固定サイズのページに分割し、ページテーブルを介して論理から物理へのマッピングを行う
- LLMへの適用: KVキャッシュをGPUメモリ内のどこにでも配置できるブロックに分割し、シーケンスは論理から物理へのマッピングのためにブロックテーブルを維持
実現される機能
- シーケンスがトークンを生成する際のオンデマンドブロック割り当て
- シーケンス完了時の解放されたブロックの即座の再利用
- コピーオンライトメカニズムによるシーケンス間のメモリ共有
- 事前割り当ての無駄なしに可変長シーケンスを効率的に処理
ブロックテーブル(ページテーブル)
ブロックテーブルは、各シーケンスの論理トークン位置を物理メモリブロックにマッピングするデータ構造です。
機能
- 物理ブロックの場所に関係なく、シーケンスがアテンション計算のためにコンテキストを再構築できるようにする
- 推論中の高速ルックアップと効率的な割り当て/解放をサポート
- 複数の推論ステップにわたってシーケンス状態を維持
パフォーマンスに関する考慮事項
- テーブルルックアップによる小さな計算オーバーヘッドが発生
- オーバーヘッドは、メモリ無駄の劇的な削減によって大幅に相殺される
- 連続的な割り当てでは不可能なメモリ最適化を実現
高度な機能
メモリ共有メカニズム
PagedAttentionは、シーケンスとリクエスト間の洗練されたメモリ共有を可能にし、特に並列サンプリングと高度なデコーディング戦略に有益です。
並列サンプリング
同じプロンプトから複数の補完を効率的に生成します。
- メカニズム: 共有プロンプトKVキャッシュブロックがすべての出力シーケンスによって参照される
- 利点: メモリの重複を排除し、使用量を削減し、より高いスループットを実現
- 実装: 各サンプルのブロックテーブルは、共有シーケンス部分に対して同じ物理ブロックを指す
ビームサーチの最適化
複数のビームが同じプレフィックスを共有することが多く、効率的なメモリ使用が可能になります。
- 共有プレフィックス: 共通のビームプレフィックスは共有メモリブロックを使用
- 分岐処理: コピーオンライトメカニズムがビーム固有の変更を処理
- 混合デコーディング: 冗長な割り当てなしに、貪欲法、サンプリング、ビームサーチ戦略を同時に提供
コピーオンライト
共有メモリブロックの安全な変更を保証するメモリ管理技術です。
動作
- シーケンスが共有ブロックを変更する場合、そのシーケンスのみに新しいコピーが作成される
- 他のシーケンスは元のブロックを参照し続ける
- 共有メモリシナリオでのデータ破損と競合状態を防止
利点
- 正確性を犠牲にすることなく積極的なメモリ共有を実現
- 共有リソースを最大化しながら、必要に応じて分離を維持
- 並列サンプリングとビームサーチの効率に不可欠
技術的実装
ページングを使用したアテンション計算
従来のアテンションカーネルは、キーとバリューのための連続したメモリを期待します。PagedAttentionは、分散したブロックからKVペアを効率的に取得するカスタムカーネルを導入します。
実装の詳細
- カーネルはブロックテーブルを参照して、潜在的に分散したブロックから必要なキーバリューベクトルを収集
- 最適化されたメモリアクセスパターンによりオーバーヘッドを最小化
- 単一シーケンスとバッチ推論の両方をサポート
パフォーマンス特性
- 分散メモリアクセスによる小さな計算オーバーヘッド
- 無駄なメモリの大幅な削減により、より大きなバッチが可能に
- 小さな計算コストにもかかわらず、正味の結果として大幅なスループット向上
vLLM統合
vLLMは、PagedAttentionをコアメモリ管理システムとして実装するオープンソースの高性能LLM推論エンジンです。
主な機能
- 最先端のスループット:HuggingFace Transformersの最大24倍高速
- メモリ無駄の劇的な削減:60-80%から4%未満へ
- 大規模なバッチサイズと長いシーケンスのサポート
- 高度なデコーディング戦略:並列サンプリング、ビームサーチ、混合デコーディング
- HuggingFaceモデル、PyTorch、OpenAI APIとの互換性
デプロイメントオプション
- pipインストールによるローカルデプロイメント
- 主要プラットフォームでのクラウドネイティブデプロイメント
- RunPodなどのプロバイダーを通じたサーバーレスオプション
- Databricks、Dropboxなどとのエンタープライズ統合
モデルサポート
vLLMは幅広いモデルアーキテクチャをサポートしています:
- クラシックトランスフォーマー:Llama、GPT-2、GPT-J、Mistral
- Mixture-of-Experts:Mixtral、Qwen2MoE
- マルチモーダルLLM:LLaVA、StableLM、Qwen-VL
使用方法と実装
インストールとセットアップ
# vLLMをインストール
pip install vllm
# 推論サーバーを実行
python -m vllm.entrypoints.openai.api_server --model <model_name>
<model_name>をサポートされているモデル識別子に置き換えてください。
主なユースケース
高スループットチャットボット
- LMSYS Chatbot Arenaのような本番チャットボットサービス
- 数千の会話の同時処理
- コスト効率の高いスケーリングのための効率的なリソース利用
バッチドキュメント処理
- 大規模なドキュメントコレクションの同時処理
- 広範な知識ベースに対する質問応答
- 要約と分析パイプライン
高度なデコーディングワークロード
- 多様な出力生成のための並列サンプリング
- 高品質なテキスト生成のためのビームサーチ
- 異なるリクエストタイプに対する混合デコーディング戦略
リソース制約のあるデプロイメント
- 限られたGPUメモリでのエッジ推論シナリオ
- 利用可能なハードウェアでのスループット最大化
- クラウドデプロイメントのコスト最適化
パフォーマンスとベンチマーク
定量的改善
スループットの向上
- HuggingFace Transformersに対して最大24倍の改善
- Text Generation Inference(TGI)に対して3.5倍の改善
- 同じハードウェアでより多くのリクエストを処理可能
メモリ効率
- メモリ無駄を60-80%から4%未満に削減
- より高いスループットのためのより大きなバッチサイズを実現
- メモリ圧力なしでより長いシーケンスをサポート
実世界での影響
- LMSYS Chatbot Arena:1秒あたり2-3倍のリクエスト
- 同じワークロードに対してGPU使用量を50%削減
- リソース利用の改善による大幅なコスト削減
業界での採用
本番デプロイメント
- モデル評価のためのLMSYS Chatbot Arena
- エンタープライズLLMアプリケーションのためのDatabricks
- ドキュメント理解のためのDropbox
- クラウドおよびハードウェアプラットフォームのためのAWS、AMD、NVIDIA
- ゲームおよびソーシャルアプリケーションのためのRoblox
オープンソースエコシステム
- 20,000以上のGitHubスター
- アクティブなコミュニティ貢献
- 頻繁なアップデートと改善
- 拡張機能と統合の成長するエコシステム
ベストプラクティス
モデル選択
- 互換性を保証するため、vLLMのサポートリストからモデルを選択
- 利用可能なGPUメモリに対するモデルサイズを考慮
- ハードウェアサイジングのためのスループット要件を評価
バッチサイズの最適化
- 控えめなバッチサイズから始め、モニタリングに基づいて増加
- GPUメモリ利用率を監視し、それに応じて調整
- レイテンシ要件とバッチサイズのバランスを取る
パフォーマンスモニタリング
- スループットメトリクスを追跡:1秒あたりのリクエスト数、1秒あたりのトークン数
- GPUメモリ利用率とフラグメンテーションを監視
- サービス品質のためのレイテンシ分布を分析
デプロイメントに関する考慮事項
- 適切なハードウェアを使用:十分なVRAMを持つNVIDIA GPU
- ワークロード特性に基づいてリソース制限を設定
- 適切な監視とアラートを実装
- 本番環境でのフェイルオーバーと高可用性を計画
参考文献
- Efficient Memory Management for Large Language Model Serving with PagedAttention (arXiv)
- vLLM Blog: Easy, Fast, and Cheap LLM Serving with PagedAttention
- Hopsworks Dictionary: PagedAttention
- Doubleword: Optimizing GPU Memory for LLMs
- Runpod Blog: Introduction to vLLM and PagedAttention
- KV Cache Optimization Deep Dive - Medium
- Red Hat Developer: How PagedAttention Resolves Memory Waste
- YouTube: vLLM/PagedAttention Technical Explanation
- vLLM Documentation
- vLLM GitHub Repository
- Operating Systems: Three Easy Pieces - Virtual Memory
- vLLM Supported Models
関連用語
投機的デコーディング
投機的デコーディングは、高速なドラフトモデルを使用してトークンを提案し、より大規模なターゲットモデルで検証することで、LLM推論を高速化します。出力品質を維持しながらレイテンシを削減します。...