連続バッチ処理アーキテクチャ

従来のバッチ処理では、固定のグループの要求をまとめて処理し、最も遅い要求が完了するまで待ってから次のグループを開始します。これにより、バッチの間にGPUがアイドル状態になります。

Continuous batchingでは、各生成ステップにおいて、スケジューラーが完了したリクエストを確認し、待機中のリクエストと即座に置き換えます。短いリクエストは完了した時点で即座に排出され、GPUは常に稼働した状態を維持します。その結果、スループットが大幅に向上し、平均レイテンシが低減されます。

リクエストのライフサイクル

リクエストは、送信から完了まで4つの状態を経由します。

       WAIT IN QUEUE          LOAD PROMPT INTO KV       STREAM OUTPUT           DONE
    ┌────────────────┐    ┌─────────────────────┐    ┌─────────────────┐    ┌────────┐
    │    PENDING     │───▶│     PREFILLING      │───▶│    DECODING     │───▶│FINISHED│
    │  ○ ○ ○ · · ·   │    │ ████████░░░░░░░░░░  │    │  → → → → · ·    │    │   ✓    │
    └────────────────┘    └─────────────────────┘    └─────────────────┘    └────────┘
         others ahead            prompt chunks              +1 token/step

1. Pending — リクエストがキューに追加され、スケジュールの割り当てを待っています。
2. Prefilling — プロンプトのトークンをフォワードパスで処理します。
3. デコーディング — 出力トークンは1つずつ生成されます。
4. 完了 — 生成が完了し、結果が利用可能になりました。

スケジューラがリクエストを許可するのに十分なトークン予算とキャッシュスペースを見つけた時点で、リクエストは pending から prefilling に移動します。プロンプトが1ステップに収まらないほど長い場合、スケジューラはそれを複数の forward pass に分割します（chunk prefill）。

チャンク化プリフィル

1ステップで長いプロンプトを処理するのは、他のリクエストの生成をブロックするためコストが高くなります。チャンク化プリフィルは、大きなプリフィルを複数のステップに分割することで、この問題を解決します。

Blocking (long prefill blocks the batch)
  Step:   1        2        3        4
  Req A  [▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓]   ← one huge prefill
  Req B  ················ wait ················
  Req C  ················ wait ················

Chunked (same prompt split and interleave with decode)
  Step:   1          2          3          4
  Req A  [▓▓ pre]    [▓▓ pre]   [▓▓ pre]   [→ dec]
  Req B  [→ dec]     [→ dec]    [▓▓ pre]   [→ dec]
  Req C  [→ dec]     · idle ·   [→ dec]    · idle ·

新しいリクエストのプロンプトが利用可能なトークン予算（ContinuousBatchingConfigのmax_batch_tokensで設定）を超えた場合、スケジューラは可能な限り多くのトークンを処理し、残りは保持します。その後のステップでは、中断した箇所から処理を再開し、prefill作業と進行中のdecodeステップを交互に実行します。これにより、バッチの生産性を維持し、他のリクエストの最初のトークンまでの時間を短縮できます。

スケジューラー

スケジューラは、2つの予算とリクエスト上限に基づいて、各フォワードパスに参加するリクエストを決定します。

• トークンバジェット — 1回のフォワードパスで処理されるクエリトークンの最大数。ContinuousBatchingConfig の max_batch_tokens で設定されます。
• Cache budget — 1回のパスで読み取れるKVページの総数。これはキャッシュの総サイズによって上限が設定されます。
• リクエスト上限 — 1回のフォワードパスにおけるリクエストの最大数で、max_requests_per_batchによって設定されます。logitsテンソルは語彙サイズに比例して大きくなるため、この上限により大規模な語彙を持つモデルにおけるピークメモリが制限されます。

FIFO

FIFOSchedulerはデフォルトのスケジューラーです。優先度順にバッチを埋めます。まずアクティブなデコードリクエスト、次にアクティブなプリフィルリクエスト、最後に到着順で待機中の新規リクエストが処理されます。これにより、新規リクエストを取り込む前に既存のリクエストを処理しきることによってスループットを最大化します。

PrefillFirst

PrefillFirstScheduler は、デコードを再開する前にチャンク化されたプリフィル操作を完了します。これにより、長いプロンプトが中心のワークロードにおけるフラグメンテーションが軽減されます。

ContinuousBatchingConfigでスケジューラを設定します。

cb_config = ContinuousBatchingConfig(scheduler_type="prefill_first")

メモリ管理

KVキャッシュはページングされており、異なる長さのリクエストがフラグメンテーションなしで固定のメモリプールを共有できます。CachePoolとして実装されたメモリチャンクの共有プールが、各アテンションタイプのアロケーターにチャンクを供給し、アロケーターはそれらをブロックとページに分割します。

Shared pool (first two chunks reserved for padding / never allocated)

  chunk N  [ ======================================== ]
               |                            |
    full attn: [ FULL BLOCK 0 | FULL BLOCK 1 | … ]
               [ page L0 | page L0 | … ]   (one page per full layer)

    sliding:   [ ----------- SLIDING BLOCK 0 ----------- ]
               [ page L1 | page L2 | page L3 | … ]  (one page per sliding layer)

  page  = page_size tokens of KV for one layer
  block = one page per layer of that attention type
  chunk = pool unit sized so full and sliding blocks tile the same bytes

Continuous batching は 2 種類の attention 型のみをサポートします: full_attention (FullAttentionCacheAllocator) と sliding_attention (SlidingAttentionCacheAllocator)。ハイブリッドモデルには、それぞれの種類の allocator が 1 つずつ割り当てられます。attention 型は ContinuousBatchingConfig ではなく、モデル (layer_types / sliding_window) から取得されます。

Sliding layersは、小さなページセットを再利用することで、直近のトークンの固定ウィンドウのみを保持します。カーネルブロックテーブルやブロック共有はサポートされていません。

キャッシュは、リクエストが決して使用しないパディング読み書きのために、データチャンクの前方に2つのチャンクを確保します。これらのチャンクはリクエストに割り当てられることはありません。

キャッシュサイズ設定

マネージャーは、起動時に空き GPU メモリから収容可能なチャンク数を推定します。KV テンソル、アテンションマスク、アクティベーション、および管理用のインデックスを考慮した上で、利用可能なメモリの max_memory_percent 以内に入るようにプールのサイズを決定します。未設定の場合は、ハードウェアおよび連続バッチングのパラメータに基づいた値がデフォルトとして使用されます。MPS では、CPU と GPU がメモリを共有するため、スワッピングを回避するためにキャッシュに割り当てる領域を小さくします。num_blocks を明示的に設定した場合、page_size トークンからなるモデルブロックの整数倍として容量を表します。

ContinuousBatchingConfigの値を固定できます。

cb_config = ContinuousBatchingConfig(
    page_size=256,
    num_blocks=4096,
    max_batch_tokens=512,
    max_memory_percent=0.8,
)

受け入れ

リクエストがバッチに追加される前に、スケジューラーはすべてのアテンションタイプのアロケーターに対して十分な空き容量があることを確認します。いずれかのアロケーターで容量が不足する場合は、そのリクエストは拒否され、何も割り当てられません。

オフローディングが唯一の選択肢となるのを避けるため、スケジューラはsafety_marginを強制します。空きキャッシュ容量がsafety_marginを下回ると、新しいプリフィルは保留され、アクティブなデコードのみが継続されます。FIFO スケジューラのデフォルト値は0.15(容量の 15% をリザーブ)であり、prefill_firstのデフォルト値は0.0です。ContinuousBatchingConfigでsafety_marginを設定することで調整でき、0に設定するとマージンが無効になります。

プレフィックスキャッシュ

2つのリクエストがプロンプトの接頭辞を共有している場合、その接頭辞のKVを保持するブロックを共有できます。各完了したブロックは内容ハッシュが計算されます。一致する接頭辞を持つ後続のリクエストは、そのブロックを再利用し、該当するトークンのprefillをスキップします。共有ブロックは参照カウント管理され、それらを使用しているすべてのリクエストが完了した時点で初めてフリープールに返却されます。

allow_block_sharing のデフォルト値は True です。プレフィックス共有は、すべてのアロケーターが共有をサポートしている場合にのみ有効になるため、フルアテンションのみを使用するモデルでのみ利用されます。純粋なスライディングウィンドウモデルやハイブリッドモデルでは無効化されます。短いプロンプトと長い生成が特徴のワークロードでは、管理コストが節約効果を上回るため、allow_block_sharing=False を設定してください。

エビクション

解放されたブロックは消去されず、その内容とハッシュが保持されたまま「初期化済み」の状態を維持されるため、後続のプレフィックス一致によって再利用できます。未初期化プールが不足すると、マネージャーは最も最近の初期化済みブロックを遅延評価で未初期化に降格し、そのハッシュを破棄します。プレフィックスキャッシュは、アロケーターが新しいリクエストの処理を継続するために破棄できるベストエフォート層です。

ソフトリセット

ソフトリセットは、オフローディングのフォールバック経路です。マネージャーが、CPU オフローディングが無効であるか、プールが満杯であるため、リクエストの KV キャッシュを CPU スワッププールにコピーできない場合、代わりにそのリクエストをソフトリセットします。リクエストの生成トークンを元のプロンプトに追加し、そのキャッシュブロックを解放して、キューに再追加します。どのリクエストが選択されるかは、Offloadingセクションで説明しています。

キャッシュに空きが再びできた場合、リクエストはプロンプトにエンコードされた生成履歴に基づいて、中断した箇所から再開されます。

CUDA グラフ

すべての生成ステップには、バッチの組み立てから GPU カーネルのディスパッチ、結果の読み取りまで、CPU によるオーバーヘッドが伴います。CUDA グラフは、GPU 実行シーケンス全体を一度だけ記録し、形状が一致するバッチに対してそれを再生することで、このオーバーヘッドを排除します。

バッチの形状はステップごとに変わるため、continuous batching システムはパディングとキャッシュによってこれに対応しています。

1. クエリの長さは、q_padding_interval_sizeの最も近い倍数にパディングされます。
2. KV の長さは、kv_padding_interval_size の最も近い倍数にパディングされます。
3. 記録されたグラフはキャッシュに保存されます。

バッチのパディング済み形状がキャッシュされたグラフと一致する場合、CPUディスパッチなしでグラフがリプレイされます。新しい形状がトリガーされると、新しいグラフのキャプチャが実行されます。

CUDA graphs は静的な制御フローを必要とし、attention mask とは互換性がありません。デフォルトでは自動的に検出され、条件が満たされない場合は無効化されます。

非同期バッチ処理

Async batching は、2 つの I/O バッファペアと 2 つの CUDA ストリームを使用して、CPU と GPU の処理をオーバーラップさせます。

Sequential
  CPU  ── [prep N] ·········idle········· [prep N+1] ·········idle·········
  GPU  ·········idle········· [compute N] ·········idle········· [compute N+1]

Async
  CPU  ── [prep N] [prep N+1] [prep N+2] ──
  GPU  ── ········ [compute N] [compute N+1] ──

GPUがバッチNを計算している間、CPUはバッチN+1の準備を行います。ただし、この2つの処理をオーバーラップさせるには、VRAMが約2倍必要になります。

Async batching は CUDA graphs が有効である必要があります。グラフの再生が、ストリームオーバーラップを正しく機能させるために必要な安定したテンソルアドレスを提供するためです。

オフローディング

長い出力を生成するリクエストは、KVキャッシュに次のブロックを割り当てるための空き領域がなくなるまで増大する可能性があります。クラッシュを回避するため、マネージャーはアクティブなリクエストを待機キューに戻してオフロードし、残りのリクエストのためにブロックを解放します。

Offloading は需要駆動です。スケジューラーは、stalled 状態の要求、すなわち必要なブロックの割り当てに失敗したアクティブな要求と、それぞれの要求が何ブロックを必要としているかを報告します。マネージャーは、その需要を賄うのに十分な数のブロックを解放するために必要な最小限の要求をオフロードします。オフロードは新しいものから行われるため、現在のバッチにすでにスケジュールされた要求はキャッシュを保持し、最後に残るアクティブな要求はオフロードされません。

オフローディングがトリガーされるのは、2つの状況です。

• バッチがスケジュールされましたが、一部のアクティブなリクエストがスタックしています。マネージャーはスタックしたリクエストを十分オフロードし、次のバッチが残り分に対してブロックを割り当てられるようにします。
• キャッシュが満杯で、バッチがまったくスケジュールできません。マネージャーはリクエストをオフロードし、同じステップ内でループを回してスケジュールを再試行します。バッチが収まるか、オフロード可能なリクエストがなくなるまで続行します。回復のために次のフォワードパスを待たなくなります。

オフロードされたリクエストは、2つのパスのいずれかをたどります。cpu_offload_spaceが0.0より大きい場合、マネージャーはリクエストのKVキャッシュブロックを事前割り当てされたピン留め済みCPUバッファにコピーし、スワッププールに収まるすべてのリクエストを単一の転送にバッチ処理します。GPUのキャッシュ空間が利用可能になると、マネージャーはブロックをGPUにコピーし、プロンプトや生成済みトークンの再計算なしでリクエストを再開します。プールに収まらないリクエスト、またはCPUオフロードが無効な場合のすべてのオフロードリクエストは、ソフトリセットにフォールバックします。

次のステップ

• Continuous batching ブログ記事では、KV キャッシング、チャンク化プリフィル、動的スケジューリングについて、パフォーマンスベンチマークの数値とともに解説しています。
• 使用例については、Continuous batchingのドキュメントを参照してください。