ubatch と cudagraph:固定形状グラフで decode/prefill を加速する仕組み

源码版本v0.25.1

役割

cudagraph は CUDA が提供する「一度録いて何度でも再再生」する仕組みです:一連の kernel シーケンスをグラフとして録画し、以降の再再生で kernel launch オーバーヘッドを省きます。decode 段では各トークンにつきたった一つか二つの新トークンを計算し、kernel は小さく数が多いため、まさに cudagraph の出番です。v1 は CUDAGraphWrapper(cuda_graph.py:145-233)でモデル forward を包み、__call__ で初めてある BatchDescriptor に遭遇したときにグラフを録画し(cuda_graph.py:265-344)、以降同じ形状なら直接 entry.cudagraph.replay()(cuda_graph.py:357-361)します。

ubatch(マイクロバッチ、micro-batch)は v1 が DP>1 のとき一ステップの大バッチを N 個の固定形状小バッチに切り、各小バッチを別々に一枚のグラフに録画します。これで異なる DP rank が同時に走り、同時に TP/EP all2all 通信を行い、通信と計算をオーバーラップできます。このロジックは UBatchWrapper(gpu_ubatch_wrapper.py:113-211)と UBatchContext(ubatching.py:20-148)に実装されています。両者を合わせたものが DBO(Distributed Batch Overlap)です:マイクロバッチ分割 + cudagraph 録画 + マルチスレッドイベントセマフォによる協調。

cudagraph の実行モードは CUDAGraphMode(compilation.py:53-87)で決まります:NONE(録画しない)、PIECEWISE(分段録画、attention はグラフに入れない)、FULL(全段録画)、FULL_DECODE_ONLY(decode のみ full、prefill は piecewise)など。GPUModelRunner.load_model の末尾で設定に従い self.model に対応する wrapper を被せます(gpu_model_runner.py:5350-5379)。

設計動機

• cudagraph は形状駆動:CUDAGraphWrapper.__call__ は BatchDescriptor をキーにし(cuda_graph.py:257-263)、同じ形状は一度だけ録画し次回は直接 replay します。CUDAGraphDispatcher が runner 内で「どの形状にどのモードを録画するか」を一元管理し(gpu_model_runner.py:844-845)、_determine_batch_execution_and_padding が num_tokens を最も近い録画可能形状に pad します(gpu_model_runner.py:3879-3894)。
• PIECEWISE は attention をスキップ:full グラフでは attention kernel の形状が柔軟すぎ(可変 query/key)て録画しても次回役立たず、v1 は piecewise モードでモデルを複数段に切り、attention 段は eager で走らせ他の段をグラフに入れます。BreakableCUDAGraphWrapper(breakable_cudagraph.py:246)で実装します。
• ubatch は均等分割:maybe_create_ubatch_slices は num_tokens_padded // num_ubatches で分割点を算出し(ubatch_utils.py:63-114)、各 ubatch の形状を一致させます。これで同じ cudagraph を複数回 replay でき、異なる DP rank 間でも整列できます。
• マルチスレッド + CUDA Stream イベント協調:UBatchContext は threading.Barrier + cpu_wait_event/cpu_signal_event + gpu_comm_done_event/gpu_compute_done_event(ubatching.py:25-48)で CPU と GPU 両側で二つの ubatch スレッドを同期します:yield_and_switch_from_compute_to_comm が現在の stream を compute から comm に切り替え GPU compute の完了を待ち(ubatching.py:133-139)、逆も同様です。
• キャプチャ期は単独スレッドで CUDA context を初期化:_capture_ubatches は各 ubatch ごとにスレッドを立ち上げ、まず torch.cuda.current_blas_handle() でそのスレッド上に CUDA context を初期化してから UBatchContext の barrier に入ります(gpu_ubatch_wrapper.py:236-251)。さもないとキャプチャ時にメインスレッドの stream と ubatch スレッドの context が合いません。
• SM を通信と計算に分割:SMControlContextManager(gpu_ubatch_wrapper.py:68-112)は set_comm_sms/set_compute_sms で DeepEP all2all 通信に固定数の SM を予約し、MoE の大規模 all2all が GPU を占有し尽くして compute が SM を待てなくなるのを防ぎます。
• グラフプール共有:UBatchWrapper.graph_pool は CUDAGraphWrapper.graph_pool に転送され(gpu_ubatch_wrapper.py:143-147)、全形状のグラフは同一メモリプールを共有します。大グラフを先に採り小グラフは再利用します(gpu_model_runner.py:6662-6668)。

主要ファイル

• CUDAGraphMode:53-87 — 列挙型:NONE/PIECEWISE/FULL/FULL_DECODE_ONLY/FULL_AND_PIECEWISE、decode_mode/mixed_mode/has_full_cudagraphs などのヘルパーを提供。
• cudagraph_batch_sizes:729-736 — compilation_config.cudagraph_capture_sizes から録画可能形状を読み取りソート。
• cudagraph_dispatcher:844-845 — CUDagraphDispatcher(vllm_config) が dispatch と capture を一元管理。
• _determine_batch_execution_and_padding:3836-3948 — uniform_decode/has_lora/has_encoder_output を算出し、cudagraph_dispatcher.dispatch(num_tokens, ...) を呼んで mode を選び num_tokens を BatchDescriptor.num_tokens に pad。
• maybe_create_ubatch_slices:4197-4203 — should_ubatch=True のとき num_tokens_padded // num_ubatches で切り、(ubatch_slices, ubatch_slices_padded) の二つを返す。padded 版は attention 専用。
• load_model wrapper selection:5350-5379 — cudagraph_mode + use_ubatching に応じて self.model に BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper を被せる。
• capture_model:6647-6712 — 大形状優先でキャプチャ、set_cudagraph_capturing_enabled(True)、cudagraph_dispatcher.get_capture_descs() を進行、最後に lock_workspace()。
• UBatchSlice:13-29 — request_slice + token_slice の dataclass。is_empty() と num_tokens プロパティ。
• check_ubatch_thresholds:38-46 — use_ubatching が無効なら直接 False、uniform decode なら dbo_decode_token_threshold、それ以外は dbo_prefill_token_threshold。
• maybe_create_ubatch_slices:63-114 — cu_num_tokens で searchsorted して各 ubatch の request_slice を見つけ、最後に _pad_out_ubatch_slices で末尾を num_tokens_padded まで pad。
• UBatchContext:20-148 — CPU+GPU イベント協調器、yield_and_switch_from_compute_to_comm/yield_and_switch_from_comm_to_compute が核心(ubatching.py:133-147)。
• dbo_enabled / dbo_current_ubatch_id:150-157 — _THREAD_ID_TO_CONTEXT グローバル辞書で現在のスレッドがどの ubatch にいるかを探し、kernel はこれを使って現在 ubatch のバッファを選ぶ。
• UBatchWrapper.__init__:113-150 — comm_stream、ready_barrier(num_ubatches+1)、cudagraphs: dict[int, CUDAGraphMetaData] を構築し、runtime_mode に応じて CUDAGraphWrapper を作るか決める。
• _capture_ubatches:212-303 — ubatch スレッドを立ち上げ CUDA context を初期化、メインスレッドが torch.cuda.graph(...) で全スレッドを包み、キャプチャ後に self.cudagraphs[num_tokens] = cudagraph_metadata。
• _run_ubatches:305-341 — cudagraph を使わないフォールバックパス。純マルチスレッドでモデルを走らせる。
• UBatchWrapper.__call__:441-537 — forward_context から ubatch_slices/cudagraph_runtime_mode を取得。FULL モードかつ num_tokens in self.cudagraphs なら cudagraph_metadata.cudagraph.replay()(gpu_ubatch_wrapper.py:516-521)、さもなくば capture か _run_ubatches。
• CUDAGraphWrapper class:145-233 — concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry]、runtime_mode で FULL/PIECEWISE を区別。
• CUDAGraphWrapper.__call__:233-361 — 初めてある BatchDescriptor が来たとき torch.cuda.graph(...) で録画、以降は entry.cudagraph.replay()。replay 前に get_offloader().sync_prev_onload() で offloader のプリフェッチ完了を待つ。
• BreakableCUDAGraphWrapper:246 — PIECEWISE モードの実装。attention 段を切り出して eager で走らせる。

データフロー

runner の execute_model が _determine_batch_execution_and_padding まで進むと cudagraph_mode と batch_descriptor を選び、maybe_create_ubatch_slices でマイクロバッチを切り、最後に set_forward_context がこれらの値を forward context に詰め込み、self.model(...) を呼ぶと wrapper に入ります:

# vllm/v1/worker/gpu_model_runner.py L4197-L4203
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
    should_ubatch,
    num_scheduled_tokens_np,
    num_tokens_padded,
    num_reqs_padded,
    self.parallel_config.num_ubatches,
)

UBatchWrapper.__call__ に入ると、forward_context.ubatch_slices が None かを確認します。None なら単一グラフパス(または plain eager)に行き、さもなくば ubatch マルチスレッドパスに行きます:

# vllm/v1/worker/gpu_ubatch_wrapper.py L493-L521
if (
    num_tokens not in self.cudagraphs
    and cudagraph_runtime_mode is CUDAGraphMode.FULL
):
    ubatch_metadata = self._make_ubatch_metadata(
        ubatch_slices=ubatch_slices,
        attn_metadata=attn_metadata,
        slot_mapping=slot_mapping,
        input_ids=input_ids,
        positions=positions,
        intermediate_tensors=intermediate_tensors,
        inputs_embeds=inputs_embeds,
        compute_stream=compute_stream,
        dp_metadata=ubatch_dp_metadata,
        batch_descriptor=batch_descriptor,
        cudagraph_runtime_mode=CUDAGraphMode.NONE,
    )
    with self.sm_control:
        return self._capture_ubatches(ubatch_metadata, self.runnable)
elif (
    num_tokens in self.cudagraphs
    and cudagraph_runtime_mode is CUDAGraphMode.FULL
):
    cudagraph_metadata = self.cudagraphs[num_tokens]
    get_offloader().sync_prev_onload()
    cudagraph_metadata.cudagraph.replay()
    return cudagraph_metadata.outputs

_capture_ubatches は N 個のスレッドを立ち上げ、各スレッドが自身の UBatchContext に入ります。CPU 側は cpu_wait_event/cpu_signal_event で交互に一方を走らせ他方を待たせ、GPU 側は gpu_comm_done_event/gpu_compute_done_event で comm_stream と compute_stream 上で record + wait します(ubatching.py:82-92)。最終的にメインスレッドが torch.cuda.graph(...) コンテキストで全段を一枚のグラフに録画し self.cudagraphs[num_tokens] に格納します。以降同じ形状が来たら直接 replay() で Python オーバーヘッドなし。

境界と失敗

• 空バッチフォールバック:maybe_create_ubatch_slices は should_ubatch=False のとき直接 return None, None(ubatch_utils.py:71-72)し、UBatchWrapper.__call__ は ubatch_slices is None を見て単一グラフパスに進みます(gpu_ubatch_wrapper.py:447-464)。
• full モード形状衝突:ubatching を有効化した形状は ubatch グラフのみを録画し、ubatch なしの同じ形状は CUDAGraphWrapper による二重録画を避けるため、__call__ は if batch_descriptor.num_tokens in self.cudagraphs: cudagraph_runtime_mode = NONE と明示します(gpu_ubatch_wrapper.py:455-458)。
• DP 不一致は直接 assert:UBatchWrapper.__call__ は ubatch パス上で assert dp_metadata is not None(gpu_ubatch_wrapper.py:477-478)。ubatch 自体が DP>1 のために設計されており、単一 DP ではここに入りません。
• 末尾 ubatch が空になった:is_last_ubatch_empty は padded_num_tokens // num_ubatches * (num_ubatches-1) >= orig_num_tokens を判定し(ubatch_utils.py:32-35)、このとき最後の ubatch には実際のトークンがなくコードは skip する必要があります。
• キャプチャ期の offloader 同期:_capture_ubatches は capture 前に get_offloader().sync_prev_onload()(gpu_ubatch_wrapper.py:283-285)、capture 後に get_offloader().join_after_forward()(gpu_ubatch_wrapper.py:298-301)を呼びます。さもないと offloader のプリフェッチ stream が接続されず unjoined stream エラーになります。
• cudagraph は実行期の予期せぬ capture を許可しない:set_cudagraph_capturing_enabled(False) が capture_model の末尾で(gpu_model_runner.py:6689-6694)、validate_cudagraph_capturing_enabled が CUDAGraphWrapper capture 前に check し(cuda_graph.py:276-277)、実行期に予期せず capture パスに入るとエラーを投げます。
• DP マルチ rank 形状同期:coordinate_batch_across_dp が rank 間で should_ubatch と num_tokens_across_dp を決定し(gpu_model_runner.py:3907-3918)、全 rank が同じ ubatch 形状に切り、replay 時に all2all がずれないようにします。
• 入力アドレスの安定性:CUDAGraphWrapper は debug モードで assert new_input_addresses == entry.input_addresses(cuda_graph.py:346-355)。persistent_buffers は runner が用意し、さもないと replay は古いアドレスを読むことになります。

まとめ

cudagraph は v1 が decode と prefill を高速化する中核手段で、固定形状を一度録画して何度も replay します。CUDAGraphWrapper は BatchDescriptor をキーにし、CUDAGraphDispatcher が各バッチでどの mode を使うか決めます。ubatch は DP>1 のとき一バッチを N 個の等形状小バッチに切り、複数 DP rank 間で通信と計算をオーバーラップさせ、UBatchContext の CPU event + GPU stream event で二つのスレッド間を交互に切り替えます。両者を合わせて DBO となり、MoE の大規模デプロイではほぼ必須です。forward がどうこの層に至るかは /worker/gpu-model-runner を、worker が runner をどう呼ぶかは /worker/worker を、実行器レイヤーは /executor/executor を参照してください。

公式資料:vLLM ドキュメント · README