ubatch and cudagraph: how fixed-shape graphs accelerate decode/prefill

源码版本v0.25.1

Responsibilities

cudagraph is CUDA's "record once, replay countless times" mechanism: a kernel sequence is recorded into a graph, and every subsequent replay skips kernel launch overhead. In the decode phase each token only computes one or two new tokens, and the kernels are small and numerous — exactly cudagraph's sweet spot. v1 wraps the model forward with CUDAGraphWrapper(cuda_graph.py:145-233); __call__ records a graph the first time it sees a BatchDescriptor(cuda_graph.py:265-344), and for subsequent same-shape batches just calls entry.cudagraph.replay()(cuda_graph.py:357-361).

ubatch (micro-batch) is v1's mechanism for slicing a step's large batch into N fixed-shape micro-batches when DP>1; each micro-batch records its own graph so that different DP ranks can run concurrently, run TP/EP all2all simultaneously, and overlap communication with computation. This logic lives in UBatchWrapper(gpu_ubatch_wrapper.py:113-211) and UBatchContext(ubatching.py:20-148). Combined, these form DBO (Distributed Batch Overlap): micro-batch splitting + cudagraph recording + multi-threaded event-semaphore coordination.

cudagraph's runtime mode is decided by CUDAGraphMode(compilation.py:53-87): NONE(do not record), PIECEWISE(segmented recording, attention excluded from the graph), FULL(record the whole thing), FULL_DECODE_ONLY(record full only for decode, prefill goes piecewise), etc. At the end of GPUModelRunner.load_model, the matching wrapper is wrapped around self.model based on the config(gpu_model_runner.py:5350-5379).

Design motivation

• cudagraph is shape-driven: CUDAGraphWrapper.__call__ uses BatchDescriptor as the key(cuda_graph.py:257-263); a shape is recorded only once and replayed next time. CUDAGraphDispatcher centrally manages "which shape uses which mode" in the runner(gpu_model_runner.py:844-845), and _determine_batch_execution_and_padding pads num_tokens to the nearest recordable shape(gpu_model_runner.py:3879-3894).
• PIECEWISE skips attention: in a full graph, the attention kernel's shape is too flexible (variable query/key), so a recorded graph would be unusable on replay. v1's piecewise mode splits the model into segments — the attention segment runs eager, the rest goes into the graph — via BreakableCUDAGraphWrapper(breakable_cudagraph.py:246).
• ubatch sliced evenly: maybe_create_ubatch_slices splits at num_tokens_padded // num_ubatches(ubatch_utils.py:63-114); every ubatch has an identical shape, so the same cudagraph can be replayed multiple times and aligned across DP ranks.
• Multi-thread + CUDA Stream event coordination: UBatchContext uses threading.Barrier + cpu_wait_event/cpu_signal_event + gpu_comm_done_event/gpu_compute_done_event(ubatching.py:25-48) to synchronize the two ubatch threads on both CPU and GPU sides: yield_and_switch_from_compute_to_comm switches the current stream from compute to comm and waits for GPU compute to finish(ubatching.py:133-139), and vice versa.
• Capture initializes the CUDA context on a dedicated thread: _capture_ubatches spawns one thread per ubatch, first calling torch.cuda.current_blas_handle() so the CUDA context is initialized on that thread, then entering UBatchContext to wait at the barrier(gpu_ubatch_wrapper.py:236-251); otherwise the main thread's stream and the ubatch thread's context do not line up during capture.
• SM partition between communication and computation: SMControlContextManager(gpu_ubatch_wrapper.py:68-112) reserves a fixed number of SMs for DeepEP all2all via set_comm_sms/set_compute_sms, preventing large-scale MoE all2all from saturating the GPU and starving compute of SMs.
• Graph pool shared: UBatchWrapper.graph_pool forwards to CUDAGraphWrapper.graph_pool(gpu_ubatch_wrapper.py:143-147); all shape graphs share the same memory pool, with large graphs captured first and small ones reusing it(gpu_model_runner.py:6662-6668).

Key files

• CUDAGraphMode:53-87 — enum: NONE/PIECEWISE/FULL/FULL_DECODE_ONLY/FULL_AND_PIECEWISE; provides helpers such as decode_mode/mixed_mode/has_full_cudagraphs.
• cudagraph_batch_sizes:729-736 — reads recordable shapes from compilation_config.cudagraph_capture_sizes and sorts them.
• cudagraph_dispatcher:844-845 — CUDagraphDispatcher(vllm_config) centrally manages dispatch and capture.
• _determine_batch_execution_and_padding:3836-3948 — computes uniform_decode/has_lora/has_encoder_output, calls cudagraph_dispatcher.dispatch(num_tokens, ...) to pick a mode, and pads num_tokens to BatchDescriptor.num_tokens.
• maybe_create_ubatch_slices:4197-4203 — when should_ubatch=True, splits at num_tokens_padded // num_ubatches and returns two copies (ubatch_slices, ubatch_slices_padded); the padded version is used by attention specifically.
• load_model wrapper selection:5350-5379 — based on cudagraph_mode + use_ubatching, wraps self.model with BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper.
• capture_model:6647-6712 — captures large shapes first, set_cudagraph_capturing_enabled(True), walks cudagraph_dispatcher.get_capture_descs(), then lock_workspace().
• UBatchSlice:13-29 — dataclass of request_slice + token_slice, with is_empty() and num_tokens attributes.
• check_ubatch_thresholds:38-46 — if use_ubatching is off, returns False directly; uniform decode uses dbo_decode_token_threshold, otherwise dbo_prefill_token_threshold.
• maybe_create_ubatch_slices:63-114 — uses searchsorted on cu_num_tokens to find each ubatch's request_slice, then _pad_out_ubatch_slices pads the tail up to num_tokens_padded.
• UBatchContext:20-148 — CPU+GPU event coordinator; yield_and_switch_from_compute_to_comm/yield_and_switch_from_comm_to_compute are the core(ubatching.py:133-147).
• dbo_enabled / dbo_current_ubatch_id:150-157 — looks up which ubatch the current thread is in via the global _THREAD_ID_TO_CONTEXT dict; the code uses it to pick the current ubatch's buffer for kernels.
• UBatchWrapper.__init__:113-150 — creates comm_stream, ready_barrier(num_ubatches+1), cudagraphs: dict[int, CUDAGraphMetaData], and conditionally creates a CUDAGraphWrapper based on runtime_mode.
• _capture_ubatches:212-303 — spawns ubatch threads to initialize the CUDA context, the main thread wraps everything with torch.cuda.graph(...) and joins all threads; after capture, self.cudagraphs[num_tokens] = cudagraph_metadata.
• _run_ubatches:305-341 — fallback path without cudagraph, runs the model with plain multi-threading.
• UBatchWrapper.__call__:441-537 — forward_context retrieves ubatch_slices/cudagraph_runtime_mode; in FULL mode with num_tokens in self.cudagraphs, calls cudagraph_metadata.cudagraph.replay()(gpu_ubatch_wrapper.py:516-521), otherwise captures or falls back to _run_ubatches.
• CUDAGraphWrapper class:145-233 — concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry]; runtime_mode distinguishes FULL/PIECEWISE.
• CUDAGraphWrapper.__call__:233-361 — the first time a BatchDescriptor comes in, torch.cuda.graph(...) records; afterwards entry.cudagraph.replay(), and before replay get_offloader().sync_prev_onload() waits for the offloader's prefetch to finish.
• BreakableCUDAGraphWrapper:246 — PIECEWISE mode implementation; the attention segment is carved out and run eager.

Data flow

The runner's execute_model reaches _determine_batch_execution_and_padding to pick cudagraph_mode and batch_descriptor, then maybe_create_ubatch_slices slices the micro-batches, and finally set_forward_context stuffs these values into the forward context; once self.model(...) is called, it enters the 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,
)

When UBatchWrapper.__call__ is entered, it checks whether forward_context.ubatch_slices is None. If None, it goes the single-graph path (or plain eager); otherwise it goes the ubatch multi-thread path:

# 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 spawns N threads, each entering its own UBatchContext; the CPU side uses cpu_wait_event/cpu_signal_event to alternate between letting one run and the other wait, and the GPU side uses gpu_comm_done_event/gpu_compute_done_event to record + wait on the comm_stream and compute_stream(ubatching.py:82-92). The main thread then records the whole sequence as one graph inside the torch.cuda.graph(...) context and stores it in self.cudagraphs[num_tokens]. Subsequent same-shape batches just call replay(), with no Python overhead.

Boundaries and failure modes

• Empty batch fallback: maybe_create_ubatch_slices directly return None, None when should_ubatch=False(ubatch_utils.py:71-72); UBatchWrapper.__call__ takes the single-graph path when it sees ubatch_slices is None(gpu_ubatch_wrapper.py:447-464).
• Full-mode shape conflict: a shape with ubatching enabled only records the ubatch graph; the same shape without ubatch must avoid being captured twice by CUDAGraphWrapper, so __call__ explicitly does if batch_descriptor.num_tokens in self.cudagraphs: cudagraph_runtime_mode = NONE(gpu_ubatch_wrapper.py:455-458).
• DP mismatch asserts directly: UBatchWrapper.__call__ does assert dp_metadata is not None on the ubatch path(gpu_ubatch_wrapper.py:477-478); ubatch is designed for DP>1, so a single DP never reaches this path.
• Trailing ubatch is empty: is_last_ubatch_empty checks padded_num_tokens // num_ubatches * (num_ubatches-1) >= orig_num_tokens(ubatch_utils.py:32-35); in this case the last ubatch has no real tokens and is skipped in the code.
• Offloader sync during capture: _capture_ubatches calls get_offloader().sync_prev_onload() before capture(gpu_ubatch_wrapper.py:283-285) and get_offloader().join_after_forward() after(gpu_ubatch_wrapper.py:298-301); otherwise the offloader's prefetch stream is not connected and reports an unjoined stream.
• cudagraph forbids accidental runtime capture: set_cudagraph_capturing_enabled(False) at the end of capture_model(gpu_model_runner.py:6689-6694), and validate_cudagraph_capturing_enabled is checked before CUDAGraphWrapper capture(cuda_graph.py:276-277); accidentally entering the capture path at runtime raises an error.
• Multi-rank DP shape sync: coordinate_batch_across_dp decides should_ubatch and num_tokens_across_dp across ranks(gpu_model_runner.py:3907-3918); all ranks slice the same ubatch shapes so that replay's all2all does not misalign.
• Input addresses must be stable: CUDAGraphWrapper in debug mode does assert new_input_addresses == entry.input_addresses(cuda_graph.py:346-355); persistent_buffers are prepared by the runner, otherwise replay reads stale addresses.

Summary

cudagraph is v1's core mechanism for making decode and prefill fly: a fixed shape is recorded once and replayed many times; CUDAGraphWrapper uses BatchDescriptor as the key, and CUDAGraphDispatcher decides which mode each batch uses. ubatch is what, when DP>1, splits a batch into N equal-shape micro-batches so that multi-DP-rank communication and computation can overlap, alternating between two threads via UBatchContext's CPU event + GPU stream event. Combined they form DBO, which is almost mandatory for large-scale MoE deployments. For how the forward reaches this layer, see /worker/gpu-model-runner; for how the worker invokes the runner, see /worker/worker; for the executor layer, see /executor/executor.

See official docs: vLLM docs · README