連続バッチ処理における非同期性の解放

同期バッチ処理では GPU が 1/4 の時間アイドル状態になる

同期連続バッチ処理では CPU と GPU が交互に動作します。GPU がフォワードパスを計算している間は CPU が待機し、CPU が次のバッチを準備している間は GPU が待機します。バッチサイズ 32 で 8K トークンを生成する 8B モデルをプロファイルした結果、総実行時間は 300.6 秒で、そのうち 24.0% が CPU を待つ GPU のアイドル時間でした。このアイドルギャップは、連続バッチ処理だけでは解消できないスループット損失を表しています。

同時実行を実現するにはデフォルト以外の CUDA ストリームが必要

CPU と GPU の作業を重ね合わせるには、CPU をブロックせずに GPU の操作を起動する必要があります。PyTorch では明示的なストリームが指定されていない操作はデフォルトストリームで実行され、すべての他のストリームと同期して CPU が GPU の完了を待つことになります。デフォルト以外のストリームを使用すると、CPU は作業をキューイングした直後に制御を取り戻すことができます。ここでは、ホスト‑デバイス(H2D)転送用、GPU 計算用、デバイス‑ホスト(D2H)転送用の 3 つの独立したストリームが必要です。これらの操作は互いに独立しており、別々のストリームに配置すれば同時に実行できます。

ストリーム間の順序付けを強制するには CUDA イベントが必要

3 つのストリームを独立に起動するとレースコンディションが発生します。計算が H2D 転送完了前に開始したり、D2H が計算終了前に結果を転送したりする可能性があります。CUDA イベントはストリーム内にマーカーを記録し、別のストリームがそのマーカーを待機できる仕組みです。H2D 転送後にイベントを記録し、計算ストリームでそのイベントを待機させ、計算後に別のイベントを記録して D2H ストリームで待機させることで、GPU が正しい順序を保証しつつ CPU はすべての作業をキューイングした後は自由に動作できます。

データ破損を防ぎ、キャリーオーバーを処理するには二重バッファリングとマスクが必要

バッチ N の計算中にバッチ N+1 を準備すると、GPU がまだ読み取っている入力バッファが上書きされるリスクがあります。解決策は、ホスト側とデバイス側に 2 つのテンソルセット(スロット A と B)を用意し、交互に使用する二重バッファリングです。これによりメモリ使用量は倍になりますが、レースコンディションを防げます。また、あるリクエストが両バッチにまたがって出現する場合、次バッチの入力に新しく生成されたトークンが必要になるため、バッチ N+1 の入力バッファにプレースホルダー(値 0)を挿入し、後でバッチ N の出力トークンに置き換えます。この置き換え(キャリーオーバー)は、転送先位置を示すキャリーオーバーマスクによって行われ、CUDA グラフ内部でほぼコストなしに適用されます。

完全な非同期ループは CPU のバッチ準備と GPU 計算を重ね合わせる

ステップ 0(コールドスタート)では同期パスでバッチ 0 をディスパッチします。ステップ 1 以降は、GPU がスロット A でバッチ N を計算している間に CPU がスロット B でバッチ N+1 を準備します:完了したリクエストの除去、新規リクエストの受け入れ、KV キャッシュの更新、キャリーオーバーマスクの構築です。入力が整ったら CPU は H2D 転送を H2D ストリームにキューイングし、イベントを記録して計算ストリームで待機させ、フォワードパスを起動し、さらに別のイベントを記録して D2H ストリームで待機した後に出力転送を開始します。CPU は最終的な同期イベントで結果をコピーするまでブロックされ、その後すぐに次バッチの準備に入ります。このパターンが繰り返され、バッチ N が終了したときにバッチ N+1 の入力が準備できていれば GPU は常に稼働し続けます。

実験結果は 22% の速度向上を示す

同期ベースラインと同じワークロード(8K トークン、バッチサイズ 32、8B モデル)を非同期パイプラインで実行すると、CPU と GPU がほぼ連続して重なり合うタイムラインが得られます。GPU の稼働率は総実行時間の 99.4% に達し、同期ケースの 76.0% から大幅に向上しました。総生成時間は 300.6 秒から 234.5 秒へ短縮され、約 22% の速度向上となります。残りのギャップは CPU が D2H 転送完了を待つ不可避な同期ポイントに起因し、新たなカーネルやモデル変更は不要です。

結論

CUDA ストリームとイベントを用いてスケジュールベースの依存関係をデータベースの依存関係に置き換え、二重バッファリングとキャリーオーバーマスクでレースコンディションを解消することで、連続バッチ処理における CPU と GPU の作業を分離しました。これにより GPU が飽和状態を保ち、LLM 推論のスループットが大幅に向上しつつ、モデル精度は維持されます。

Sources