| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Data transfer objects for encoder CUDA graph management.""" |
| 4 | |
| 5 | from collections.abc import Callable, Hashable |
| 6 | from dataclasses import dataclass, field |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | EncoderCudaGraphPaddingLogic = Callable[[torch.Tensor, torch.Tensor], None] |
| 11 | |
| 12 | # Reserved mm_kwargs key: models with ``EncoderCudaGraphConfig.capture_axes`` |
| 13 | # put the resolved per-axis keys (one per axis, in order) into the dict |
| 14 | # returned by ``select_encoder_cudagraph_items()``; the manager pops it before |
| 15 | # the kwargs are used anywhere else. |
| 16 | ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG = "encoder_cudagraph_axis_keys" |
| 17 | |
| 18 | |
| 19 | @dataclass |
| 20 | class EncoderItemSpec: |
| 21 | """Description of a single encoder input item. |
| 22 | |
| 23 | Returned by ``get_encoder_cudagraph_item_specs()`` to describe each |
| 24 | image or video in a batch without the manager needing to understand |
| 25 | model-specific input formats. |
| 26 | """ |
| 27 | |
| 28 | input_size: int |
| 29 | """Number of input patches/rows for this item.""" |
| 30 | |
| 31 | output_tokens: int |
| 32 | """Number of output tokens after encoder processing (e.g. after |
| 33 | spatial merge).""" |
| 34 | |
| 35 | path_output_tokens: dict[str, int] = field(default_factory=dict) |
| 36 | """Per-path output token counts for multi-path encoders. |
| 37 | |
| 38 | Single-path encoders leave this empty and use ``output_tokens`` for the |
| 39 | default path. |
| 40 | """ |
| 41 | |
| 42 | def get_path_output_tokens(self, path: str) -> int: |
| 43 | if path == "default" and not self.path_output_tokens: |
| 44 | return self.output_tokens |
| 45 | return self.path_output_tokens.get(path, 0) |
| 46 | |
| 47 | |
| 48 | @dataclass(frozen=True) |
| 49 | class EncoderCudaGraphPathConfig: |
| 50 | """Capture policy for one independently replayable encoder path.""" |
| 51 | |
| 52 | min_token_budget: int | None = None |
| 53 | """Smallest capture budget, or the model default minimum when unset.""" |
| 54 | |
| 55 | allow_zero_tokens: bool = False |
| 56 | """Whether a batch may omit this path entirely.""" |
| 57 | |
| 58 | |
| 59 | @dataclass |
| 60 | class EncoderCudaGraphConfig: |
| 61 | """Configuration for encoder CUDA graph management. |
| 62 | |
| 63 | Provided by the model at init time via |
| 64 | ``get_encoder_cudagraph_config()``. Values are fixed for the |
| 65 | lifetime of the manager. |
| 66 | """ |
| 67 | |
| 68 | modalities: list[str] |
| 69 | """Supported modalities (e.g. ["image"]).""" |
| 70 | |
| 71 | buffer_keys: list[str] |
| 72 | """Keys for the tensor buffers recorded into the CUDA graph. |
| 73 | Before replay the manager zeros then slice-copies new data |
| 74 | into these buffers.""" |
| 75 | |
| 76 | out_hidden_size: int |
| 77 | """Output hidden dim of the vision encoder. |
| 78 | Used for DP gather buffer allocation.""" |
| 79 | |
| 80 | padding_logics: dict[str, EncoderCudaGraphPaddingLogic] = field( |
| 81 | default_factory=dict |
| 82 | ) |
| 83 | """Optional per-buffer replay padding/copy logic. |
| 84 | If absent for a key, the manager zeros the capture buffer and slice-copies |
| 85 | the replay buffer into it.""" |
| 86 | |
| 87 | max_frames_per_video: int = 1 |
| 88 | """Maximum number of frames per video. |
| 89 | Only relevant when "video" is in ``modalities``. |
| 90 | Image-only models can use the default of 1.""" |
| 91 | |
| 92 | paths: dict[str, EncoderCudaGraphPathConfig] = field( |
| 93 | default_factory=lambda: {"default": EncoderCudaGraphPathConfig()} |
| 94 | ) |
| 95 | """Independently captured encoder paths keyed by their forward name.""" |
| 96 | |
| 97 | capture_axes: tuple[tuple[Hashable, ...], ...] = () |
| 98 | """Extra capture axes beyond the token budget; empty to disable. |
| 99 | |
| 100 | Each entry is the ordered key set of one axis. When non-empty, one graph |
| 101 | is captured per token budget per combination of axis keys (cartesian |
| 102 | product), so the total number of captured graphs is |
| 103 | ``num_budgets * prod(len(axis) for axis in capture_axes)``.""" |
| 104 | |
| 105 | |
| 106 | @dataclass |
| 107 | class EncoderCudaGraphCaptureInputs: |
| 108 | """Everything needed for one CUDA graph capture. |
| 109 | |
| 110 | Returned by ``prepare_encoder_cudagraph_capture_inputs()``. |
| 111 | """ |
| 112 | |
| 113 | values: dict[str, torch.Tensor] |
| 114 | """Precomputed tensor buffers that will be recorded into the |
| 115 | CUDA graph. The manager stores references to these exact |
| 116 | tensor objects and copies new data into them before each |
| 117 | ``graph.replay()`` call (buffer identity invariant).""" |
| 118 | |
| 119 | |
| 120 | @dataclass |
| 121 | class EncoderCudaGraphReplayBuffers: |
| 122 | """New buffer values for graph replay, computed by the model from |
| 123 | actual batch inputs. |
| 124 | |
| 125 | Returned by ``prepare_encoder_cudagraph_replay_buffers()``. |
| 126 | Keys match ``EncoderCudaGraphConfig.buffer_keys``. |
| 127 | """ |
| 128 | |
| 129 | values: dict[str, torch.Tensor | None] |
| 130 | """Data to copy into the captured buffers before replay. |
| 131 | ``None`` values leave the corresponding captured buffer |
| 132 | unchanged.""" |