From 73d2a8cf86c46e21af723ffd34ce0d56668fdd14 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:42:33 +0800 Subject: [PATCH] [2/2][Model Runner V2] FULL CUDA graph capture for microbatched steps (DBO) (#51700) Signed-off-by: specture724 Co-authored-by: Codex --- tests/v1/worker/test_gpu_ubatch_slicing.py | 163 ++++++++++++++++++++- vllm/config/vllm.py | 5 +- vllm/v1/worker/gpu/cudagraph_utils.py | 146 ++++++++++++++---- vllm/v1/worker/gpu/dp_utils.py | 46 ++++-- vllm/v1/worker/gpu/model_runner.py | 7 +- vllm/v1/worker/gpu/ubatch_utils.py | 120 +++++++++++---- 6 files changed, 407 insertions(+), 80 deletions(-) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index f54e746710..e735953d8b 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -20,13 +20,22 @@ import pytest import torch from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata -from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig +from vllm.config import ( + CompilationConfig, + CUDAGraphMode, + ModelConfig, + ParallelConfig, + VllmConfig, +) from vllm.forward_context import create_forward_context from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils -from vllm.v1.worker.gpu import dp_utils -from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu import cudagraph_utils, dp_utils +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, +) from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.ubatch_utils import ( @@ -272,6 +281,7 @@ def _sync_dp( num_tokens_per_rank: list[int], uniform_token_count_per_rank: list[int] | None = None, allow_ubatching: bool = True, + cudagraph_manager: Any = None, ) -> tuple[BatchExecutionDescriptor, dp_utils.DPSyncState | None]: """Run the DP handshake with the all-reduce stubbed out. @@ -295,7 +305,7 @@ def _sync_dp( patch.object(dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)), ): return dp_utils.sync_cudagraph_and_dp_padding( - cudagraph_manager=None, + cudagraph_manager=cudagraph_manager, desired_batch_desc=BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens_per_rank[0], @@ -718,3 +728,148 @@ def test_slicing_drops_stale_dcp_metadata_when_dcp_is_off(): ] assert all(ubatch.dcp_local_seq_lens is None for ubatch in ubatches) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU") +def test_capturable_run_replays_as_a_cudagraph(): + """Microbatched graph replay reads updated persistent input buffers.""" + vllm_config = VllmConfig( + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0), + parallel_config=ParallelConfig( + enable_dbo=True, all2all_backend="deepep_low_latency" + ), + ) + device = torch.device("cuda:0") + runner = _make_execution_runner(vllm_config) + + num_tokens = 16 + input_ids = torch.arange(num_tokens, device=device) + positions = torch.arange(num_tokens, device=device) + model_inputs = { + "input_ids": input_ids, + "positions": positions, + "inputs_embeds": None, + "intermediate_tensors": None, + } + ubatch_state = _make_ubatch_state( + vllm_config, + [ + UBatchSlice(slice(0, 4), slice(0, 8)), + UBatchSlice(slice(4, 8), slice(8, 16)), + ], + ) + + model = _YieldingModel([]) + graph = torch.cuda.CUDAGraph() + finish = runner.begin_capturable_run( + model, model_inputs, ubatch_state, for_capture=True + ) + with torch.cuda.graph(graph, stream=runner.capture_stream): + captured = finish() + + input_ids.fill_(3) + positions.fill_(5) + graph.replay() + torch.accelerator.synchronize() + + expected = torch.full( + (num_tokens, 1), 3 * 2 + 5, dtype=torch.float32, device=device + ) + torch.testing.assert_close(captured, expected) + + +def _request_slices(input_batch: InputBatch) -> list[slice]: + return [s.request_slice for s in create_ubatch_slices(input_batch, 2)] + + +def test_captured_split_survives_a_replay_with_enough_requests(): + """Replay preserves the captured split only if real requests reach it.""" + buffers = _make_buffers() + captured = _request_slices(InputBatch.make_dummy(16, 16, buffers)) + + def replay(num_reqs: int) -> list[slice]: + return _request_slices( + _make_input_batch( + [1] * num_reqs, + [128] * num_reqs, + buffers, + num_reqs_padded=16, + num_tokens_padded=16, + ) + ) + + assert replay(16) == captured + assert replay(9) == captured + assert replay(4) != captured + + +def _make_cudagraph_manager(capture_sizes: list[int]) -> CudaGraphManager: + vllm_config = VllmConfig( + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0), + parallel_config=ParallelConfig( + enable_dbo=True, + all2all_backend="deepep_low_latency", + dbo_decode_token_threshold=DECODE_THRESHOLD, + dbo_prefill_token_threshold=PREFILL_THRESHOLD, + ), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.FULL, + cudagraph_capture_sizes=capture_sizes, + ), + ) + with patch.object( + cudagraph_utils, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ): + manager = CudaGraphManager( + vllm_config, + torch.device("cuda:0"), + CUDAGraphMode.FULL, + decode_query_len=1, + ubatch_runner=cast(UBatchRunner, object()), + ) + manager._graphs_captured = True + return manager + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool") +def test_microbatched_graphs_are_only_offered_to_uniform_batches(): + """Mixed batches cannot reuse graphs captured with uniform query lengths.""" + manager = _make_cudagraph_manager([64, 128]) + + desc = manager.dispatch( + num_reqs=64, + num_tokens=64, + uniform_token_count=1, + num_active_loras=0, + num_ubatches=2, + ) + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_ubatches == 2 + + desc = manager.dispatch( + num_reqs=8, + num_tokens=64, + uniform_token_count=None, + num_active_loras=0, + num_ubatches=2, + ) + assert desc.cg_mode == CUDAGraphMode.NONE + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool") +def test_microbatched_graph_needs_every_rank_to_reach_the_split(): + """Fall back to eager on all ranks if any rank cannot reach the captured split.""" + manager = _make_cudagraph_manager([64, 128]) + uniform = [1, 1] + + desc, _ = _sync_dp([64, 64], uniform, cudagraph_manager=manager) + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_ubatches == 2 + + # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33. + desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager) + assert desc.cg_mode == CUDAGraphMode.NONE + assert desc.num_ubatches == 2 + assert dp_sync is not None and dp_sync.eager diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index fbe0213855..2fcc144af0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3034,8 +3034,9 @@ class VllmConfig: unsupported.append("dual batch overlap with multimodal models") if model_config is not None and model_config.is_hybrid: unsupported.append("dual batch overlap with hybrid models") - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: - unsupported.append("dual batch overlap with CUDA graphs") + if self.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE: + # DBO captures FULL graphs only. + unsupported.append("dual batch overlap with PIECEWISE CUDA graphs") if self.is_mm_encoder_only: unsupported.append("dual batch overlap with encoder only models") diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 12354f3c22..d197dbbc6d 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -4,7 +4,8 @@ import gc import itertools from collections import defaultdict from collections.abc import Callable, Iterable -from dataclasses import dataclass +from contextlib import ExitStack +from dataclasses import dataclass, replace from itertools import groupby, product from typing import TYPE_CHECKING, Any, NamedTuple, Protocol @@ -42,11 +43,13 @@ from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches if TYPE_CHECKING: from vllm.v1.worker.gpu.model_runner import GPUModelRunner from vllm.v1.worker.gpu.pcp_manager import PCPManager + from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner logger = init_logger(__name__) @@ -144,6 +147,7 @@ class CudaGraphManager: decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + ubatch_runner: "UBatchRunner | None" = None, ): self.vllm_config = vllm_config self.device = device @@ -153,6 +157,8 @@ class CudaGraphManager: self.cudagraph_mode = cudagraph_mode self.decode_query_len = decode_query_len self.varlen_decode = varlen_decode + # DBO supports FULL CUDA graphs only. + self.ubatch_runner = ubatch_runner self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -212,6 +218,33 @@ class CudaGraphManager: # Counts above the largest captured case clamp to it. return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case) + def _maybe_ubatch_twin( + self, desc: BatchExecutionDescriptor + ) -> BatchExecutionDescriptor | None: + """Return a microbatched capture candidate when eligible. + + Uniform query lengths preserve the captured request split. Use the DP + dispatch thresholds so all ranks generate the same candidates. + """ + if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL: + return None + if desc.num_reqs is None: + return None + uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs) + if remainder or desc.uniform_token_count not in (None, uniform_token_count): + return None + parallel_config = self.vllm_config.parallel_config + num_ubatches = get_num_ubatches(parallel_config) + if desc.num_tokens < num_ubatches: + return None + if not check_ubatch_thresholds( + parallel_config, desc.num_tokens, uniform_decode=True + ): + return None + return replace( + desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count + ) + def _init_candidates(self) -> None: """Build priority-ordered candidate lists for each token count.""" capture_sizes = self.compilation_config.cudagraph_capture_sizes @@ -300,6 +333,12 @@ class CudaGraphManager: if desc not in descs_by_mode[decode_mode]: descs_by_mode[decode_mode].append(desc) + ubatch_desc = self._maybe_ubatch_twin(desc) + if ubatch_desc is not None and ( + ubatch_desc not in descs_by_mode[decode_mode] + ): + descs_by_mode[decode_mode].append(ubatch_desc) + # recoverSSM cannot capture a dummy query wider than its workspace. if mixed_mode and ( not self.vllm_config.cache_config.use_kda_recoverssm @@ -321,6 +360,10 @@ class CudaGraphManager: ) descs_by_mode[mixed_mode].append(desc) + ubatch_desc = self._maybe_ubatch_twin(desc) + if ubatch_desc is not None: + descs_by_mode[mixed_mode].append(ubatch_desc) + for mode, descs in descs_by_mode.items(): descs.sort(key=lambda d: d.num_tokens, reverse=True) self._capture_descs[mode] = descs @@ -344,6 +387,13 @@ class CudaGraphManager: def needs_capture(self) -> bool: return len(self._capture_descs) > 0 + def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream: + """Capture on the stream used by the microbatch threads.""" + if desc.num_ubatches > 1: + assert self.ubatch_runner is not None + return self.ubatch_runner.capture_stream + return current_stream() + @torch.inference_mode() def capture( self, @@ -359,7 +409,10 @@ class CudaGraphManager: because attention backends may mutate or lazily initialize metadata during warmup. """ - with graph_capture(device=self.device): + with graph_capture(device=self.device), ExitStack() as stack: + if self.ubatch_runner is not None: + # Join parked threads on failure to avoid blocking later captures. + stack.callback(self.ubatch_runner.abort_pending_run) # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger # activations so FULL activations should fit in already allocated # buffers in the graph pool. @@ -387,7 +440,9 @@ class CudaGraphManager: # Capture logger.debug( - "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc + "CG Capture: mode=%s, batch_desc=%s", + desc.cg_mode.name, + desc, ) if ( desc.cg_mode == CUDAGraphMode.PIECEWISE @@ -415,7 +470,7 @@ class CudaGraphManager: torch.accelerator.synchronize() free_before = torch.accelerator.get_memory_info()[0] with torch.cuda.graph( - graph, self.pool, stream=current_stream() + graph, self.pool, stream=self._capture_stream(desc) ): forward_fn(CUDAGraphMode.NONE) # Join offloader's copy stream after forward to avoid @@ -429,6 +484,7 @@ class CudaGraphManager: self._capture_mem_samples.append(free_before - free_after) self.graphs[desc] = graph compilation_counter.num_cudagraph_captured += 1 + self._graphs_captured = True def captured_token_counts(self) -> list[int]: @@ -510,6 +566,7 @@ class ModelCudaGraphManager(CudaGraphManager): decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + ubatch_runner: "UBatchRunner | None" = None, ): super().__init__( vllm_config, @@ -518,6 +575,7 @@ class ModelCudaGraphManager(CudaGraphManager): decode_query_len, lora_capture_cases=lora_capture_cases, varlen_decode=varlen_decode, + ubatch_runner=ubatch_runner, ) self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] @@ -554,6 +612,35 @@ class ModelCudaGraphManager(CudaGraphManager): "Set VLLM_USE_BREAKABLE_CUDAGRAPH=1 or cudagraph_mode=NONE/FULL." ) + def store_capture_output(num_tokens: int, model_output: Any) -> None: + """Copy outputs to persistent buffers, allocating on first use.""" + if self.is_last_pp_rank: + # Last PP rank (common case). + if self.use_aux_hidden_state_outputs: + hidden_states, aux_hidden_states = model_output + else: + hidden_states = model_output + aux_hidden_states = [] + if self.hidden_states is None: + self.hidden_states = torch.empty_like(hidden_states) + self.hidden_states[:num_tokens] = hidden_states + if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: + self.aux_hidden_states = [ + torch.empty_like(x) for x in aux_hidden_states + ] + for i, aux in enumerate(aux_hidden_states): + self.aux_hidden_states[i][:num_tokens] = aux + else: + # Non-last PP rank. + assert isinstance(model_output, IntermediateTensors) + intermediate_tensors = model_output + if self.intermediate_tensors is None: + self.intermediate_tensors = IntermediateTensors.empty_like( + intermediate_tensors + ) + for k, v in intermediate_tensors.tensors.items(): + self.intermediate_tensors[k][:num_tokens] = v + def create_forward_fn( desc: BatchExecutionDescriptor, warmup: bool, @@ -583,6 +670,30 @@ class ModelCudaGraphManager(CudaGraphManager): assert intermediate_tensors is not None model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] + if desc.num_ubatches > 1: + # Prepare and park threads before capture; finish runs inside it. + assert self.ubatch_runner is not None + ubatch_state = self.ubatch_runner.prepare( + InputBatch.make_dummy(num_reqs, num_tokens, input_buffers), + block_tables.get_dummy_block_tables(num_reqs), + block_tables.get_dummy_slot_mappings(num_tokens), + cg_mode=CUDAGraphMode.FULL, + for_capture=True, + ) + # Capture with dummy rows marked as padding. + input_buffers.is_padding.fill_(True) + finish = self.ubatch_runner.begin_capturable_run( + model, model_inputs, ubatch_state, for_capture=True + ) + + def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None: + assert cg_mode != CUDAGraphMode.PIECEWISE, ( + "DBO does not support PIECEWISE cudagraphs" + ) + store_capture_output(num_tokens, finish()) + + return ubatch_forward_fn + attn_metadata, slot_mappings = prepare_inputs_to_capture( num_reqs, num_tokens, @@ -629,32 +740,7 @@ class ModelCudaGraphManager(CudaGraphManager): # model outputs. No need to keep track of the hidden states. return None - if self.is_last_pp_rank: - # Last PP rank (common case). - if self.use_aux_hidden_state_outputs: - hidden_states, aux_hidden_states = model_output - else: - hidden_states = model_output - aux_hidden_states = [] - if self.hidden_states is None: - self.hidden_states = torch.empty_like(hidden_states) - self.hidden_states[:num_tokens] = hidden_states - if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: - self.aux_hidden_states = [ - torch.empty_like(x) for x in aux_hidden_states - ] - for i, aux in enumerate(aux_hidden_states): - self.aux_hidden_states[i][:num_tokens] = aux - else: - # Non-last PP rank. - assert isinstance(model_output, IntermediateTensors) - intermediate_tensors = model_output - if self.intermediate_tensors is None: - self.intermediate_tensors = IntermediateTensors.empty_like( - intermediate_tensors - ) - for k, v in intermediate_tensors.tensors.items(): - self.intermediate_tensors[k][:num_tokens] = v + store_capture_output(num_tokens, model_output) return forward_fn diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index b7c73789de..80309bc009 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -112,23 +112,45 @@ def sync_cudagraph_and_dp_padding( int(num_tokens_across_dp.min()), uniform_decode=uniform_decode_across_dp, ): - # Microbatching is all-or-nothing: the expert all-to-all is - # collective, so every rank splits and pads to the same token count. - # A rank too small to fill a microbatch does no work in it, like a - # dummy run. Microbatched steps run eager; nothing is captured yet. + # Expert all-to-all requires every rank to split and pad equally. + # Empty microbatches run as dummy batches. ubatch_num_tokens = int(num_tokens_across_dp.max()) - return BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=ubatch_num_tokens, - num_reqs=num_reqs, - num_ubatches=get_num_ubatches(parallel_config), - ), DPSyncState( + num_ubatches = get_num_ubatches(parallel_config) + ubatch_desc = None + if cudagraph_manager is not None: + # Match a FULL microbatched CUDA graph, falling back to an eager + # descriptor if no graph matches. + ubatch_desc = cudagraph_manager.dispatch( + num_reqs, + ubatch_num_tokens, + synced_uniform_token_count, + num_active_loras=num_active_loras, + num_ubatches=num_ubatches, + ) + if 2 * int(num_tokens_across_dp.min()) < ubatch_desc.num_tokens: + # If one rank has an empty second microbatch, run without + # CUDA graphs. + ubatch_desc = None + if ubatch_desc is None: + ubatch_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=ubatch_num_tokens, + num_reqs=num_reqs, + num_ubatches=num_ubatches, + ) + # Refresh the token count to include CUDA graph padding. + ubatch_num_tokens = ubatch_desc.num_tokens + num_reqs = int(num_reqs_across_dp.max()) + if ubatch_desc.cg_mode == CUDAGraphMode.FULL: + assert ubatch_desc.num_reqs is not None + num_reqs = ubatch_desc.num_reqs + return ubatch_desc, DPSyncState( num_tokens_across_dp=torch.full_like( num_tokens_across_dp, ubatch_num_tokens ), uniform_token_count=synced_uniform_token_count, - eager=True, - num_reqs=int(num_reqs_across_dp.max()), + eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE, + num_reqs=num_reqs, ) synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item())) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9d7918126f..d66fc6c90f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -723,6 +723,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, varlen_decode=self.adaptive_verification is not None, + ubatch_runner=self.ubatch_runner, ) if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None: self.fast_prefill = FastPrefillHelper( @@ -1775,7 +1776,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): assert self.ubatch_runner is not None assert block_tables is not None and slot_mappings is not None ubatch_state = self.ubatch_runner.prepare( - input_batch, block_tables, slot_mappings + input_batch, + block_tables, + slot_mappings, + cg_mode=batch_desc.cg_mode, + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) elif not (dummy_run and skip_attn_for_dummy_run): assert slot_mappings is not None diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 8007e2dbb4..e3cade2d0b 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -3,6 +3,8 @@ """Microbatching (DBO) helpers for the V2 GPU model runner.""" import threading +from collections.abc import Callable +from contextlib import ExitStack from dataclasses import replace from typing import Any, NamedTuple @@ -317,8 +319,11 @@ class UBatchRunner: for _ in range(self.num_ubatches) ] self.comm_stream = torch.cuda.Stream(device=device) + # Threads and graph capture must use the same compute stream. + self.capture_stream = torch.cuda.Stream(device=device) # The microbatch threads plus the thread that starts them. self.ready_barrier = threading.Barrier(self.num_ubatches + 1) + self._pending_finish: Callable[[], Any] | None = None self.sm_control = create_sm_control_context(self.parallel_config) def prepare( @@ -326,11 +331,16 @@ class UBatchRunner: input_batch: InputBatch, block_tables: tuple[torch.Tensor, ...], slot_mappings: torch.Tensor, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, ) -> UBatchState: """Split the batch into the microbatches the step will run on. Attention metadata is built per microbatch and carried in the forward contexts the threads install, not in the caller's. + + FULL uses padded sizes; `for_capture` refreshes metadata from dummy + block tables. """ ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches) @@ -355,12 +365,13 @@ class UBatchRunner: attn_metadata.append( self.model_state.prepare_attn( ubatch, - CUDAGraphMode.NONE, + cg_mode, ubatch_block_tables, ubatch_slot_mappings, self.attn_groups, self.kv_cache_config, ubatch_idx=i, + for_capture=for_capture, ) ) slot_mappings_by_layer.append( @@ -415,14 +426,34 @@ class UBatchRunner: model_inputs: dict[str, Any], ubatch_state: UBatchState, ) -> Any: + finish = self.begin_capturable_run(model, model_inputs, ubatch_state) + return finish() + + def begin_capturable_run( + self, + model: Any, + model_inputs: dict[str, Any], + ubatch_state: UBatchState, + for_capture: bool = False, + ) -> Callable[[], Any]: + """Start microbatch threads and return a callback to run and join them. + + For capture, start threads outside the graph and call the callback inside. + `for_capture` selects the capture stream and initializes cuBLAS workspaces + before capture. + + Call the callback exactly once on this thread, or call `abort_pending_run` + on failure to release parked threads. + """ ubatch_slices = ubatch_state.slices assert len(ubatch_slices) == len(ubatch_state.forward_contexts) assert len(ubatch_slices) == self.num_ubatches + compute_stream = self.capture_stream if for_capture else current_stream() ubatch_contexts = make_ubatch_contexts( num_micro_batches=self.num_ubatches, comm_stream=self.comm_stream, - compute_stream=current_stream(), + compute_stream=compute_stream, forward_contexts=ubatch_state.forward_contexts, ready_barrier=self.ready_barrier, ) @@ -433,6 +464,13 @@ class UBatchRunner: @torch.inference_mode() def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None: try: + # New threads must select the worker's device. + torch.accelerator.set_device_index(self.device.index) + if for_capture: + # Allocate cuBLAS workspaces before capture. + for stream in (compute_stream, self.comm_stream): + with torch.cuda.stream(stream): + torch.cuda.current_blas_handle() with ubatch_context: outputs[ubatch_context.id] = model(**inputs) except BaseException as e: # noqa: BLE001 @@ -445,33 +483,56 @@ class UBatchRunner: # is deliberately out of scope here. errors[ubatch_context.id] = e - # The threads manage the forward context themselves; clear it here so - # it is restored correctly once they are done. - with override_forward_context(None), self.sm_control: - threads = [] - for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices): - thread = threading.Thread( - target=run_ubatch, - args=( - ubatch_context, - slice_model_inputs(model_inputs, ubatch_slice.token_slice), - ), - ) - threads.append(thread) - thread.start() + # Threads own the forward context; keep it cleared and SMs reserved + # until finish joins them. + stack = ExitStack() + stack.enter_context(override_forward_context(None)) + stack.enter_context(self.sm_control) - # Wait for every thread to reach its context, then start the first. - self.ready_barrier.wait() - ubatch_contexts[0].cpu_wait_event.set() - for thread in threads: - thread.join() + threads = [] + for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices): + thread = threading.Thread( + target=run_ubatch, + args=( + ubatch_context, + slice_model_inputs(model_inputs, ubatch_slice.token_slice), + ), + ) + threads.append(thread) + thread.start() - if errors: - failed = min(errors) - raise RuntimeError( - f"Microbatch {failed} of {self.num_ubatches} failed" - ) from errors[failed] - return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)]) + # Park all threads before the caller starts capture. + self.ready_barrier.wait() + + def finish() -> Any: + try: + ubatch_contexts[0].cpu_wait_event.set() + for thread in threads: + thread.join() + finally: + self._pending_finish = None + stack.close() + + if errors: + failed = min(errors) + raise RuntimeError( + f"Microbatch {failed} of {self.num_ubatches} failed" + ) from errors[failed] + return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)]) + + self._pending_finish = finish + return finish + + def abort_pending_run(self) -> None: + """Release and join parked threads after a failed capture.""" + if self._pending_finish is None: + return + try: + self._pending_finish() + except Exception: + logger.exception("Failed to join microbatch threads after an abort") + finally: + self._pending_finish = None def maybe_build_ubatch_runner( @@ -491,10 +552,7 @@ def maybe_build_ubatch_runner( if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1: return None - logger.info_once( - "Dual batch overlap is enabled. Microbatched steps run without " - "CUDA graphs on the V2 model runner." - ) + logger.info_once("Dual batch overlap is enabled.") return UBatchRunner( vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs )