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。