# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """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 import numpy as np import torch from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed import get_dcp_group from vllm.forward_context import ( DPMetadata, ForwardContext, create_forward_context, override_forward_context, ) from vllm.logger import init_logger from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import current_stream from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.ubatch_utils import ( UBatchSlice, UBatchSlices, create_sm_control_context, maybe_create_ubatch_slices, ) from vllm.v1.worker.ubatching import make_ubatch_contexts from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) class UBatchState(NamedTuple): """Per-microbatch inputs for one dual-batch-overlap step.""" slices: UBatchSlices forward_contexts: list[ForwardContext] def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: """Split a DP-padded batch into the slices its microbatches run on. Splitting at the midpoint of the DP-padded token count leaves the trailing microbatch anywhere from full to entirely inside the padding. Clamping its request slice onto the last request leaves it holding that request with zero query tokens -- the shape a straddling request already takes -- so it stays well-formed with no work to do, like a dummy run. """ _, ubatch_slices = maybe_create_ubatch_slices( True, input_batch.num_scheduled_tokens, input_batch.num_tokens_after_padding, input_batch.num_reqs_after_padding, num_ubatches, ) assert ubatch_slices is not None return [ UBatchSlice( slice( min(s.request_slice.start, input_batch.num_reqs - 1), min(s.request_slice.stop, input_batch.num_reqs_after_padding), ), s.token_slice, ) for s in ubatch_slices ] def _slice_input_batch( input_batch: InputBatch, ubatch_slice: UBatchSlice, query_start_loc_buf: torch.Tensor, seq_lens_buf: torch.Tensor, dcp_local_seq_lens_buf: torch.Tensor | None = None, *, dcp_size: int = 1, dcp_rank: int = 0, cp_interleave: int = 1, ) -> InputBatch: """Build the sub-`InputBatch` a single microbatch runs on. A request straddling the boundary appears in both microbatches, its query truncated to the tokens each one owns. Query offsets and sequence lengths (including DCP-local lengths) are written to the caller's persistent buffers after truncation. Other GPU inputs are views of the full batch's buffers. The sub-batch describes the forward pass only. Sampling runs once over the merged batch, so `logits_indices`, `cu_num_logits` and the draft-token fields are carried over unchanged and must not be read off a sub-batch. """ assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty" req_start = ubatch_slice.request_slice.start req_stop = ubatch_slice.request_slice.stop tok_start = ubatch_slice.token_slice.start tok_stop = ubatch_slice.token_slice.stop num_reqs_padded = req_stop - req_start num_tokens_padded = tok_stop - tok_start # The trailing microbatch may extend past the real batch into CUDA graph # padding; the unpadded counts stop at the last real request/token. num_reqs = max(0, min(req_stop, input_batch.num_reqs) - req_start) num_tokens = max(0, min(tok_stop, input_batch.num_tokens) - tok_start) query_start_loc = _slice_query_start_loc( input_batch.query_start_loc, query_start_loc_buf[: num_reqs_padded + 1], req_start, req_stop, tok_start, num_tokens_padded, ) query_start_loc_np = np.clip( input_batch.query_start_loc_np[req_start : req_stop + 1] - tok_start, 0, num_tokens_padded, ).astype(np.int32) seq_lens = _slice_seq_lens( input_batch.seq_lens, seq_lens_buf[:num_reqs_padded], input_batch.query_start_loc, req_start, req_stop, tok_stop, ) # Same truncation as `_slice_seq_lens`, on the host-side upper bound. seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[ req_start:req_stop ].clone() tokens_truncated = max(0, int(input_batch.query_start_loc_np[req_stop]) - tok_stop) if tokens_truncated: seq_lens_cpu_upper_bound[-1] -= tokens_truncated # Query lengths of the truncated requests, so consumers that derive # max_query_len from this array see the microbatch's own lengths. num_scheduled_tokens = np.diff(query_start_loc_np)[:num_reqs] dcp_local_seq_lens = None if dcp_size > 1: assert dcp_local_seq_lens_buf is not None dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens( dcp_local_seq_lens_buf, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave, num_reqs_padded=num_reqs_padded, ) return replace( input_batch, req_ids=input_batch.req_ids[req_start : min(req_stop, input_batch.num_reqs)], num_reqs=num_reqs, num_reqs_after_padding=num_reqs_padded, idx_mapping=input_batch.idx_mapping[req_start:req_stop], idx_mapping_np=input_batch.idx_mapping_np[req_start:req_stop], num_scheduled_tokens=num_scheduled_tokens, num_tokens=num_tokens, num_tokens_after_padding=num_tokens_padded, query_start_loc=query_start_loc, query_start_loc_np=query_start_loc_np, seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=dcp_local_seq_lens, num_computed_tokens_np=input_batch.num_computed_tokens_np[req_start:req_stop], prefill_len_np=input_batch.prefill_len_np[req_start:req_stop], num_computed_prefill_tokens_np=input_batch.num_computed_prefill_tokens_np[ req_start:req_stop ], is_prefilling_np=input_batch.is_prefilling_np[req_start:req_stop], input_ids=input_batch.input_ids[tok_start:tok_stop], positions=input_batch.positions[tok_start:tok_stop], is_padding=input_batch.is_padding[tok_start:tok_stop], prompt_lens=( None if input_batch.prompt_lens is None else input_batch.prompt_lens[req_start:req_stop] ), fast_prefill=None, ) def _slice_query_start_loc( query_start_loc: torch.Tensor, out: torch.Tensor, req_start: int, req_stop: int, tok_start: int, num_tokens_padded: int, ) -> torch.Tensor: """Rebase query_start_loc onto the microbatch's token range. The clamp truncates the requests straddling either boundary to the tokens this microbatch owns. """ torch.sub(query_start_loc[req_start : req_stop + 1], tok_start, out=out) return out.clamp_(0, num_tokens_padded) def _slice_seq_lens( seq_lens: torch.Tensor, out: torch.Tensor, query_start_loc: torch.Tensor, req_start: int, req_stop: int, tok_stop: int, ) -> torch.Tensor: """Slice seq_lens, shortening the request truncated at the boundary. Only tokens past the end of this microbatch are missing; tokens it owns in an *earlier* microbatch are computed before this one reads them, so they still count. Deriving the truncation from this microbatch's own query lengths would subtract them twice. It is computed from tensors, not Python ints, so this stays valid inside a CUDA graph. """ out.copy_(seq_lens[req_start:req_stop]) last = req_stop - req_start - 1 out[last] -= (query_start_loc[req_stop] - tok_stop).clamp_(min=0) return out def slice_model_inputs( model_inputs: dict[str, Any], token_slice: slice ) -> dict[str, Any]: """Narrow the model's per-token inputs to one microbatch.""" sliced = dict(model_inputs) for key in ("input_ids", "inputs_embeds"): value = model_inputs.get(key) if value is not None: sliced[key] = value[token_slice] positions = model_inputs["positions"] # M-RoPE carries a leading section dim. sliced["positions"] = ( positions[:, token_slice] if positions.ndim == 2 else positions[token_slice] ) intermediate_tensors = model_inputs.get("intermediate_tensors") if intermediate_tensors is not None: sliced["intermediate_tensors"] = intermediate_tensors[token_slice] return sliced def merge_ubatch_outputs(outputs: list[Any]) -> Any: """Reassemble the full-batch model output from the per-microbatch ones. Models return hidden states, a tuple of those plus auxiliary hidden states (EAGLE3), or IntermediateTensors on non-final pipeline ranks. """ first = outputs[0] if isinstance(first, IntermediateTensors): return IntermediateTensors( { key: torch.cat([out.tensors[key] for out in outputs], dim=0) for key in first.tensors } ) if isinstance(first, tuple): return tuple(torch.cat(parts, dim=0) for parts in zip(*outputs)) return torch.cat(outputs, dim=0) class UBatchRunner: """Runs the model on two or more microbatches that overlap each other. Each microbatch gets a thread and its own forward context. Only one thread holds the GPU at a time; they hand off to each other at the communication points inside the model (see `vllm.v1.worker.ubatching`), so one microbatch's expert all-to-all overlaps the other's compute. """ def __init__( self, vllm_config: VllmConfig, device: torch.device, model_state: ModelState, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, max_num_reqs: int, ): self.vllm_config = vllm_config self.parallel_config = vllm_config.parallel_config self.num_ubatches = self.parallel_config.num_ubatches self.dcp_size = self.parallel_config.decode_context_parallel_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0 self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size self.device = device self.model_state = model_state self.attn_groups = attn_groups self.kv_cache_config = kv_cache_config # `query_start_loc` and `seq_lens` are rebased onto each microbatch's # own token range, so they cannot be views of the full batch's buffers. # Allocating up front keeps their addresses stable across replays. self.ubatch_query_start_loc = [ torch.zeros(max_num_reqs + 1, dtype=torch.int32, device=device) for _ in range(self.num_ubatches) ] self.ubatch_seq_lens = [ torch.zeros(max_num_reqs, dtype=torch.int32, device=device) for _ in range(self.num_ubatches) ] self.ubatch_dcp_local_seq_lens = [ torch.zeros(max_num_reqs, dtype=torch.int32, device=device) if self.dcp_size > 1 else None 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( self, 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) attn_metadata = [] slot_mappings_by_layer = [] is_padding = [] for i, ubatch_slice in enumerate(ubatch_slices): ubatch = _slice_input_batch( input_batch, ubatch_slice, self.ubatch_query_start_loc[i], self.ubatch_seq_lens[i], self.ubatch_dcp_local_seq_lens[i], dcp_size=self.dcp_size, dcp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, ) ubatch_slot_mappings = slot_mappings[:, ubatch_slice.token_slice] ubatch_block_tables = tuple( block_table[ubatch_slice.request_slice] for block_table in block_tables ) attn_metadata.append( self.model_state.prepare_attn( ubatch, 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( build_slot_mappings_by_layer(ubatch_slot_mappings, self.kv_cache_config) ) is_padding.append(ubatch.is_padding) return UBatchState( slices=ubatch_slices, forward_contexts=self._make_forward_contexts( ubatch_slices, attn_metadata, slot_mappings_by_layer, is_padding ), ) def _make_forward_contexts( self, ubatch_slices: UBatchSlices, attn_metadata: list[dict[str, Any]], slot_mappings_by_layer: list[dict[str, torch.Tensor]], is_padding: list[torch.Tensor | None], ) -> list[ForwardContext]: """Build one forward context per microbatch. The expert all-to-all runs once per microbatch, so each context's DP metadata carries that microbatch's token count. Microbatching pads every DP rank alike, so each rank can assume that count for the others too. """ dp_size = self.parallel_config.data_parallel_size forward_contexts = [] for i, ubatch_slice in enumerate(ubatch_slices): num_tokens = ubatch_slice.num_tokens num_tokens_across_dp = torch.full( (dp_size,), num_tokens, dtype=torch.int32, device="cpu" ) forward_contexts.append( create_forward_context( attn_metadata[i], self.vllm_config, dp_metadata=DPMetadata.make( self.parallel_config, num_tokens, num_tokens_across_dp ), cudagraph_runtime_mode=CUDAGraphMode.NONE, slot_mapping=slot_mappings_by_layer[i], is_padding=is_padding[i], ) ) return forward_contexts def run( self, model: Any, 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=compute_stream, forward_contexts=ubatch_state.forward_contexts, ready_barrier=self.ready_barrier, ) outputs: dict[int, Any] = {} errors: dict[int, BaseException] = {} @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 # Recorded so the join below can report which microbatch broke. # Note the siblings are not told to unwind: microbatches hand # control to each other, so one that dies mid-forward leaves # the others parked on a handoff that never comes and the step # hangs. The V1 runner behaves the same way; fixing it means # changing the shared handoff protocol in ubatching.py, which # is deliberately out of scope here. errors[ubatch_context.id] = e # 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) 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() # 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( vllm_config: VllmConfig, device: torch.device, model_state: ModelState, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, max_num_reqs: int, ) -> UBatchRunner | None: """Build the microbatch runner, or None when DBO is not in use. Microbatching needs the DP handshake to agree on it, so it is only available with more than one DP rank (as in the V1 runner). """ parallel_config = vllm_config.parallel_config if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1: return None logger.info_once("Dual batch overlap is enabled.") return UBatchRunner( vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs )