| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import logging |
| 5 | import threading |
| 6 | from collections.abc import Callable |
| 7 | from dataclasses import dataclass |
| 8 | from typing import Any |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.compilation.cuda_graph import CUDAGraphWrapper |
| 13 | from vllm.config import CUDAGraphMode, VllmConfig |
| 14 | from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id |
| 15 | from vllm.forward_context import ( |
| 16 | DPMetadata, |
| 17 | create_forward_context, |
| 18 | get_forward_context, |
| 19 | override_forward_context, |
| 20 | ) |
| 21 | from vllm.logger import init_logger |
| 22 | from vllm.model_executor.offloader.base import get_offloader |
| 23 | from vllm.platforms import current_platform |
| 24 | from vllm.sequence import IntermediateTensors |
| 25 | from vllm.v1.worker.ubatch_utils import create_sm_control_context |
| 26 | from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | |
| 31 | def _cat_ubatch_outputs( |
| 32 | sorted_results: list, |
| 33 | ) -> "torch.Tensor | tuple[torch.Tensor, ...]": |
| 34 | """Concatenate per-ubatch model outputs along the batch dim. |
| 35 | |
| 36 | Most models return a single hidden-states tensor per ubatch. Target |
| 37 | models running with auxiliary output (e.g. EAGLE3 speculative decoding, |
| 38 | which collects aux hidden states for the drafter) return a tuple of |
| 39 | tensors instead. Fan out over tuple components so `torch.cat` sees |
| 40 | matching shapes and the caller receives the same structure the model |
| 41 | produced for a single ubatch (#40769). |
| 42 | """ |
| 43 | if sorted_results and isinstance(sorted_results[0], tuple): |
| 44 | return tuple(torch.cat(parts, dim=0) for parts in zip(*sorted_results)) |
| 45 | return torch.cat(sorted_results, dim=0) |
| 46 | |
| 47 | |
| 48 | @dataclass |
| 49 | class UbatchMetadata: |
| 50 | context: UBatchContext |
| 51 | input_ids: torch.Tensor |
| 52 | positions: torch.Tensor |
| 53 | inputs_embeds: torch.Tensor | None |
| 54 | intermediate_tensors: IntermediateTensors | None |
| 55 | num_tokens: int |
| 56 | |
| 57 | |
| 58 | @dataclass |
| 59 | class CUDAGraphMetaData: |
| 60 | cudagraph: torch.cuda.CUDAGraph |
| 61 | ubatch_metadata: UbatchMetadata |
| 62 | outputs: Any | None = None |
| 63 | |
| 64 | |
| 65 | class UBatchWrapper: |
| 66 | def __init__( |
| 67 | self, |
| 68 | runnable: Callable, |
| 69 | vllm_config: VllmConfig, |
| 70 | runtime_mode: CUDAGraphMode, |
| 71 | device: torch.cuda.device, |
| 72 | ): |
| 73 | self.runnable = runnable |
| 74 | self.vllm_config = vllm_config |
| 75 | self.compilation_config = vllm_config.compilation_config |
| 76 | self.comm_stream = torch.cuda.Stream(device=device) |
| 77 | # Ubatch threads plus the main thread |
| 78 | self.ready_barrier = threading.Barrier( |
| 79 | self.vllm_config.parallel_config.num_ubatches + 1 |
| 80 | ) |
| 81 | |
| 82 | self.cudagraphs: dict[int, CUDAGraphMetaData] = {} |
| 83 | |
| 84 | self.cudagraph_wrapper = None |
| 85 | if runtime_mode is not CUDAGraphMode.NONE: |
| 86 | self.cudagraph_wrapper = CUDAGraphWrapper( |
| 87 | runnable, vllm_config, runtime_mode=runtime_mode |
| 88 | ) |
| 89 | |
| 90 | self.sm_control = create_sm_control_context(vllm_config.parallel_config) |
| 91 | self.device = device |
| 92 | self.is_debugging_mode = logger.isEnabledFor(logging.DEBUG) |
| 93 | self._runnable_str = str(runnable) if self.is_debugging_mode else None |
| 94 | |
| 95 | @property |
| 96 | def graph_pool(self): |
| 97 | if self.cudagraph_wrapper is not None: |
| 98 | return self.cudagraph_wrapper.graph_pool |
| 99 | return None |
| 100 | |
| 101 | def clear_graphs(self) -> None: |
| 102 | self.cudagraphs.clear() |
| 103 | if self.cudagraph_wrapper is not None: |
| 104 | self.cudagraph_wrapper.clear_graphs() |
| 105 | |
| 106 | def __getattr__(self, key: str): |
| 107 | # allow accessing the attributes of the runnable. |
| 108 | if hasattr(self.runnable, key): |
| 109 | return getattr(self.runnable, key) |
| 110 | if self.is_debugging_mode: |
| 111 | raise AttributeError( |
| 112 | f"Attribute {key} not exists in the runnable of " |
| 113 | f"cudagraph wrapper: {self._runnable_str}" |
| 114 | ) |
| 115 | raise AttributeError |
| 116 | |
| 117 | def unwrap(self) -> Callable: |
| 118 | # in case we need to access the original runnable. |
| 119 | return self.runnable |
| 120 | |
| 121 | def _capture_ubatches(self, ubatch_metadata, model) -> torch.Tensor: |
| 122 | """Capture a cudagraph for a microbatched run. |
| 123 | |
| 124 | The logic here is somewhat complicated because we need to make sure that |
| 125 | each of the ubatch threads initialize the cuda context before we start |
| 126 | the graph capture. |
| 127 | |
| 128 | The flow is as follows: |
| 129 | 1. The main thread starts up each ubatch thread. Each thread will |
| 130 | initialize its cuda context (torch.cuda.current_blas_handle()) |
| 131 | before going to sleep upon entering the ubatch_context. |
| 132 | |
| 133 | 2. The main thread starts the graph capture and wakes up the first |
| 134 | ubatch thread. |
| 135 | |
| 136 | 3. Each ubatch thread runs the model to completion and returns the |
| 137 | completed output tensors back to the main thread. |
| 138 | |
| 139 | 4. The main thread stores the captured cudagraph along with its metadata |
| 140 | and returns |
| 141 | """ |
| 142 | |
| 143 | @torch.inference_mode() |
| 144 | def _capture_ubatch_thread(results, ubatch_metadata): |
| 145 | torch.accelerator.set_device_index(self.device) |
| 146 | ubatch_context = ubatch_metadata.context |
| 147 | with torch.cuda.stream(ubatch_context.compute_stream): |
| 148 | _ = torch.cuda.current_blas_handle() |
| 149 | with torch.cuda.stream(ubatch_context.comm_stream): |
| 150 | _ = torch.cuda.current_blas_handle() |
| 151 | with ubatch_context: |
| 152 | model_output = model( |
| 153 | input_ids=ubatch_metadata.input_ids, |
| 154 | positions=ubatch_metadata.positions, |
| 155 | intermediate_tensors=ubatch_metadata.intermediate_tensors, |
| 156 | inputs_embeds=ubatch_metadata.inputs_embeds, |
| 157 | ) |
| 158 | |
| 159 | results.append((ubatch_metadata.context.id, model_output)) |
| 160 | |
| 161 | results: list[tuple[int, torch.Tensor]] = [] |
| 162 | compute_stream = ubatch_metadata[0].context.compute_stream |
| 163 | num_tokens = sum(m.num_tokens for m in ubatch_metadata) |
| 164 | |
| 165 | # Ubatches will manually manage the forward context, so we override |
| 166 | # it to None here so we can have it restored correctly later |
| 167 | with override_forward_context(None): |
| 168 | ubatch_threads = [] |
| 169 | for metadata in ubatch_metadata: |
| 170 | thread = threading.Thread( |
| 171 | target=_capture_ubatch_thread, |
| 172 | args=( |
| 173 | results, |
| 174 | metadata, |
| 175 | ), |
| 176 | ) |
| 177 | ubatch_threads.append(thread) |
| 178 | thread.start() |
| 179 | self.ready_barrier.wait() # Wait for all ubatch threads to be ready |
| 180 | |
| 181 | # Capture the cudagraph |
| 182 | cudagraph_metadata = CUDAGraphMetaData( |
| 183 | cudagraph=torch.cuda.CUDAGraph(), |
| 184 | ubatch_metadata=ubatch_metadata, |
| 185 | ) |
| 186 | if self.graph_pool is not None: |
| 187 | set_graph_pool_id(self.graph_pool) |
| 188 | else: |
| 189 | set_graph_pool_id(current_platform.graph_pool_handle()) |
| 190 | |
| 191 | # Sync offloader's copy stream before capture. |
| 192 | # Ensure any pre-capture prefetches from offloader are complete. |
| 193 | get_offloader().sync_prev_onload() |
| 194 | |
| 195 | with torch.cuda.graph( |
| 196 | cudagraph_metadata.cudagraph, |
| 197 | stream=compute_stream, |
| 198 | pool=self.graph_pool, |
| 199 | ): |
| 200 | ubatch_metadata[0].context.cpu_wait_event.set() |
| 201 | for thread in ubatch_threads: |
| 202 | thread.join() |
| 203 | sorted_results = [value for position, value in sorted(results)] |
| 204 | result = _cat_ubatch_outputs(sorted_results) |
| 205 | cudagraph_metadata.outputs = result |
| 206 | # Join offloader's copy stream after forward to avoid unjoined |
| 207 | # stream error. The last layer's start_prefetch forks copy_stream, |
| 208 | # but wait_prefetch only happens in the next forward pass. |
| 209 | get_offloader().join_after_forward() |
| 210 | self.cudagraphs[num_tokens] = cudagraph_metadata |
| 211 | return cudagraph_metadata.outputs |
| 212 | |
| 213 | def _run_ubatches(self, ubatch_metadata, model) -> torch.Tensor: |
| 214 | @torch.inference_mode() |
| 215 | def _ubatch_thread(results, model, ubatch_metadata): |
| 216 | with ubatch_metadata.context: |
| 217 | model_output = model( |
| 218 | input_ids=ubatch_metadata.input_ids, |
| 219 | positions=ubatch_metadata.positions, |
| 220 | intermediate_tensors=ubatch_metadata.intermediate_tensors, |
| 221 | inputs_embeds=ubatch_metadata.inputs_embeds, |
| 222 | ) |
| 223 | results.append((ubatch_metadata.context.id, model_output)) |
| 224 | |
| 225 | results: list[tuple[int, torch.Tensor]] = [] |
| 226 | |
| 227 | # Ubatch threads will manually manage the forward context, so we |
| 228 | # override it to None here so we can have it restored correctly |
| 229 | # after both threads have finished |
| 230 | with override_forward_context(None): |
| 231 | ubatch_threads = [] |
| 232 | for metadata in ubatch_metadata: |
| 233 | thread = threading.Thread( |
| 234 | target=_ubatch_thread, |
| 235 | args=( |
| 236 | results, |
| 237 | model, |
| 238 | metadata, |
| 239 | ), |
| 240 | ) |
| 241 | ubatch_threads.append(thread) |
| 242 | thread.start() |
| 243 | self.ready_barrier.wait() # Wait for all ubatch threads to be ready |
| 244 | ubatch_metadata[0].context.cpu_wait_event.set() |
| 245 | for thread in ubatch_threads: |
| 246 | thread.join() |
| 247 | sorted_results = [value for position, value in sorted(results)] |
| 248 | result = _cat_ubatch_outputs(sorted_results) |
| 249 | return result |
| 250 | |
| 251 | def _make_ubatch_metadata( |
| 252 | self, |
| 253 | ubatch_slices, |
| 254 | attn_metadata, |
| 255 | slot_mapping, |
| 256 | input_ids, |
| 257 | positions, |
| 258 | inputs_embeds, |
| 259 | intermediate_tensors, |
| 260 | compute_stream, |
| 261 | dp_metadata, |
| 262 | batch_descriptor, |
| 263 | cudagraph_runtime_mode, |
| 264 | is_padding, |
| 265 | ) -> list[UbatchMetadata]: |
| 266 | # Create one forward context per ubatch |
| 267 | forward_contexts = [] |
| 268 | # slot_mapping can be None, an empty dict (from create_forward_context |
| 269 | # converting None to {}), or a list of dicts (one per ubatch) |
| 270 | has_slot_mapping = slot_mapping and isinstance(slot_mapping, list) |
| 271 | for i, ubatch_slice in enumerate(ubatch_slices): |
| 272 | forward_contexts.append( |
| 273 | create_forward_context( |
| 274 | attn_metadata[i] if attn_metadata is not None else None, |
| 275 | self.vllm_config, |
| 276 | dp_metadata=dp_metadata[i], |
| 277 | batch_descriptor=batch_descriptor, |
| 278 | cudagraph_runtime_mode=cudagraph_runtime_mode, |
| 279 | slot_mapping=slot_mapping[i] if has_slot_mapping else None, |
| 280 | is_padding=( |
| 281 | is_padding[ubatch_slice.token_slice] |
| 282 | if is_padding is not None |
| 283 | else None |
| 284 | ), |
| 285 | ) |
| 286 | ) |
| 287 | |
| 288 | ubatch_ctxs = make_ubatch_contexts( |
| 289 | num_micro_batches=len(ubatch_slices), |
| 290 | comm_stream=self.comm_stream, |
| 291 | compute_stream=compute_stream, |
| 292 | forward_contexts=forward_contexts, |
| 293 | ready_barrier=self.ready_barrier, |
| 294 | ) |
| 295 | |
| 296 | ubatch_metadata: list[UbatchMetadata] = [] |
| 297 | for i, ubatch_slice in enumerate(ubatch_slices): |
| 298 | ( |
| 299 | sliced_input_ids, |
| 300 | sliced_positions, |
| 301 | sliced_inputs_embeds, |
| 302 | sliced_intermediate_tensors, |
| 303 | ) = self._slice_model_inputs( |
| 304 | ubatch_slice.token_slice, |
| 305 | input_ids, |
| 306 | positions, |
| 307 | inputs_embeds, |
| 308 | intermediate_tensors, |
| 309 | ) |
| 310 | ubatch_metadata.append( |
| 311 | UbatchMetadata( |
| 312 | context=ubatch_ctxs[i], |
| 313 | input_ids=sliced_input_ids, |
| 314 | positions=sliced_positions, |
| 315 | inputs_embeds=sliced_inputs_embeds, |
| 316 | intermediate_tensors=sliced_intermediate_tensors, |
| 317 | num_tokens=ubatch_slice.token_slice.stop |
| 318 | - ubatch_slice.token_slice.start, |
| 319 | ) |
| 320 | ) |
| 321 | |
| 322 | return ubatch_metadata |
| 323 | |
| 324 | def _slice_model_inputs( |
| 325 | self, |
| 326 | tokens_slice: slice, |
| 327 | input_ids, |
| 328 | positions, |
| 329 | inputs_embeds, |
| 330 | intermediate_tensors, |
| 331 | ): |
| 332 | sliced_input_ids = input_ids[tokens_slice] if input_ids is not None else None |
| 333 | # if we are using mrope. Mrope adds an additional dimension to the |
| 334 | # positions tensor |
| 335 | if positions.ndim == 2: |
| 336 | sliced_positions = positions[:, tokens_slice] |
| 337 | else: |
| 338 | sliced_positions = positions[tokens_slice] |
| 339 | sliced_inputs_embeds = ( |
| 340 | inputs_embeds[tokens_slice] if inputs_embeds is not None else None |
| 341 | ) |
| 342 | sliced_intermediate_tensors = ( |
| 343 | intermediate_tensors[tokens_slice] |
| 344 | if intermediate_tensors is not None |
| 345 | else None |
| 346 | ) |
| 347 | |
| 348 | return ( |
| 349 | sliced_input_ids, |
| 350 | sliced_positions, |
| 351 | sliced_inputs_embeds, |
| 352 | sliced_intermediate_tensors, |
| 353 | ) |
| 354 | |
| 355 | def __call__(self, *args, **kwargs): |
| 356 | forward_context = get_forward_context() |
| 357 | batch_descriptor = forward_context.batch_descriptor |
| 358 | ubatch_slices = forward_context.ubatch_slices |
| 359 | cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode |
| 360 | |
| 361 | # If there's no ubatching, just run the runnable object |
| 362 | if ubatch_slices is None: |
| 363 | # This is to account for the case where ubatching was aborted. |
| 364 | # When we capture full graphs we only capture one graph per shape, |
| 365 | # meaning that if we have a ubatched cudagraph for the current |
| 366 | # num_tokens, we don't have a non-ubatched one. Without this |
| 367 | # check, the cudagraph wrapper will try to capture a cudagraph |
| 368 | # for this shape during a normal run. |
| 369 | if cudagraph_runtime_mode is CUDAGraphMode.FULL: |
| 370 | assert batch_descriptor is not None |
| 371 | if batch_descriptor.num_tokens in self.cudagraphs: |
| 372 | cudagraph_runtime_mode = CUDAGraphMode.NONE |
| 373 | |
| 374 | if cudagraph_runtime_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE): |
| 375 | return self.runnable(*args, **kwargs) |
| 376 | else: |
| 377 | assert self.cudagraph_wrapper is not None |
| 378 | return self.cudagraph_wrapper(*args, **kwargs) |
| 379 | |
| 380 | attn_metadata = forward_context.attn_metadata |
| 381 | slot_mapping = forward_context.slot_mapping |
| 382 | is_padding = forward_context.is_padding |
| 383 | num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices) |
| 384 | input_ids = kwargs["input_ids"] |
| 385 | positions = kwargs["positions"] |
| 386 | intermediate_tensors = kwargs["intermediate_tensors"] |
| 387 | inputs_embeds = kwargs["inputs_embeds"] |
| 388 | compute_stream = torch.cuda.current_stream() |
| 389 | |
| 390 | dp_metadata = forward_context.dp_metadata |
| 391 | |
| 392 | # We shouldn't be here unless we are running with multiple DP ranks |
| 393 | assert dp_metadata is not None |
| 394 | ubatch_dp_metadata = [] |
| 395 | for ubatch_slice in ubatch_slices: |
| 396 | dp_size = self.vllm_config.parallel_config.data_parallel_size |
| 397 | ubatch_num_tokens_across_dp = torch.tensor( |
| 398 | [ubatch_slice.num_tokens] * dp_size, device="cpu", dtype=torch.int32 |
| 399 | ) |
| 400 | ubatch_dp_metadata.append( |
| 401 | DPMetadata.make( |
| 402 | self.vllm_config.parallel_config, |
| 403 | ubatch_slice.num_tokens, |
| 404 | ubatch_num_tokens_across_dp, |
| 405 | ) |
| 406 | ) |
| 407 | |
| 408 | if ( |
| 409 | num_tokens not in self.cudagraphs |
| 410 | and cudagraph_runtime_mode is CUDAGraphMode.FULL |
| 411 | ): |
| 412 | ubatch_metadata = self._make_ubatch_metadata( |
| 413 | ubatch_slices=ubatch_slices, |
| 414 | attn_metadata=attn_metadata, |
| 415 | slot_mapping=slot_mapping, |
| 416 | input_ids=input_ids, |
| 417 | positions=positions, |
| 418 | intermediate_tensors=intermediate_tensors, |
| 419 | inputs_embeds=inputs_embeds, |
| 420 | compute_stream=compute_stream, |
| 421 | dp_metadata=ubatch_dp_metadata, |
| 422 | batch_descriptor=batch_descriptor, |
| 423 | cudagraph_runtime_mode=CUDAGraphMode.NONE, |
| 424 | is_padding=is_padding, |
| 425 | ) |
| 426 | with self.sm_control: |
| 427 | return self._capture_ubatches(ubatch_metadata, self.runnable) |
| 428 | elif ( |
| 429 | num_tokens in self.cudagraphs |
| 430 | and cudagraph_runtime_mode is CUDAGraphMode.FULL |
| 431 | ): |
| 432 | cudagraph_metadata = self.cudagraphs[num_tokens] |
| 433 | # Sync offloader before replay - ensures any external dependencies |
| 434 | # from pre-capture prefetches are satisfied. |
| 435 | get_offloader().sync_prev_onload() |
| 436 | cudagraph_metadata.cudagraph.replay() |
| 437 | return cudagraph_metadata.outputs |
| 438 | else: |
| 439 | ubatch_metadata = self._make_ubatch_metadata( |
| 440 | ubatch_slices=ubatch_slices, |
| 441 | attn_metadata=attn_metadata, |
| 442 | slot_mapping=slot_mapping, |
| 443 | input_ids=input_ids, |
| 444 | positions=positions, |
| 445 | intermediate_tensors=intermediate_tensors, |
| 446 | inputs_embeds=inputs_embeds, |
| 447 | compute_stream=compute_stream, |
| 448 | dp_metadata=ubatch_dp_metadata, |
| 449 | batch_descriptor=batch_descriptor, |
| 450 | cudagraph_runtime_mode=CUDAGraphMode.NONE, |
| 451 | is_padding=is_padding, |
| 452 | ) |
| 453 | with self.sm_control: |
| 454 | return self._run_ubatches(ubatch_metadata, self.runnable) |