| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CUDA graph manager for vision encoder budget-batch execution.""" |
| 4 | |
| 5 | import itertools |
| 6 | import math |
| 7 | from collections.abc import Hashable |
| 8 | from dataclasses import dataclass |
| 9 | from typing import Any, TypeAlias |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm.config import VllmConfig |
| 14 | from vllm.distributed import ( |
| 15 | get_tensor_model_parallel_rank, |
| 16 | get_tensor_model_parallel_world_size, |
| 17 | tensor_model_parallel_all_gather, |
| 18 | ) |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.model_executor.models.interfaces import ( |
| 21 | SupportsEncoderCudaGraph, |
| 22 | ) |
| 23 | from vllm.model_executor.models.utils import scatter_output_slices |
| 24 | from vllm.model_executor.models.vision import get_load_balance_assignment |
| 25 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 26 | from vllm.utils.torch_utils import current_stream |
| 27 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 28 | ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, |
| 29 | EncoderCudaGraphConfig, |
| 30 | EncoderItemSpec, |
| 31 | ) |
| 32 | |
| 33 | logger = init_logger(__name__) |
| 34 | |
| 35 | CaptureAxisKeys: TypeAlias = tuple[Hashable, ...] |
| 36 | |
| 37 | BudgetGraphMapKey: TypeAlias = int | tuple[int, CaptureAxisKeys] |
| 38 | |
| 39 | # Warn when capture axes multiply the captured graph count past this point. |
| 40 | _CAPTURE_GRAPH_COUNT_WARNING_THRESHOLD = 64 |
| 41 | |
| 42 | |
| 43 | @dataclass |
| 44 | class BudgetGraphMetadata: |
| 45 | """Metadata for a single budget graph. |
| 46 | |
| 47 | CUDA graph replay pattern: |
| 48 | * Copy precomputed values into input_buffers |
| 49 | * Replay graph |
| 50 | * Read encoder outputs from output_buffer |
| 51 | """ |
| 52 | |
| 53 | token_budget: int |
| 54 | max_batch_size: int # Max number of images/videos per batch |
| 55 | max_frames_per_batch: int # Max total frames per batch (for video) |
| 56 | graph: torch.cuda.CUDAGraph |
| 57 | # Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata). |
| 58 | # Before replay the manager updates these in-place. By default buffers are |
| 59 | # zeroed before slice-copying the actual values; model-specific padding |
| 60 | # behavior is provided by EncoderCudaGraphConfig.padding_logics. |
| 61 | input_buffers: dict[str, torch.Tensor] |
| 62 | # Output written by graph, read after replay |
| 63 | output_buffer: torch.Tensor |
| 64 | # Per-axis capture keys, empty when no extra capture axis is used. |
| 65 | axis_keys: CaptureAxisKeys = () |
| 66 | |
| 67 | |
| 68 | class EncoderCudaGraphManager: |
| 69 | """Budget-based CUDA graph capture/replay for vision encoders.""" |
| 70 | |
| 71 | def __init__( |
| 72 | self, |
| 73 | vllm_config: VllmConfig, |
| 74 | device: torch.device, |
| 75 | dtype: torch.dtype, |
| 76 | model: SupportsEncoderCudaGraph, |
| 77 | ): |
| 78 | """Initialize CUDA graph manager with provided token budgets |
| 79 | and max batch size.""" |
| 80 | self.vllm_config = vllm_config |
| 81 | self.device = device |
| 82 | self.dtype = dtype |
| 83 | self.model = model |
| 84 | self.config: EncoderCudaGraphConfig = model.get_encoder_cudagraph_config() |
| 85 | |
| 86 | comp_config = vllm_config.compilation_config |
| 87 | user_budgets = comp_config.encoder_cudagraph_token_budgets |
| 88 | user_max_vision_items = comp_config.encoder_cudagraph_max_vision_items_per_batch |
| 89 | user_max_frames = comp_config.encoder_cudagraph_max_frames_per_batch |
| 90 | |
| 91 | multimodal_config = vllm_config.model_config.multimodal_config |
| 92 | |
| 93 | self._capture_axes: tuple[CaptureAxisKeys, ...] = tuple( |
| 94 | tuple(axis) for axis in self.config.capture_axes |
| 95 | ) |
| 96 | if any(len(axis) == 0 for axis in self._capture_axes): |
| 97 | raise ValueError("Encoder cudagraph capture axes must be non-empty.") |
| 98 | |
| 99 | if len(self.config.paths) > 1 and self._capture_axes: |
| 100 | raise NotImplementedError( |
| 101 | "Combining dual-path encoder CUDA graphs with capture axes " |
| 102 | "is not supported yet." |
| 103 | ) |
| 104 | |
| 105 | # Invariant: max_batch_size <= min_token_budget. |
| 106 | # This ensures per_image_output = budget // max_batch_size >= 1 |
| 107 | # for every captured budget, preventing reshape crashes on empty |
| 108 | # tensors during CUDA graph capture. Validated/enforced below for |
| 109 | # each configuration path. |
| 110 | if user_budgets and user_max_vision_items > 0: |
| 111 | # Fully user-specified: validate the invariant. |
| 112 | self.token_budgets = sorted(user_budgets) |
| 113 | self.max_batch_size = user_max_vision_items |
| 114 | min_tok = min(self.token_budgets) |
| 115 | if self.max_batch_size > min_tok: |
| 116 | raise ValueError( |
| 117 | f"encoder_cudagraph_max_vision_items_per_batch " |
| 118 | f"({self.max_batch_size}) must be <= smallest token " |
| 119 | f"budget ({min_tok}). With budgets=" |
| 120 | f"{self.token_budgets}, per_image_output = " |
| 121 | f"{min_tok} // {self.max_batch_size} = " |
| 122 | f"{min_tok // self.max_batch_size}, which would cause " |
| 123 | f"a capture failure. Either increase the smallest " |
| 124 | f"budget or decrease max_vision_items_per_batch." |
| 125 | ) |
| 126 | else: |
| 127 | # Auto-infer missing values from model. |
| 128 | min_budget, max_budget = model.get_encoder_cudagraph_budget_range( |
| 129 | vllm_config |
| 130 | ) |
| 131 | if min_budget <= 0 or max_budget <= 0: |
| 132 | raise ValueError( |
| 133 | f"Invalid encoder cudagraph budget range: " |
| 134 | f"min_budget={min_budget}, max_budget={max_budget}. " |
| 135 | f"Both must be positive." |
| 136 | ) |
| 137 | if min_budget > max_budget: |
| 138 | raise ValueError( |
| 139 | f"Invalid encoder cudagraph budget range: " |
| 140 | f"min_budget={min_budget} > max_budget={max_budget}." |
| 141 | ) |
| 142 | |
| 143 | if user_max_vision_items > 0: |
| 144 | # User provided max_vision_items only; adjust auto-inferred |
| 145 | # budgets so min(budgets) >= max_batch_size. |
| 146 | self.max_batch_size = user_max_vision_items |
| 147 | effective_min = max(min_budget, user_max_vision_items) |
| 148 | self.token_budgets = self._generate_budgets(effective_min, max_budget) |
| 149 | elif user_budgets: |
| 150 | # User provided budgets only; cap auto-inferred |
| 151 | # max_batch_size to min(user_budgets). |
| 152 | self.token_budgets = sorted(user_budgets) |
| 153 | self.max_batch_size = min( |
| 154 | max_budget // min_budget, |
| 155 | min(self.token_budgets), |
| 156 | ) |
| 157 | else: |
| 158 | # Fully auto-inferred. |
| 159 | self.token_budgets = self._generate_budgets(min_budget, max_budget) |
| 160 | self.max_batch_size = min( |
| 161 | max_budget // min_budget, |
| 162 | min(self.token_budgets), |
| 163 | ) |
| 164 | |
| 165 | assert multimodal_config is not None |
| 166 | if multimodal_config.get_limit_per_prompt("video") == 0: |
| 167 | self.max_frames_per_batch = 0 |
| 168 | elif user_max_frames is not None: |
| 169 | self.max_frames_per_batch = user_max_frames |
| 170 | else: |
| 171 | # Set it to the model-specific value from config. |
| 172 | max_frames_per_video = self.config.max_frames_per_video |
| 173 | self.max_frames_per_batch = self.max_batch_size * max_frames_per_video |
| 174 | |
| 175 | mm_config = vllm_config.model_config.multimodal_config |
| 176 | self.use_dp = ( |
| 177 | mm_config is not None |
| 178 | and mm_config.mm_encoder_tp_mode == "data" |
| 179 | and vllm_config.parallel_config.tensor_parallel_size > 1 |
| 180 | ) |
| 181 | |
| 182 | self.budget_graphs: dict[str, dict[BudgetGraphMapKey, BudgetGraphMetadata]] = {} |
| 183 | self.graph_pool: Any | None = None |
| 184 | self.graph_hits = 0 |
| 185 | self.graph_misses = 0 |
| 186 | self.log_stats_interval = 100 |
| 187 | |
| 188 | max_budget = self.token_budgets[-1] |
| 189 | if not self.config.paths: |
| 190 | raise ValueError("Encoder CUDA graph config must define at least one path") |
| 191 | self.path_token_budgets: dict[str, list[int]] = {} |
| 192 | for path, path_config in self.config.paths.items(): |
| 193 | min_path_budget = path_config.min_token_budget |
| 194 | if min_path_budget is None: |
| 195 | budgets = list(self.token_budgets) |
| 196 | else: |
| 197 | if min_path_budget <= 0 or min_path_budget > max_budget: |
| 198 | raise ValueError( |
| 199 | f"Invalid minimum budget {min_path_budget} for encoder " |
| 200 | f"CUDA graph path {path!r}; max budget is {max_budget}" |
| 201 | ) |
| 202 | budgets = self._generate_budgets(min_path_budget, max_budget) |
| 203 | if path_config.allow_zero_tokens: |
| 204 | budgets.insert(0, 0) |
| 205 | self.path_token_budgets[path] = budgets |
| 206 | |
| 207 | logger.info( |
| 208 | "EncoderCudaGraphManager initialized with paths=%s, " |
| 209 | "max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s, " |
| 210 | "capture_axes=%s", |
| 211 | self.path_token_budgets, |
| 212 | self.max_batch_size, |
| 213 | self.max_frames_per_batch, |
| 214 | self.use_dp, |
| 215 | self._capture_axes, |
| 216 | ) |
| 217 | |
| 218 | @staticmethod |
| 219 | def _generate_budgets(min_budget: int, max_budget: int) -> list[int]: |
| 220 | """Generate power-of-2 token budgets from min_budget to max_budget.""" |
| 221 | budgets: list[int] = [] |
| 222 | b = min_budget |
| 223 | while b <= max_budget: |
| 224 | budgets.append(b) |
| 225 | b *= 2 |
| 226 | # Always include max_budget if it's not already a power-of-2 boundary |
| 227 | if not budgets or budgets[-1] < max_budget: |
| 228 | budgets.append(max_budget) |
| 229 | return budgets |
| 230 | |
| 231 | def supports_modality(self, modality: str) -> bool: |
| 232 | """Check if a modality is supported by this manager.""" |
| 233 | return modality in self.config.modalities |
| 234 | |
| 235 | def is_captured(self) -> bool: |
| 236 | """Return whether a CUDA graph pool is active.""" |
| 237 | return self.graph_pool is not None |
| 238 | |
| 239 | def clear(self) -> None: |
| 240 | """Release captured encoder CUDA graphs and the manager-local pool.""" |
| 241 | for graph_set in self.budget_graphs.values(): |
| 242 | graph_set.clear() |
| 243 | self.graph_pool = None |
| 244 | |
| 245 | def capture(self, graph_pool: Any): |
| 246 | """Capture CUDA graphs for every configured path and token budget.""" |
| 247 | self.graph_pool = graph_pool |
| 248 | |
| 249 | num_graphs = self.get_num_graphs_to_capture() |
| 250 | if num_graphs > _CAPTURE_GRAPH_COUNT_WARNING_THRESHOLD: |
| 251 | logger.warning( |
| 252 | "Capturing %d encoder CUDA graphs; capture axes multiply the " |
| 253 | "per-budget graph count, consider pruning axis keys.", |
| 254 | num_graphs, |
| 255 | ) |
| 256 | |
| 257 | for path, budgets in self.path_token_budgets.items(): |
| 258 | for token_budget in sorted(budgets, reverse=True): |
| 259 | if token_budget == 0: |
| 260 | continue |
| 261 | for axis_keys in itertools.product(*self._capture_axes): |
| 262 | self._capture_budget_graph( |
| 263 | token_budget, |
| 264 | path=path, |
| 265 | axis_keys=axis_keys, |
| 266 | ) |
| 267 | |
| 268 | logger.info( |
| 269 | "Encoder CUDA graph capture complete. Captured %d graphs across %d paths.", |
| 270 | num_graphs, |
| 271 | len(self.path_token_budgets), |
| 272 | ) |
| 273 | |
| 274 | def _num_graphs_for_budgets(self, budgets: list[int]) -> int: |
| 275 | num_budgets = sum(1 for budget in budgets if budget != 0) |
| 276 | return num_budgets * math.prod(len(axis) for axis in self._capture_axes) |
| 277 | |
| 278 | def get_num_graphs_to_capture(self) -> int: |
| 279 | return sum( |
| 280 | self._num_graphs_for_budgets(budgets) |
| 281 | for budgets in self.path_token_budgets.values() |
| 282 | ) |
| 283 | |
| 284 | def _get_graph_set( |
| 285 | self, path: str = "default" |
| 286 | ) -> dict[BudgetGraphMapKey, BudgetGraphMetadata]: |
| 287 | """Return the captured graphs for one encoder path.""" |
| 288 | if path not in self.budget_graphs: |
| 289 | self.budget_graphs[path] = {} |
| 290 | return self.budget_graphs[path] |
| 291 | |
| 292 | def _capture_budget_graph( |
| 293 | self, |
| 294 | token_budget: int, |
| 295 | path: str = "default", |
| 296 | axis_keys: CaptureAxisKeys = (), |
| 297 | ): |
| 298 | """Capture CUDA graph for a single token budget.""" |
| 299 | logger.debug( |
| 300 | "Capturing encoder cudagraph for budget=%d, max_batch_size=%d, " |
| 301 | "max_frames_per_batch=%d, axis_keys=%s", |
| 302 | token_budget, |
| 303 | self.max_batch_size, |
| 304 | self.max_frames_per_batch, |
| 305 | axis_keys, |
| 306 | ) |
| 307 | |
| 308 | graph_set = self._get_graph_set(path) |
| 309 | |
| 310 | capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs( |
| 311 | token_budget, |
| 312 | self.max_batch_size, |
| 313 | self.max_frames_per_batch, |
| 314 | self.device, |
| 315 | self.dtype, |
| 316 | path, |
| 317 | axis_keys, |
| 318 | ) |
| 319 | |
| 320 | values = capture_inputs.values |
| 321 | |
| 322 | with torch.inference_mode(): |
| 323 | output = self.model.encoder_cudagraph_forward({**values}, path=path) |
| 324 | output_buffer = torch.empty_like(output) |
| 325 | |
| 326 | graph = torch.cuda.CUDAGraph() |
| 327 | with ( |
| 328 | torch.inference_mode(), |
| 329 | torch.cuda.graph(graph, pool=self.graph_pool, stream=current_stream()), |
| 330 | ): |
| 331 | output = self.model.encoder_cudagraph_forward({**values}, path=path) |
| 332 | output_buffer.copy_(output) |
| 333 | |
| 334 | graph_map_key: BudgetGraphMapKey = ( |
| 335 | token_budget if not axis_keys else (token_budget, axis_keys) |
| 336 | ) |
| 337 | graph_set[graph_map_key] = BudgetGraphMetadata( |
| 338 | token_budget=token_budget, |
| 339 | max_batch_size=self.max_batch_size, |
| 340 | max_frames_per_batch=self.max_frames_per_batch, |
| 341 | graph=graph, |
| 342 | input_buffers=values, |
| 343 | output_buffer=output_buffer, |
| 344 | axis_keys=axis_keys, |
| 345 | ) |
| 346 | |
| 347 | def _find_smallest_fitting_budget_given_tokens( |
| 348 | self, total_tokens: int, budgets: list[int] | None = None |
| 349 | ) -> int | None: |
| 350 | """Find smallest budget >= total_tokens. |
| 351 | |
| 352 | Returns: |
| 353 | Token budget if found, None if no fitting budget. |
| 354 | |
| 355 | """ |
| 356 | budgets = budgets if budgets is not None else self.token_budgets |
| 357 | for budget in budgets: |
| 358 | if budget >= total_tokens: |
| 359 | return budget |
| 360 | return None |
| 361 | |
| 362 | def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]: |
| 363 | """Get item specs from the model.""" |
| 364 | # Implementations read per-item grid/patch counts off device tensors |
| 365 | # to size the cudagraph buffers, so the D2H is inherent here. |
| 366 | with gpu_sync_allowed(): |
| 367 | return self.model.get_encoder_cudagraph_item_specs(mm_kwargs) |
| 368 | |
| 369 | def _select_items( |
| 370 | self, mm_kwargs: dict[str, Any], indices: list[int] |
| 371 | ) -> tuple[dict[str, Any], CaptureAxisKeys]: |
| 372 | """Select the mm kwargs for `indices` from the model. |
| 373 | |
| 374 | Returns the sliced mm_kwargs together with the resolved capture-axis |
| 375 | keys (popped out of the returned dict; empty without capture axes). |
| 376 | """ |
| 377 | # Same as `_get_item_specs`: implementations re-read the per-item |
| 378 | # grid/patch counts to slice the batch, so the D2H is inherent. |
| 379 | with gpu_sync_allowed(): |
| 380 | selected = self.model.select_encoder_cudagraph_items(mm_kwargs, indices) |
| 381 | axis_keys = selected.pop(ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, ()) |
| 382 | if len(axis_keys) != len(self._capture_axes): |
| 383 | raise ValueError( |
| 384 | f"Model returned {len(axis_keys)} capture axis keys, " |
| 385 | f"expected {len(self._capture_axes)} " |
| 386 | f"(one per axis of capture_axes)." |
| 387 | ) |
| 388 | return selected, axis_keys |
| 389 | |
| 390 | def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]: |
| 391 | """Get per-item output token counts as plain ints.""" |
| 392 | return [spec.output_tokens for spec in self._get_item_specs(mm_kwargs)] |
| 393 | |
| 394 | @staticmethod |
| 395 | def _copy_padded_buffer( |
| 396 | dst: torch.Tensor, |
| 397 | src: torch.Tensor, |
| 398 | ) -> None: |
| 399 | dst.zero_() |
| 400 | dst[: src.shape[0]].copy_(src) |
| 401 | |
| 402 | def _run_budget_graph( |
| 403 | self, |
| 404 | mm_kwargs: dict[str, Any], |
| 405 | token_budget: int, |
| 406 | path: str = "default", |
| 407 | axis_keys: CaptureAxisKeys = (), |
| 408 | ) -> torch.Tensor | None: |
| 409 | """Execute budget graph. |
| 410 | |
| 411 | Args: |
| 412 | mm_kwargs: Multimodal inputs for the batch. |
| 413 | token_budget: Token budget to use. |
| 414 | path: Configured encoder path. |
| 415 | axis_keys: Resolved capture-axis keys for this batch. |
| 416 | |
| 417 | Returns: |
| 418 | Encoder outputs, or None if graph not captured. |
| 419 | |
| 420 | """ |
| 421 | graph_set = self._get_graph_set(path) |
| 422 | num_items = len(self._get_item_specs(mm_kwargs)) |
| 423 | |
| 424 | graph_map_key: BudgetGraphMapKey = ( |
| 425 | token_budget if not axis_keys else (token_budget, axis_keys) |
| 426 | ) |
| 427 | if graph_map_key not in graph_set: |
| 428 | self.graph_misses += num_items |
| 429 | return None |
| 430 | |
| 431 | graph_meta = graph_set[graph_map_key] |
| 432 | |
| 433 | replay = self.model.prepare_encoder_cudagraph_replay_buffers( |
| 434 | mm_kwargs, |
| 435 | self.max_batch_size, |
| 436 | self.max_frames_per_batch, |
| 437 | path, |
| 438 | ) |
| 439 | |
| 440 | # Copy replay values into the buffers recorded for this path. |
| 441 | for key, buf in graph_meta.input_buffers.items(): |
| 442 | src = replay.values.get(key) |
| 443 | if src is None: |
| 444 | continue |
| 445 | if src.ndim == 0: |
| 446 | buf.copy_(src) |
| 447 | else: |
| 448 | padding_logic = self.config.padding_logics.get( |
| 449 | key, self._copy_padded_buffer |
| 450 | ) |
| 451 | padding_logic(buf, src) |
| 452 | |
| 453 | graph_meta.graph.replay() |
| 454 | |
| 455 | self.graph_hits += num_items |
| 456 | return graph_meta.output_buffer |
| 457 | |
| 458 | def _execute_local( |
| 459 | self, |
| 460 | mm_kwargs: dict[str, Any], |
| 461 | ) -> list[torch.Tensor]: |
| 462 | """Execute locally using greedy packing across all configured paths.""" |
| 463 | item_specs = self._get_item_specs(mm_kwargs) |
| 464 | num_items = len(item_specs) |
| 465 | paths = tuple(self.path_token_budgets) |
| 466 | per_item_path_tokens = { |
| 467 | path: [spec.get_path_output_tokens(path) for spec in item_specs] |
| 468 | for path in paths |
| 469 | } |
| 470 | per_item_out_tokens = [spec.output_tokens for spec in item_specs] |
| 471 | max_path_budgets = {path: max(self.path_token_budgets[path]) for path in paths} |
| 472 | |
| 473 | sorted_indices = sorted(range(num_items), key=lambda i: per_item_out_tokens[i]) |
| 474 | batches: list[tuple[list[int], dict[str, int | None]]] = [] |
| 475 | current_batch: list[int] = [] |
| 476 | current_tokens = dict.fromkeys(paths, 0) |
| 477 | |
| 478 | def append_current_batch() -> None: |
| 479 | if not current_batch: |
| 480 | return |
| 481 | path_budgets = { |
| 482 | path: ( |
| 483 | 0 |
| 484 | if current_tokens[path] == 0 |
| 485 | else self._find_smallest_fitting_budget_given_tokens( |
| 486 | current_tokens[path], self.path_token_budgets[path] |
| 487 | ) |
| 488 | ) |
| 489 | for path in paths |
| 490 | } |
| 491 | batches.append((list(current_batch), path_budgets)) |
| 492 | |
| 493 | for orig_idx in sorted_indices: |
| 494 | item_tokens = {path: per_item_path_tokens[path][orig_idx] for path in paths} |
| 495 | fits = len(current_batch) < self.max_batch_size and all( |
| 496 | current_tokens[path] + item_tokens[path] <= max_path_budgets[path] |
| 497 | for path in paths |
| 498 | ) |
| 499 | if current_batch and not fits: |
| 500 | append_current_batch() |
| 501 | current_batch = [] |
| 502 | current_tokens = dict.fromkeys(paths, 0) |
| 503 | |
| 504 | current_batch.append(orig_idx) |
| 505 | for path in paths: |
| 506 | current_tokens[path] += item_tokens[path] |
| 507 | |
| 508 | append_current_batch() |
| 509 | |
| 510 | outputs_by_orig_idx: dict[int, torch.Tensor] = {} |
| 511 | for batch_indices, path_budgets in batches: |
| 512 | batch_mm_kwargs, axis_keys = self._select_items(mm_kwargs, batch_indices) |
| 513 | graph_outputs: dict[str, torch.Tensor] = {} |
| 514 | all_eager = True |
| 515 | |
| 516 | for path in paths: |
| 517 | path_tokens = sum(per_item_path_tokens[path][i] for i in batch_indices) |
| 518 | if path_tokens == 0: |
| 519 | continue |
| 520 | |
| 521 | token_budget = path_budgets[path] |
| 522 | if token_budget is None: |
| 523 | with torch.inference_mode(): |
| 524 | output = self.model.encoder_eager_forward( |
| 525 | batch_mm_kwargs, path=path |
| 526 | ) |
| 527 | else: |
| 528 | all_eager = False |
| 529 | graph_output = self._run_budget_graph( |
| 530 | batch_mm_kwargs, |
| 531 | token_budget, |
| 532 | path=path, |
| 533 | axis_keys=axis_keys, |
| 534 | ) |
| 535 | assert graph_output is not None |
| 536 | output = graph_output |
| 537 | graph_outputs[path] = output |
| 538 | |
| 539 | if all_eager: |
| 540 | self.graph_misses += len(batch_indices) |
| 541 | |
| 542 | self.model.postprocess_encoder_output( |
| 543 | graph_outputs, |
| 544 | batch_indices, |
| 545 | per_item_out_tokens, |
| 546 | outputs_by_orig_idx, |
| 547 | clone=True, |
| 548 | batch_mm_kwargs=batch_mm_kwargs, |
| 549 | ) |
| 550 | |
| 551 | return [outputs_by_orig_idx[i] for i in range(num_items)] |
| 552 | |
| 553 | def _dp_shard( |
| 554 | self, |
| 555 | mm_kwargs: dict[str, Any], |
| 556 | per_item_out_tokens: list[int], |
| 557 | ) -> tuple[dict[str, Any], list[int], list[int], int]: |
| 558 | """Distribute items across TP ranks for data-parallel execution. |
| 559 | |
| 560 | Uses get_load_balance_assignment() to balance load by input size, |
| 561 | then select_encoder_cudagraph_items() to extract each rank's inputs. |
| 562 | |
| 563 | Returns: |
| 564 | local_mm_kwargs: Inputs for this rank. |
| 565 | image_rank_assignment: Flattened assignment order across all ranks. |
| 566 | images_per_rank: Number of items per rank. |
| 567 | max_output_tokens_per_rank: Max output tokens across all ranks |
| 568 | (for padding during all_gather). |
| 569 | |
| 570 | """ |
| 571 | tp_size = get_tensor_model_parallel_world_size() |
| 572 | current_rank = get_tensor_model_parallel_rank() |
| 573 | |
| 574 | item_specs = self._get_item_specs(mm_kwargs) |
| 575 | per_item_input_sizes = [spec.input_size for spec in item_specs] |
| 576 | |
| 577 | (image_rank_assignment, images_per_rank, input_patches_per_rank) = ( |
| 578 | get_load_balance_assignment(per_item_input_sizes, tp_size) |
| 579 | ) |
| 580 | |
| 581 | # Extract local indices for this rank |
| 582 | cum_images_per_rank = [0] |
| 583 | for count in images_per_rank: |
| 584 | cum_images_per_rank.append(cum_images_per_rank[-1] + count) |
| 585 | |
| 586 | local_indices = image_rank_assignment[ |
| 587 | cum_images_per_rank[current_rank] : cum_images_per_rank[current_rank + 1] |
| 588 | ] |
| 589 | |
| 590 | if len(local_indices) > 0: |
| 591 | local_mm_kwargs = self._select_items(mm_kwargs, local_indices)[0] |
| 592 | else: |
| 593 | local_mm_kwargs = self._select_items(mm_kwargs, [])[0] |
| 594 | |
| 595 | max_output_tokens_per_rank = ( |
| 596 | max( |
| 597 | sum( |
| 598 | per_item_out_tokens[i] |
| 599 | for i in image_rank_assignment[ |
| 600 | cum_images_per_rank[r] : cum_images_per_rank[r + 1] |
| 601 | ] |
| 602 | ) |
| 603 | for r in range(tp_size) |
| 604 | ) |
| 605 | if len(per_item_out_tokens) > 0 |
| 606 | else 0 |
| 607 | ) |
| 608 | |
| 609 | return ( |
| 610 | local_mm_kwargs, |
| 611 | image_rank_assignment, |
| 612 | images_per_rank, |
| 613 | max_output_tokens_per_rank, |
| 614 | ) |
| 615 | |
| 616 | def _dp_gather( |
| 617 | self, |
| 618 | local_outputs: list[torch.Tensor], |
| 619 | per_item_out_tokens: list[int], |
| 620 | image_rank_assignment: list[int], |
| 621 | images_per_rank: list[int], |
| 622 | max_output_tokens_per_rank: int, |
| 623 | ) -> list[torch.Tensor]: |
| 624 | """Gather outputs from all TP ranks and reorder to original sequence. |
| 625 | |
| 626 | Assumes 2D output tensors [tokens, hidden]. Follows the same |
| 627 | pad -> all_gather -> unpad -> reorder algorithm as |
| 628 | run_dp_sharded_mrope_vision_model() in the eager path. |
| 629 | """ |
| 630 | hidden_size = self.config.out_hidden_size |
| 631 | tp_size = len(images_per_rank) |
| 632 | |
| 633 | if len(local_outputs) > 0: |
| 634 | local_concat = torch.cat(local_outputs, dim=0) |
| 635 | else: |
| 636 | local_concat = torch.empty( |
| 637 | (0, hidden_size), device=self.device, dtype=self.dtype |
| 638 | ) |
| 639 | |
| 640 | # Pad to max_output_tokens_per_rank for all_gather |
| 641 | current_len = local_concat.shape[0] |
| 642 | if current_len < max_output_tokens_per_rank: |
| 643 | padding = torch.empty( |
| 644 | (max_output_tokens_per_rank - current_len, hidden_size), |
| 645 | dtype=self.dtype, |
| 646 | device=self.device, |
| 647 | ) |
| 648 | local_padded = torch.cat([local_concat, padding], dim=0) |
| 649 | else: |
| 650 | local_padded = local_concat |
| 651 | |
| 652 | gathered = tensor_model_parallel_all_gather(local_padded, dim=0) |
| 653 | |
| 654 | # Unpad each rank's contribution |
| 655 | rank_outputs: list[torch.Tensor] = [] |
| 656 | current_idx = 0 |
| 657 | for rank in range(tp_size): |
| 658 | start = rank * max_output_tokens_per_rank |
| 659 | rank_count = images_per_rank[rank] |
| 660 | rank_indices = image_rank_assignment[current_idx : current_idx + rank_count] |
| 661 | rank_tokens = sum(per_item_out_tokens[i] for i in rank_indices) |
| 662 | current_idx += rank_count |
| 663 | rank_outputs.append(gathered[start : start + rank_tokens]) |
| 664 | |
| 665 | # Reorder to original sequence |
| 666 | total_items = len(per_item_out_tokens) |
| 667 | result: list[torch.Tensor | None] = [None] * total_items |
| 668 | current_idx = 0 |
| 669 | for rank in range(tp_size): |
| 670 | count = images_per_rank[rank] |
| 671 | if count > 0: |
| 672 | rank_items = image_rank_assignment[current_idx : current_idx + count] |
| 673 | scatter_output_slices( |
| 674 | rank_outputs[rank], |
| 675 | rank_items, |
| 676 | per_item_out_tokens, |
| 677 | result, |
| 678 | ) |
| 679 | current_idx += count |
| 680 | |
| 681 | return [t for t in result if t is not None] |
| 682 | |
| 683 | def execute( |
| 684 | self, |
| 685 | mm_kwargs: dict[str, Any], |
| 686 | ) -> list[torch.Tensor]: |
| 687 | """Execute encoder using CUDA graph with optional DP. |
| 688 | |
| 689 | Args: |
| 690 | mm_kwargs: Multimodal keyword arguments containing the |
| 691 | input tensor and grid dimensions. |
| 692 | |
| 693 | Returns: |
| 694 | List of encoder outputs (one per item). |
| 695 | |
| 696 | """ |
| 697 | if self.use_dp: |
| 698 | per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs) |
| 699 | |
| 700 | ( |
| 701 | local_mm_kwargs, |
| 702 | image_rank_assignment, |
| 703 | images_per_rank, |
| 704 | max_output_tokens_per_rank, |
| 705 | ) = self._dp_shard(mm_kwargs, per_item_out_tokens) |
| 706 | |
| 707 | local_outputs = self._execute_local(local_mm_kwargs) |
| 708 | |
| 709 | result = self._dp_gather( |
| 710 | local_outputs, |
| 711 | per_item_out_tokens, |
| 712 | image_rank_assignment, |
| 713 | images_per_rank, |
| 714 | max_output_tokens_per_rank, |
| 715 | ) |
| 716 | else: |
| 717 | result = self._execute_local(mm_kwargs) |
| 718 | |
| 719 | # Log cumulative stats periodically |
| 720 | stats = self.get_cumulative_stats() |
| 721 | total_requests = self.graph_hits + self.graph_misses |
| 722 | if total_requests > 0 and total_requests % self.log_stats_interval == 0: |
| 723 | logger.debug( |
| 724 | "Encoder CUDA graph cumulative stats: " |
| 725 | "hits=%d, misses=%d, hit_rate=%.1f%%", |
| 726 | stats["graph_hits"], |
| 727 | stats["graph_misses"], |
| 728 | stats["hit_rate"] * 100, |
| 729 | ) |
| 730 | |
| 731 | return result |
| 732 | |
| 733 | def get_cumulative_stats(self) -> dict[str, Any]: |
| 734 | """Get cumulative CUDA graph statistics.""" |
| 735 | total_requests = self.graph_hits + self.graph_misses |
| 736 | hit_rate = self.graph_hits / total_requests if total_requests > 0 else 0.0 |
| 737 | |
| 738 | num_budgets = sum(len(g) for g in self.budget_graphs.values()) |
| 739 | |
| 740 | return { |
| 741 | "graph_hits": self.graph_hits, |
| 742 | "graph_misses": self.graph_misses, |
| 743 | "hit_rate": hit_rate, |
| 744 | "num_budgets": num_budgets, |
| 745 | "token_budgets": self.token_budgets, |
| 746 | } |