記事のサマリー(TL;DR)
- 同期 Continuous Batching では GPU アイドル時間が全体の 24% を占め、H200($5/時)が 1 日で $29 分無駄になる計算
- CUDA ストリーム・CUDA イベント・ダブルバッファリングを組み合わせ、CPU と GPU を完全並走させる非同期ループを構築
- 8B モデル・8K トークン生成で GPU 稼働率 76.0% → 99.4%、生成時間 300.6 秒 → 234.5 秒(22% 短縮)を達成
国内 LLM 推論基盤・クラウド GPU 運用担当者が押さえるべき実装ポイント
本手法はモデル変更もカスタムカーネルも不要で、既存の transformers ライブラリ に実装済みである点が実用上の最大の特徴です。AWS・GCP で H100/H200 インスタンスを時間課金で運用している場合、同じ GPU コストで約 22% 多くのトークンを処理できることになります。
日本国内でも、強化学習(RLHF)や長文一括処理など 16K トークン超の生成ワークロードを抱えるケースは増えており、本手法が特に有効です。実装の入口は continuous_batching.py、非同期処理の中心は ContinuousBatchingAsyncIOs クラスで、OSSとして参照・流用が可能です。
生成AI 基盤の推論コストを下げたい場面で、モデルの差し替えより先にインフラ側の同期/非同期設計を見直すアプローチとして参考になります。
詳細
同期 Continuous Batching の問題点
同期方式では CPU と GPU が交互に動作します。
- CPU がバッチを準備(リクエスト選択・KV キャッシュ更新・入力テンソル構築)
- CPU → GPU へ入力データを転送
- GPU がフォワードパスとサンプリングを実行
- 結果が CPU に返却され、サイクル繰り返し
この構造では CPU が動いている間 GPU はアイドル、GPU が動いている間 CPU はアイドルとなり、両者が同時に有用な処理を行うことがない。
8B モデル・バッチサイズ 32・8K トークン生成のプロファイルでは、合計生成時間 300.6 秒のうち 24.0%(約 72 秒)が GPU アイドル。仮に CPU オーバーヘッドをゼロにできれば 300 秒 → 228 秒(24% 高速化)が理論上限となります。
非同期化の設計思想
解決すべき技術的課題は以下の 3 点です。
- GPU にジョブを投入しながら CPU に制御を返す方法
- 各タスク開始時にデータが準備完了していることを保証する方法
- バッチ N の予測結果を使うバッチ N+1 をどう準備するか
CUDA ストリームとは何か
CUDA ストリームは GPU 操作(カーネル起動・メモリコピー・同期バリア)を順序付けるキューです。
- 同一ストリーム内:逐次実行(前の操作が終わるまで次は開始しない)
- 異なるストリーム間:互いに独立し、並列実行可能
デフォルトストリームと非デフォルトストリーム
PyTorch でストリームを明示しない操作はすべてデフォルトストリームに乗ります。デフォルトストリームは同期的で、「他のすべてのストリームがフラッシュされるまで待機」「他のすべての操作はデフォルトストリームがフラッシュされるまで待機」という双方向の同期を強制します。
並列処理を実現するには非デフォルトストリームを使い、カーネル起動やノンブロッキングメモリコピーが CPU に即座に制御を返すようにする必要があります。
Continuous Batching での 3 ストリーム構成
GPU 操作は以下の 3 種類に分類されます。
| 操作 | ストリーム名 |
|---|---|
| CPU → GPU 入力転送(Host-to-Device) | H2D ストリーム |
| GPU フォワードパス | compute ストリーム |
| GPU → CPU 出力転送(Device-to-Host) | D2H ストリーム |
転送は互いに独立しているため、それぞれに専用ストリームを割り当てます。
CUDA イベントによる順序保証
ストリームを独立させるだけでは「H2D 転送完了前に compute が始まる」などの不正実行が発生します。これを防ぐのが CUDA イベントです。
stream.record(event) # ストリームにマーカーを挿入
stream.wait(event) # イベント完了までストリームをブロック
wait はストリームを止めるのみで CPU や他のストリームはブロックしない点が重要です。
実際の依存関係の設定例:
- H2D 転送完了 →
h2d_doneイベントを record - compute ストリームが
h2d_doneを wait してからフォワードパス開始 - フォワードパス完了 →
compute_doneイベントを record - D2H ストリームが
compute_doneを wait してから出力転送開始 - D2H 転送完了 → CPU が
d2h_done_event.synchronize()で結果を受け取る
CPU がブロックするのはステップ 5 の synchronize() のみです。
レースコンディションの回避:ダブルバッファリング
バッチ N+1 の入力準備中にバッチ N がまだ同じメモリを読んでいるとデータ破損が発生します。解決策は2 つのスロット(スロット A・スロット B)を交互に使うダブルバッファリングです。
- GPU がスロット A でバッチ N を処理している間、CPU はスロット B でバッチ N+1 を準備
- 次のステップでスロットを入れ替え
トレードオフとして RAM と VRAM の使用量が 2 倍になりますが、FlashAttention を使う場合はアテンションマスクが不要なため、最大テンソルのサイズが大幅に削減され現実的なコストに収まります。
CUDA グラフとメモリプール
CUDA グラフは特定のメモリアドレスに対してキャプチャされるため、スロット A 用と B 用で 2 つのグラフが必要になります。メモリプール(shared memory buffer)を使うことで、2 つのグラフが同じプールから確保し合計 VRAM 使用量を 1 グラフ分に抑えられます(2 グラフは常に並列実行しないため安全)。
キャリーオーバー:バッチ N の出力トークンをバッチ N+1 の入力へ
同一リクエストがバッチ N にも N+1 にも含まれる場合、バッチ N で生成されたトークンをバッチ N+1 の入力に反映する必要があります。しかしバッチ N+1 を準備する時点ではまだそのトークンは存在しません。
解決策はプレースホルダー(値 0)を入力に入れておき、バッチ N の演算完了後・バッチ N+1 のフォワードパス開始前にキャリーオーバーマスクで差し替える 4 ステップです。
- バッチ N の出力から対象トークンをテンソル T に選択
- T のうちキャリーオーバー不要な位置をゼロクリア
- T をバッチ N+1 の入力長にトランケート
- T をバッチ N+1 の入力 ID に加算(プレースホルダーが 0 なので安全)
この操作はコストが低いため CUDA グラフ内にキャプチャされます。
非同期ループの全体像
ステップ 0(コールドスタート):前バッチが存在しないため、同期バッチングと同様にバッチ 0 をスロット A で起動。
ステップ 1 以降(非同期ループ):
- GPU がスロット A でバッチ 0 を実行中に、CPU がスロット B でバッチ 1 を並列準備
- バッチ 1 の入力が準備できたら H2D 転送をキューに投入し、イベントで compute・D2H に依存関係を設定
- スロット A(バッチ 0 の compute 完了 → D2H)とスロット B(バッチ 1 の H2D → compute)が並列進行
- CPU は
d2h_done_event.synchronize()でバッチ 0 の出力を受け取り次第、バッチ 2 の準備を開始
バッチ N+1 の入力が GPU 上に揃ってさえいれば、バッチ間で GPU がアイドルになることはありません。
実測結果
| 指標 | 同期方式 | 非同期方式 |
|---|---|---|
| 合計生成時間 | 300.6 秒 | 234.5 秒 |
| GPU 稼働率 | 76.0% | 99.4% |
| 改善幅 | — | 22% 高速化 |
理論上限(CPU オーバーヘッド完全排除)は 24% 高速化。残り 2% の差は不可避な同期ポイント(synchronize() 呼び出し)によるものです。
新しいカーネルもモデル変更も不要で、CPU と GPU を同時に働かせるだけで達成した結果です。
実装と今後の展望
完全な実装は transformers ライブラリの continuous_batching.py に統合済みです。非同期処理のコアは ContinuousBatchingAsyncIOs クラスにあります。
次回の記事では、SOTA スループット達成に向けた残課題(リクエストのオフロード・decode 専用カーネル・細粒度コンパイルなど)を取り上げる予定です。