[2/2][Model Runner V2] FULL CUDA graph capture for microbatched steps (DBO) (#51700)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Codex <noreply@openai.com>

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

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")
 

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
 

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()))

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

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
     )