@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 73d2a8cf86

[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>

Anjie Hou committed Sep 14, 2026 at 13:42 UTC 73d2a8cf86c46e21af723ffd34ce0d56668fdd14

6 files changed +407 -80

tests/v1/worker/test_gpu_ubatch_slicing.py

+159 -4

@@ -20,13 +20,22 @@ import pytest

20 import torch

22 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata

23 -from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig

23 +from vllm.config import (

24 + CompilationConfig,

25 + CUDAGraphMode,

26 + ModelConfig,

27 + ParallelConfig,

28 + VllmConfig,

29 +)

30 from vllm.forward_context import create_forward_context

31 from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens

32 from vllm.v1.kv_cache_interface import KVCacheConfig

33 from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils

28 -from vllm.v1.worker.gpu import dp_utils

29 -from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor

34 +from vllm.v1.worker.gpu import cudagraph_utils, dp_utils

35 +from vllm.v1.worker.gpu.cudagraph_utils import (

36 + BatchExecutionDescriptor,

37 + CudaGraphManager,

38 +)

39 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

40 from vllm.v1.worker.gpu.model_states.interface import ModelState

41 from vllm.v1.worker.gpu.ubatch_utils import (

@@ -272,6 +281,7 @@ def _sync_dp(

281 num_tokens_per_rank: list[int],

282 uniform_token_count_per_rank: list[int] | None = None,

283 allow_ubatching: bool = True,

284 + cudagraph_manager: Any = None,

285 ) -> tuple[BatchExecutionDescriptor, dp_utils.DPSyncState | None]:

286 """Run the DP handshake with the all-reduce stubbed out.

287

@@ -295,7 +305,7 @@ def _sync_dp(

305 patch.object(dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)),

306 ):

307 return dp_utils.sync_cudagraph_and_dp_padding(

298 - cudagraph_manager=None,

308 + cudagraph_manager=cudagraph_manager,

309 desired_batch_desc=BatchExecutionDescriptor(

310 cg_mode=CUDAGraphMode.NONE,

311 num_tokens=num_tokens_per_rank[0],

@@ -718,3 +728,148 @@ def test_slicing_drops_stale_dcp_metadata_when_dcp_is_off():

728 ]

729

730 assert all(ubatch.dcp_local_seq_lens is None for ubatch in ubatches)

731 +

732 +

733 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU")

734 +def test_capturable_run_replays_as_a_cudagraph():

735 + """Microbatched graph replay reads updated persistent input buffers."""

736 + vllm_config = VllmConfig(

737 + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),

738 + parallel_config=ParallelConfig(

739 + enable_dbo=True, all2all_backend="deepep_low_latency"

740 + ),

741 + )

742 + device = torch.device("cuda:0")

743 + runner = _make_execution_runner(vllm_config)

744 +

745 + num_tokens = 16

746 + input_ids = torch.arange(num_tokens, device=device)

747 + positions = torch.arange(num_tokens, device=device)

748 + model_inputs = {

749 + "input_ids": input_ids,

750 + "positions": positions,

751 + "inputs_embeds": None,

752 + "intermediate_tensors": None,

753 + }

754 + ubatch_state = _make_ubatch_state(

755 + vllm_config,

756 + [

757 + UBatchSlice(slice(0, 4), slice(0, 8)),

758 + UBatchSlice(slice(4, 8), slice(8, 16)),

759 + ],

760 + )

761 +

762 + model = _YieldingModel([])

763 + graph = torch.cuda.CUDAGraph()

764 + finish = runner.begin_capturable_run(

765 + model, model_inputs, ubatch_state, for_capture=True

766 + )

767 + with torch.cuda.graph(graph, stream=runner.capture_stream):

768 + captured = finish()

769 +

770 + input_ids.fill_(3)

771 + positions.fill_(5)

772 + graph.replay()

773 + torch.accelerator.synchronize()

774 +

775 + expected = torch.full(

776 + (num_tokens, 1), 3 * 2 + 5, dtype=torch.float32, device=device

777 + )

778 + torch.testing.assert_close(captured, expected)

779 +

780 +

781 +def _request_slices(input_batch: InputBatch) -> list[slice]:

782 + return [s.request_slice for s in create_ubatch_slices(input_batch, 2)]

783 +

784 +

785 +def test_captured_split_survives_a_replay_with_enough_requests():

786 + """Replay preserves the captured split only if real requests reach it."""

787 + buffers = _make_buffers()

788 + captured = _request_slices(InputBatch.make_dummy(16, 16, buffers))

789 +

790 + def replay(num_reqs: int) -> list[slice]:

791 + return _request_slices(

792 + _make_input_batch(

793 + [1] * num_reqs,

794 + [128] * num_reqs,

795 + buffers,

796 + num_reqs_padded=16,

797 + num_tokens_padded=16,

798 + )

799 + )

800 +

801 + assert replay(16) == captured

802 + assert replay(9) == captured

803 + assert replay(4) != captured

804 +

805 +

806 +def _make_cudagraph_manager(capture_sizes: list[int]) -> CudaGraphManager:

807 + vllm_config = VllmConfig(

808 + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),

809 + parallel_config=ParallelConfig(

810 + enable_dbo=True,

811 + all2all_backend="deepep_low_latency",

812 + dbo_decode_token_threshold=DECODE_THRESHOLD,

813 + dbo_prefill_token_threshold=PREFILL_THRESHOLD,

814 + ),

815 + compilation_config=CompilationConfig(

816 + cudagraph_mode=CUDAGraphMode.FULL,

817 + cudagraph_capture_sizes=capture_sizes,

818 + ),

819 + )

820 + with patch.object(

821 + cudagraph_utils,

822 + "get_pp_group",

823 + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),

824 + ):

825 + manager = CudaGraphManager(

826 + vllm_config,

827 + torch.device("cuda:0"),

828 + CUDAGraphMode.FULL,

829 + decode_query_len=1,

830 + ubatch_runner=cast(UBatchRunner, object()),

831 + )

832 + manager._graphs_captured = True

833 + return manager

834 +

835 +

836 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool")

837 +def test_microbatched_graphs_are_only_offered_to_uniform_batches():

838 + """Mixed batches cannot reuse graphs captured with uniform query lengths."""

839 + manager = _make_cudagraph_manager([64, 128])

840 +

841 + desc = manager.dispatch(

842 + num_reqs=64,

843 + num_tokens=64,

844 + uniform_token_count=1,

845 + num_active_loras=0,

846 + num_ubatches=2,

847 + )

848 + assert desc.cg_mode == CUDAGraphMode.FULL

849 + assert desc.num_ubatches == 2

850 +

851 + desc = manager.dispatch(

852 + num_reqs=8,

853 + num_tokens=64,

854 + uniform_token_count=None,

855 + num_active_loras=0,

856 + num_ubatches=2,

857 + )

858 + assert desc.cg_mode == CUDAGraphMode.NONE

859 +

860 +

861 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool")

862 +def test_microbatched_graph_needs_every_rank_to_reach_the_split():

863 + """Fall back to eager on all ranks if any rank cannot reach the captured split."""

864 + manager = _make_cudagraph_manager([64, 128])

865 + uniform = [1, 1]

866 +

867 + desc, _ = _sync_dp([64, 64], uniform, cudagraph_manager=manager)

868 + assert desc.cg_mode == CUDAGraphMode.FULL

869 + assert desc.num_ubatches == 2

870 +

871 + # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33.

872 + desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager)

873 + assert desc.cg_mode == CUDAGraphMode.NONE

874 + assert desc.num_ubatches == 2

875 + assert dp_sync is not None and dp_sync.eager

vllm/config/vllm.py

+3 -2

@@ -3034,8 +3034,9 @@ class VllmConfig:

3034 unsupported.append("dual batch overlap with multimodal models")

3035 if model_config is not None and model_config.is_hybrid:

3036 unsupported.append("dual batch overlap with hybrid models")

3037 - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:

3038 - unsupported.append("dual batch overlap with CUDA graphs")

3037 + if self.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE:

3038 + # DBO captures FULL graphs only.

3039 + unsupported.append("dual batch overlap with PIECEWISE CUDA graphs")

3040 if self.is_mm_encoder_only:

3041 unsupported.append("dual batch overlap with encoder only models")

3042

vllm/v1/worker/gpu/cudagraph_utils.py

+116 -30

@@ -4,7 +4,8 @@ import gc

4 import itertools

5 from collections import defaultdict

6 from collections.abc import Callable, Iterable

7 -from dataclasses import dataclass

7 +from contextlib import ExitStack

8 +from dataclasses import dataclass, replace

9 from itertools import groupby, product

10 from typing import TYPE_CHECKING, Any, NamedTuple, Protocol

@@ -42,11 +43,13 @@ from vllm.v1.worker.gpu.block_table import BlockTables

43 from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens

44 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

45 from vllm.v1.worker.gpu.model_states.interface import ModelState

46 +from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches

47 from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches

49 if TYPE_CHECKING:

50 from vllm.v1.worker.gpu.model_runner import GPUModelRunner

51 from vllm.v1.worker.gpu.pcp_manager import PCPManager

52 + from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner

54 logger = init_logger(__name__)

@@ -144,6 +147,7 @@ class CudaGraphManager:

147 decode_query_len: int,

148 lora_capture_cases: list[int] | None = None,

149 varlen_decode: bool = False,

150 + ubatch_runner: "UBatchRunner | None" = None,

151 ):

152 self.vllm_config = vllm_config

153 self.device = device

@@ -153,6 +157,8 @@ class CudaGraphManager:

157 self.cudagraph_mode = cudagraph_mode

158 self.decode_query_len = decode_query_len

159 self.varlen_decode = varlen_decode

160 + # DBO supports FULL CUDA graphs only.

161 + self.ubatch_runner = ubatch_runner

162

163 self.dp_size = vllm_config.parallel_config.data_parallel_size

164 self.tp_size = vllm_config.parallel_config.tensor_parallel_size

@@ -212,6 +218,33 @@ class CudaGraphManager:

218 # Counts above the largest captured case clamp to it.

219 return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case)

220

221 + def _maybe_ubatch_twin(

222 + self, desc: BatchExecutionDescriptor

223 + ) -> BatchExecutionDescriptor | None:

224 + """Return a microbatched capture candidate when eligible.

225 +

226 + Uniform query lengths preserve the captured request split. Use the DP

227 + dispatch thresholds so all ranks generate the same candidates.

228 + """

229 + if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL:

230 + return None

231 + if desc.num_reqs is None:

232 + return None

233 + uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs)

234 + if remainder or desc.uniform_token_count not in (None, uniform_token_count):

235 + return None

236 + parallel_config = self.vllm_config.parallel_config

237 + num_ubatches = get_num_ubatches(parallel_config)

238 + if desc.num_tokens < num_ubatches:

239 + return None

240 + if not check_ubatch_thresholds(

241 + parallel_config, desc.num_tokens, uniform_decode=True

242 + ):

243 + return None

244 + return replace(

245 + desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count

246 + )

247 +

248 def _init_candidates(self) -> None:

249 """Build priority-ordered candidate lists for each token count."""

250 capture_sizes = self.compilation_config.cudagraph_capture_sizes

@@ -300,6 +333,12 @@ class CudaGraphManager:

333 if desc not in descs_by_mode[decode_mode]:

334 descs_by_mode[decode_mode].append(desc)

335

336 + ubatch_desc = self._maybe_ubatch_twin(desc)

337 + if ubatch_desc is not None and (

338 + ubatch_desc not in descs_by_mode[decode_mode]

339 + ):

340 + descs_by_mode[decode_mode].append(ubatch_desc)

341 +

342 # recoverSSM cannot capture a dummy query wider than its workspace.

343 if mixed_mode and (

344 not self.vllm_config.cache_config.use_kda_recoverssm

@@ -321,6 +360,10 @@ class CudaGraphManager:

360 )

361 descs_by_mode[mixed_mode].append(desc)

362

363 + ubatch_desc = self._maybe_ubatch_twin(desc)

364 + if ubatch_desc is not None:

365 + descs_by_mode[mixed_mode].append(ubatch_desc)

366 +

367 for mode, descs in descs_by_mode.items():

368 descs.sort(key=lambda d: d.num_tokens, reverse=True)

369 self._capture_descs[mode] = descs

@@ -344,6 +387,13 @@ class CudaGraphManager:

387 def needs_capture(self) -> bool:

388 return len(self._capture_descs) > 0

389

390 + def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream:

391 + """Capture on the stream used by the microbatch threads."""

392 + if desc.num_ubatches > 1:

393 + assert self.ubatch_runner is not None

394 + return self.ubatch_runner.capture_stream

395 + return current_stream()

396 +

397 @torch.inference_mode()

398 def capture(

399 self,

@@ -359,7 +409,10 @@ class CudaGraphManager:

409 because attention backends may mutate or lazily initialize

410 metadata during warmup.

411 """

362 - with graph_capture(device=self.device):

412 + with graph_capture(device=self.device), ExitStack() as stack:

413 + if self.ubatch_runner is not None:

414 + # Join parked threads on failure to avoid blocking later captures.

415 + stack.callback(self.ubatch_runner.abort_pending_run)

416 # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger

417 # activations so FULL activations should fit in already allocated

418 # buffers in the graph pool.

@@ -387,7 +440,9 @@ class CudaGraphManager:

440

441 # Capture

442 logger.debug(

390 - "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc

443 + "CG Capture: mode=%s, batch_desc=%s",

444 + desc.cg_mode.name,

445 + desc,

446 )

447 if (

448 desc.cg_mode == CUDAGraphMode.PIECEWISE

@@ -415,7 +470,7 @@ class CudaGraphManager:

470 torch.accelerator.synchronize()

471 free_before = torch.accelerator.get_memory_info()[0]

472 with torch.cuda.graph(

418 - graph, self.pool, stream=current_stream()

473 + graph, self.pool, stream=self._capture_stream(desc)

474 ):

475 forward_fn(CUDAGraphMode.NONE)

476 # Join offloader's copy stream after forward to avoid

@@ -429,6 +484,7 @@ class CudaGraphManager:

484 self._capture_mem_samples.append(free_before - free_after)

485 self.graphs[desc] = graph

486 compilation_counter.num_cudagraph_captured += 1

487 +

488 self._graphs_captured = True

489

490 def captured_token_counts(self) -> list[int]:

@@ -510,6 +566,7 @@ class ModelCudaGraphManager(CudaGraphManager):

566 decode_query_len: int,

567 lora_capture_cases: list[int] | None = None,

568 varlen_decode: bool = False,

569 + ubatch_runner: "UBatchRunner | None" = None,

570 ):

571 super().__init__(

572 vllm_config,

@@ -518,6 +575,7 @@ class ModelCudaGraphManager(CudaGraphManager):

575 decode_query_len,

576 lora_capture_cases=lora_capture_cases,

577 varlen_decode=varlen_decode,

578 + ubatch_runner=ubatch_runner,

579 )

580 self.hidden_states: torch.Tensor | None = None

581 self.aux_hidden_states: list[torch.Tensor] = []

@@ -554,6 +612,35 @@ class ModelCudaGraphManager(CudaGraphManager):

612 "Set VLLM_USE_BREAKABLE_CUDAGRAPH=1 or cudagraph_mode=NONE/FULL."

613 )

614

615 + def store_capture_output(num_tokens: int, model_output: Any) -> None:

616 + """Copy outputs to persistent buffers, allocating on first use."""

617 + if self.is_last_pp_rank:

618 + # Last PP rank (common case).

619 + if self.use_aux_hidden_state_outputs:

620 + hidden_states, aux_hidden_states = model_output

621 + else:

622 + hidden_states = model_output

623 + aux_hidden_states = []

624 + if self.hidden_states is None:

625 + self.hidden_states = torch.empty_like(hidden_states)

626 + self.hidden_states[:num_tokens] = hidden_states

627 + if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:

628 + self.aux_hidden_states = [

629 + torch.empty_like(x) for x in aux_hidden_states

630 + ]

631 + for i, aux in enumerate(aux_hidden_states):

632 + self.aux_hidden_states[i][:num_tokens] = aux

633 + else:

634 + # Non-last PP rank.

635 + assert isinstance(model_output, IntermediateTensors)

636 + intermediate_tensors = model_output

637 + if self.intermediate_tensors is None:

638 + self.intermediate_tensors = IntermediateTensors.empty_like(

639 + intermediate_tensors

640 + )

641 + for k, v in intermediate_tensors.tensors.items():

642 + self.intermediate_tensors[k][:num_tokens] = v

643 +

644 def create_forward_fn(

645 desc: BatchExecutionDescriptor,

646 warmup: bool,

@@ -583,6 +670,30 @@ class ModelCudaGraphManager(CudaGraphManager):

670 assert intermediate_tensors is not None

671 model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens]

672

673 + if desc.num_ubatches > 1:

674 + # Prepare and park threads before capture; finish runs inside it.

675 + assert self.ubatch_runner is not None

676 + ubatch_state = self.ubatch_runner.prepare(

677 + InputBatch.make_dummy(num_reqs, num_tokens, input_buffers),

678 + block_tables.get_dummy_block_tables(num_reqs),

679 + block_tables.get_dummy_slot_mappings(num_tokens),

680 + cg_mode=CUDAGraphMode.FULL,

681 + for_capture=True,

682 + )

683 + # Capture with dummy rows marked as padding.

684 + input_buffers.is_padding.fill_(True)

685 + finish = self.ubatch_runner.begin_capturable_run(

686 + model, model_inputs, ubatch_state, for_capture=True

687 + )

688 +

689 + def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None:

690 + assert cg_mode != CUDAGraphMode.PIECEWISE, (

691 + "DBO does not support PIECEWISE cudagraphs"

692 + )

693 + store_capture_output(num_tokens, finish())

694 +

695 + return ubatch_forward_fn

696 +

697 attn_metadata, slot_mappings = prepare_inputs_to_capture(

698 num_reqs,

699 num_tokens,

@@ -629,32 +740,7 @@ class ModelCudaGraphManager(CudaGraphManager):

740 # model outputs. No need to keep track of the hidden states.

741 return None

742

632 - if self.is_last_pp_rank:

633 - # Last PP rank (common case).

634 - if self.use_aux_hidden_state_outputs:

635 - hidden_states, aux_hidden_states = model_output

636 - else:

637 - hidden_states = model_output

638 - aux_hidden_states = []

639 - if self.hidden_states is None:

640 - self.hidden_states = torch.empty_like(hidden_states)

641 - self.hidden_states[:num_tokens] = hidden_states

642 - if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:

643 - self.aux_hidden_states = [

644 - torch.empty_like(x) for x in aux_hidden_states

645 - ]

646 - for i, aux in enumerate(aux_hidden_states):

647 - self.aux_hidden_states[i][:num_tokens] = aux

648 - else:

649 - # Non-last PP rank.

650 - assert isinstance(model_output, IntermediateTensors)

651 - intermediate_tensors = model_output

652 - if self.intermediate_tensors is None:

653 - self.intermediate_tensors = IntermediateTensors.empty_like(

654 - intermediate_tensors

655 - )

656 - for k, v in intermediate_tensors.tensors.items():

657 - self.intermediate_tensors[k][:num_tokens] = v

743 + store_capture_output(num_tokens, model_output)

744

745 return forward_fn

746

vllm/v1/worker/gpu/dp_utils.py

+34 -12

@@ -112,23 +112,45 @@ def sync_cudagraph_and_dp_padding(

112 int(num_tokens_across_dp.min()),

113 uniform_decode=uniform_decode_across_dp,

114 ):

115 - # Microbatching is all-or-nothing: the expert all-to-all is

116 - # collective, so every rank splits and pads to the same token count.

117 - # A rank too small to fill a microbatch does no work in it, like a

118 - # dummy run. Microbatched steps run eager; nothing is captured yet.

115 + # Expert all-to-all requires every rank to split and pad equally.

116 + # Empty microbatches run as dummy batches.

117 ubatch_num_tokens = int(num_tokens_across_dp.max())

120 - return BatchExecutionDescriptor(

121 - cg_mode=CUDAGraphMode.NONE,

122 - num_tokens=ubatch_num_tokens,

123 - num_reqs=num_reqs,

124 - num_ubatches=get_num_ubatches(parallel_config),

125 - ), DPSyncState(

118 + num_ubatches = get_num_ubatches(parallel_config)

119 + ubatch_desc = None

120 + if cudagraph_manager is not None:

121 + # Match a FULL microbatched CUDA graph, falling back to an eager

122 + # descriptor if no graph matches.

123 + ubatch_desc = cudagraph_manager.dispatch(

124 + num_reqs,

125 + ubatch_num_tokens,

126 + synced_uniform_token_count,

127 + num_active_loras=num_active_loras,

128 + num_ubatches=num_ubatches,

129 + )

130 + if 2 * int(num_tokens_across_dp.min()) < ubatch_desc.num_tokens:

131 + # If one rank has an empty second microbatch, run without

132 + # CUDA graphs.

133 + ubatch_desc = None

134 + if ubatch_desc is None:

135 + ubatch_desc = BatchExecutionDescriptor(

136 + cg_mode=CUDAGraphMode.NONE,

137 + num_tokens=ubatch_num_tokens,

138 + num_reqs=num_reqs,

139 + num_ubatches=num_ubatches,

140 + )

141 + # Refresh the token count to include CUDA graph padding.

142 + ubatch_num_tokens = ubatch_desc.num_tokens

143 + num_reqs = int(num_reqs_across_dp.max())

144 + if ubatch_desc.cg_mode == CUDAGraphMode.FULL:

145 + assert ubatch_desc.num_reqs is not None

146 + num_reqs = ubatch_desc.num_reqs

147 + return ubatch_desc, DPSyncState(

148 num_tokens_across_dp=torch.full_like(

149 num_tokens_across_dp, ubatch_num_tokens

150 ),

151 uniform_token_count=synced_uniform_token_count,

130 - eager=True,

131 - num_reqs=int(num_reqs_across_dp.max()),

152 + eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE,

153 + num_reqs=num_reqs,

154 )

155

156 synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item()))

vllm/v1/worker/gpu/model_runner.py

+6 -1

@@ -723,6 +723,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

723 decode_query_len=self.decode_query_len,

724 lora_capture_cases=self.lora_capture_cases,

725 varlen_decode=self.adaptive_verification is not None,

726 + ubatch_runner=self.ubatch_runner,

727 )

728 if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None:

729 self.fast_prefill = FastPrefillHelper(

@@ -1775,7 +1776,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1776 assert self.ubatch_runner is not None

1777 assert block_tables is not None and slot_mappings is not None

1778 ubatch_state = self.ubatch_runner.prepare(

1778 - input_batch, block_tables, slot_mappings

1779 + input_batch,

1780 + block_tables,

1781 + slot_mappings,

1782 + cg_mode=batch_desc.cg_mode,

1783 + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,

1784 )

1785 elif not (dummy_run and skip_attn_for_dummy_run):

1786 assert slot_mappings is not None

vllm/v1/worker/gpu/ubatch_utils.py

+89 -31

@@ -3,6 +3,8 @@

3 """Microbatching (DBO) helpers for the V2 GPU model runner."""

5 import threading

6 +from collections.abc import Callable

7 +from contextlib import ExitStack

8 from dataclasses import replace

9 from typing import Any, NamedTuple

@@ -317,8 +319,11 @@ class UBatchRunner:

319 for _ in range(self.num_ubatches)

320 ]

321 self.comm_stream = torch.cuda.Stream(device=device)

322 + # Threads and graph capture must use the same compute stream.

323 + self.capture_stream = torch.cuda.Stream(device=device)

324 # The microbatch threads plus the thread that starts them.

325 self.ready_barrier = threading.Barrier(self.num_ubatches + 1)

326 + self._pending_finish: Callable[[], Any] | None = None

327 self.sm_control = create_sm_control_context(self.parallel_config)

328

329 def prepare(

@@ -326,11 +331,16 @@ class UBatchRunner:

331 input_batch: InputBatch,

332 block_tables: tuple[torch.Tensor, ...],

333 slot_mappings: torch.Tensor,

334 + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE,

335 + for_capture: bool = False,

336 ) -> UBatchState:

337 """Split the batch into the microbatches the step will run on.

338

339 Attention metadata is built per microbatch and carried in the forward

340 contexts the threads install, not in the caller's.

341 +

342 + FULL uses padded sizes; `for_capture` refreshes metadata from dummy

343 + block tables.

344 """

345 ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches)

346

@@ -355,12 +365,13 @@ class UBatchRunner:

365 attn_metadata.append(

366 self.model_state.prepare_attn(

367 ubatch,

358 - CUDAGraphMode.NONE,

368 + cg_mode,

369 ubatch_block_tables,

370 ubatch_slot_mappings,

371 self.attn_groups,

372 self.kv_cache_config,

373 ubatch_idx=i,

374 + for_capture=for_capture,

375 )

376 )

377 slot_mappings_by_layer.append(

@@ -415,14 +426,34 @@ class UBatchRunner:

426 model_inputs: dict[str, Any],

427 ubatch_state: UBatchState,

428 ) -> Any:

429 + finish = self.begin_capturable_run(model, model_inputs, ubatch_state)

430 + return finish()

431 +

432 + def begin_capturable_run(

433 + self,

434 + model: Any,

435 + model_inputs: dict[str, Any],

436 + ubatch_state: UBatchState,

437 + for_capture: bool = False,

438 + ) -> Callable[[], Any]:

439 + """Start microbatch threads and return a callback to run and join them.

440 +

441 + For capture, start threads outside the graph and call the callback inside.

442 + `for_capture` selects the capture stream and initializes cuBLAS workspaces

443 + before capture.

444 +

445 + Call the callback exactly once on this thread, or call `abort_pending_run`

446 + on failure to release parked threads.

447 + """

448 ubatch_slices = ubatch_state.slices

449 assert len(ubatch_slices) == len(ubatch_state.forward_contexts)

450 assert len(ubatch_slices) == self.num_ubatches

451

452 + compute_stream = self.capture_stream if for_capture else current_stream()

453 ubatch_contexts = make_ubatch_contexts(

454 num_micro_batches=self.num_ubatches,

455 comm_stream=self.comm_stream,

425 - compute_stream=current_stream(),

456 + compute_stream=compute_stream,

457 forward_contexts=ubatch_state.forward_contexts,

458 ready_barrier=self.ready_barrier,

459 )

@@ -433,6 +464,13 @@ class UBatchRunner:

464 @torch.inference_mode()

465 def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None:

466 try:

467 + # New threads must select the worker's device.

468 + torch.accelerator.set_device_index(self.device.index)

469 + if for_capture:

470 + # Allocate cuBLAS workspaces before capture.

471 + for stream in (compute_stream, self.comm_stream):

472 + with torch.cuda.stream(stream):

473 + torch.cuda.current_blas_handle()

474 with ubatch_context:

475 outputs[ubatch_context.id] = model(**inputs)

476 except BaseException as e: # noqa: BLE001

@@ -445,33 +483,56 @@ class UBatchRunner:

483 # is deliberately out of scope here.

484 errors[ubatch_context.id] = e

485

448 - # The threads manage the forward context themselves; clear it here so

449 - # it is restored correctly once they are done.

450 - with override_forward_context(None), self.sm_control:

451 - threads = []

452 - for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices):

453 - thread = threading.Thread(

454 - target=run_ubatch,

455 - args=(

456 - ubatch_context,

457 - slice_model_inputs(model_inputs, ubatch_slice.token_slice),

458 - ),

459 - )

460 - threads.append(thread)

461 - thread.start()

486 + # Threads own the forward context; keep it cleared and SMs reserved

487 + # until finish joins them.

488 + stack = ExitStack()

489 + stack.enter_context(override_forward_context(None))

490 + stack.enter_context(self.sm_control)

491 +

492 + threads = []

493 + for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices):

494 + thread = threading.Thread(

495 + target=run_ubatch,

496 + args=(

497 + ubatch_context,

498 + slice_model_inputs(model_inputs, ubatch_slice.token_slice),

499 + ),

500 + )

501 + threads.append(thread)

502 + thread.start()

503

463 - # Wait for every thread to reach its context, then start the first.

464 - self.ready_barrier.wait()

465 - ubatch_contexts[0].cpu_wait_event.set()

466 - for thread in threads:

467 - thread.join()

504 + # Park all threads before the caller starts capture.

505 + self.ready_barrier.wait()

506

469 - if errors:

470 - failed = min(errors)

471 - raise RuntimeError(

472 - f"Microbatch {failed} of {self.num_ubatches} failed"

473 - ) from errors[failed]

474 - return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)])

507 + def finish() -> Any:

508 + try:

509 + ubatch_contexts[0].cpu_wait_event.set()

510 + for thread in threads:

511 + thread.join()

512 + finally:

513 + self._pending_finish = None

514 + stack.close()

515 +

516 + if errors:

517 + failed = min(errors)

518 + raise RuntimeError(

519 + f"Microbatch {failed} of {self.num_ubatches} failed"

520 + ) from errors[failed]

521 + return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)])

522 +

523 + self._pending_finish = finish

524 + return finish

525 +

526 + def abort_pending_run(self) -> None:

527 + """Release and join parked threads after a failed capture."""

528 + if self._pending_finish is None:

529 + return

530 + try:

531 + self._pending_finish()

532 + except Exception:

533 + logger.exception("Failed to join microbatch threads after an abort")

534 + finally:

535 + self._pending_finish = None

536

537

538 def maybe_build_ubatch_runner(

@@ -491,10 +552,7 @@ def maybe_build_ubatch_runner(

552 if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1:

553 return None

554

494 - logger.info_once(

495 - "Dual batch overlap is enabled. Microbatched steps run without "

496 - "CUDA graphs on the V2 model runner."

497 - )

555 + logger.info_once("Dual batch overlap is enabled.")

556 return UBatchRunner(

557 vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs

558 )