Hugging Face: 原理から学ぶ連続バッチング

連続バッチングは、複数の会話を並列処理し、完了次第すぐに交換することでスループットを最大化する、Large Language Model (LLM) サービスのための重要な最適化です。パディングの必要性をなくし、プリフィルフェーズとデコードフェーズを同時に処理することにより、連続バッチングはAIサービスが数千人の同時ユーザーを効率的に処理できるようにします。

The Role of Attention and the Prefill Phase

LLM生成の中心は注意機構であり、これはネットワークの中で異なるトークンが相互作用する唯一の部分です。標準的なフォワードパスでは、モデルは入力シーケンス（プリフィルフェーズ）を処理して次のトークンを予測します。

このプロセスでは、入力トークンをクエリ ($Q$)、キー ($K$)、バリュー ($V$) 状態に射影します。その後、モデルは $Q K^{T}$ を介して類似度スコアを計算します。これはシーケンス長 $n$ に対して二次計算量 $\mathcal{O}(n^{2} d)$ を持ちます。因果注意マスクを適用して、各トークンがそれ以前のトークンのみと相互作用するようにし、将来のトークンが過去に影響を与えないようにします。

Optimizing Generation with KV Caching

トークンを一つずつ生成する（デコードフェーズ）ことは、モデルが新しいトークンを生成するたびにシーケンス内のすべての以前のトークンの $K$ と $V$ 状態を再計算する場合、計算的に無駄になります。

KVキャッシングは、プリフィルおよび以前のデコードステップで作成されたキーとバリューの状態を保存することでこれを解決します。これにより、トークン $n + 1$ の生成コストが $\mathcal{O}(n^{2})$ から $\mathcal{O}(n)$ に削減され、計算をメモリとトレードオフします。たとえば、Llama-2-7Bでは、1つのトークンを保存するために、32層と32ヘッドにわたってfloat16精度で約16KBのメモリが必要です。

Handling Large Prompts via Chunked Prefill

初期プロンプトが exceptionally long（非常に長い）場合、アクティベーションに必要なメモリが利用可能なGPUメモリを超えることがあり、これにより単一のフォワードパスでプリフィルを実行することが不可能になります。

チャンクプリフィルは、初期プロンプトを小さく管理可能なチャンクに分割することでこれに対処します。KVキャッシュを利用することで、モデルは最初のチャンクの状態を保存し、その後のチャンクの状態の先頭に追加します。これにより、情報を失うことなくプロンプトを増分的に処理できます。

From Batched Generation to Continuous Batching

スループット（1秒あたりに生成されるトークン数）を向上させるため、モデルは複数のプロンプトを並列で処理します。ただし、従来の方法は大きな非効率をもたらします：

  • バッチドジェネレーション: すべてのプロンプトが同じ長さである必要があり、パディングトークン（<pad>）が必要となります。プロンプトの長さが異なる場合、これは無駄です。
  • ダイナミックスケジューリング: 完了したプロンプトを新しいものと置き換えることを許可しますが、プリフィルフェーズの新しいプロンプトはデコードフェーズの他のプロンプトとバッチ処理される必要があるため、依然として大量のパディングが必要です。

ラグドバッチング

パディングの無駄をなくすため、ラグドバッチングはバッチ軸を完全に取り除き、代わりにすべてのプロンプトを単一のシーケンスに連結します。トークンの相互作用は注意マスクを介して厳密に制御され、あるプロンプトのトークンが別のプロンプトのトークンと相互作用しないようにします。

連続バッチングアルゴリズム

連続バッチングは、ラグドバッチングとダイナミックスケジューリングを組み合わせてGPUをフルに活用します。このプロセスは次の論理ステップに従います：

  1. メモリ予算を最大化: システムは、バッチあたり $m$ トークンのGPUメモリ予算を満たそうとします。
  2. デコードを優先: 現在デコードフェーズにあるすべてのプロンプトを最初にバッチに追加します（それぞれ1トークンを占有します）。
  3. プリフィルで埋める: 残りのスペースは、プリフィルフェーズにあるプロンプトで埋められ、必要に応じてチャンクプリフィルを使用して入力を分割します。
  4. ダイナミック置換: <eos> トークンを生成する完了したプロンプトはすぐに削除され、新しい着信リクエストで置き換えられます。

Summary of Key Techniques

主要な技術のまとめ

連続バッチングは、以下の3つの主要な技術コンポーネントを統合することでその効率を達成します。

テクニック 主な利益
KVキャッシング デコード中に過去のトークン表現の再計算を避ける。
チャンクプリフィル GPUメモリ制限を超えるプロンプトの処理を可能にする。
ラグドバッチング シーケンスを連結しマスクを使用することでパディングの無駄をなくす。
ダイナミックスケジューリング リアルタイムでリクエストを交換することで高いスループットを維持する。

Sources