ubatch 與 cudagraph:固定形狀圖怎麼加速 decode/prefill

源码版本v0.25.1

職責

cudagraph 是 CUDA 提供的「錄一次重放無數次」機制:把一段 kernel 序列錄成圖,以後每次 replay 省掉 kernel launch 開銷,decode 階段每個 token 只算一兩個新 token,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 時把一個 step 的大批切成 N 個固定形狀的小批,每個小批單獨錄一張圖,這樣不同 DP rank 可以同時跑、同時做 TP/EP all2all 通訊,通訊和計算 overlap 起來。這套邏輯在 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 作 key(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),錄了下次 replay 也用不上,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),反之亦然。
• capture 期單獨執行緒初始化 CUDA context:_capture_ubatches 給每個 ubatch 起一個執行緒,先 torch.cuda.current_blas_handle() 讓 CUDA context 在該執行緒上初始化,再進 UBatchContext 等 barrier(gpu_ubatch_wrapper.py:236-251),否則 capture 時主執行緒的 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 等 helper。
• 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 pad 到 BatchDescriptor.num_tokens。
• 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 把尾部 pad 到 num_tokens_padded。
• 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 的 buffer。
• 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(...) 包住 join 所有執行緒,捕獲後 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 不一致直接斷言: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 沒實際 token,代碼裡要 skip。
• capture 期 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 的 prefetch 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 才不會錯位。
• input 地址需穩定: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 做 key,CUDAGraphDispatcher 決定每批走哪種 mode。ubatch 是 DP>1 時把一批切成 N 個等形小批,讓多 DP rank 之間通訊和計算能 overlap,靠 UBatchContext 的 CPU event + GPU stream event 在兩個執行緒之間交替切換。兩者合起來就是 DBO,在 MoE 大規模部署時幾乎是必開的。前向怎麼走到這一層見 /worker/gpu-model-runner,worker 怎麼調 runner 見 /worker/worker,執行器層見 /executor/executor。

對照官方資料:vLLM 文件 · README