Skip to content

vllm.v1.worker.encoder_cudagraph_defs

Data transfer objects for encoder CUDA graph management.

Classes:

EncoderCudaGraphCaptureInputs dataclass

Everything needed for one CUDA graph capture.

Returned by prepare_encoder_cudagraph_capture_inputs().

Attributes:

Source code in vllm/v1/worker/encoder_cudagraph_defs.py 
106
107
108
109
110
111
112
113
114
115
116
117
@dataclass
class EncoderCudaGraphCaptureInputs:
    """Everything needed for one CUDA graph capture.

    Returned by ``prepare_encoder_cudagraph_capture_inputs()``.
    """

    values: dict[str, torch.Tensor]
    """Precomputed tensor buffers that will be recorded into the
    CUDA graph.  The manager stores references to these exact
    tensor objects and copies new data into them before each
    ``graph.replay()`` call (buffer identity invariant)."""

values instance-attribute

Precomputed tensor buffers that will be recorded into the CUDA graph. The manager stores references to these exact tensor objects and copies new data into them before each graph.replay() call (buffer identity invariant).

EncoderCudaGraphConfig dataclass

Configuration for encoder CUDA graph management.

Provided by the model at init time via get_encoder_cudagraph_config(). Values are fixed for the lifetime of the manager.

Attributes:

Source code in vllm/v1/worker/encoder_cudagraph_defs.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
@dataclass
class EncoderCudaGraphConfig:
    """Configuration for encoder CUDA graph management.

    Provided by the model at init time via
    ``get_encoder_cudagraph_config()``. Values are fixed for the
    lifetime of the manager.
    """

    modalities: list[str]
    """Supported modalities (e.g. ["image"])."""

    buffer_keys: list[str]
    """Keys for the tensor buffers recorded into the CUDA graph.
    Before replay the manager zeros then slice-copies new data
    into these buffers."""

    out_hidden_size: int
    """Output hidden dim of the vision encoder.
    Used for DP gather buffer allocation."""

    padding_logics: dict[str, EncoderCudaGraphPaddingLogic] = field(
        default_factory=dict
    )
    """Optional per-buffer replay padding/copy logic.
    If absent for a key, the manager zeros the capture buffer and slice-copies
    the replay buffer into it."""

    max_frames_per_video: int = 1
    """Maximum number of frames per video.
    Only relevant when "video" is in ``modalities``.
    Image-only models can use the default of 1."""

    paths: dict[str, EncoderCudaGraphPathConfig] = field(
        default_factory=lambda: {"default": EncoderCudaGraphPathConfig()}
    )
    """Independently captured encoder paths keyed by their forward name."""

    capture_axes: tuple[tuple[Hashable, ...], ...] = ()
    """Extra capture axes beyond the token budget; empty to disable.

    Each entry is the ordered key set of one axis. When non-empty, one graph
    is captured per token budget per combination of axis keys (cartesian
    product), so the total number of captured graphs is
    ``num_budgets * prod(len(axis) for axis in capture_axes)``."""

buffer_keys instance-attribute

Keys for the tensor buffers recorded into the CUDA graph. Before replay the manager zeros then slice-copies new data into these buffers.

capture_axes = () class-attribute instance-attribute

Extra capture axes beyond the token budget; empty to disable.

Each entry is the ordered key set of one axis. When non-empty, one graph is captured per token budget per combination of axis keys (cartesian product), so the total number of captured graphs is num_budgets * prod(len(axis) for axis in capture_axes).

max_frames_per_video = 1 class-attribute instance-attribute

Maximum number of frames per video. Only relevant when "video" is in modalities. Image-only models can use the default of 1.

modalities instance-attribute

Supported modalities (e.g. ["image"]).

out_hidden_size instance-attribute

Output hidden dim of the vision encoder. Used for DP gather buffer allocation.

padding_logics = field(default_factory=dict) class-attribute instance-attribute

Optional per-buffer replay padding/copy logic. If absent for a key, the manager zeros the capture buffer and slice-copies the replay buffer into it.

paths = field(default_factory=lambda: {'default': EncoderCudaGraphPathConfig()}) class-attribute instance-attribute

Independently captured encoder paths keyed by their forward name.

EncoderCudaGraphPathConfig dataclass

Capture policy for one independently replayable encoder path.

Attributes:

Source code in vllm/v1/worker/encoder_cudagraph_defs.py 
48
49
50
51
52
53
54
55
56
@dataclass(frozen=True)
class EncoderCudaGraphPathConfig:
    """Capture policy for one independently replayable encoder path."""

    min_token_budget: int | None = None
    """Smallest capture budget, or the model default minimum when unset."""

    allow_zero_tokens: bool = False
    """Whether a batch may omit this path entirely."""

allow_zero_tokens = False class-attribute instance-attribute

Whether a batch may omit this path entirely.

min_token_budget = None class-attribute instance-attribute

Smallest capture budget, or the model default minimum when unset.

EncoderCudaGraphReplayBuffers dataclass

New buffer values for graph replay, computed by the model from actual batch inputs.

Returned by prepare_encoder_cudagraph_replay_buffers(). Keys match EncoderCudaGraphConfig.buffer_keys.

Attributes:

Source code in vllm/v1/worker/encoder_cudagraph_defs.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
@dataclass
class EncoderCudaGraphReplayBuffers:
    """New buffer values for graph replay, computed by the model from
    actual batch inputs.

    Returned by ``prepare_encoder_cudagraph_replay_buffers()``.
    Keys match ``EncoderCudaGraphConfig.buffer_keys``.
    """

    values: dict[str, torch.Tensor | None]
    """Data to copy into the captured buffers before replay.
    ``None`` values leave the corresponding captured buffer
    unchanged."""

values instance-attribute

Data to copy into the captured buffers before replay. None values leave the corresponding captured buffer unchanged.

EncoderItemSpec dataclass

Description of a single encoder input item.

Returned by get_encoder_cudagraph_item_specs() to describe each image or video in a batch without the manager needing to understand model-specific input formats.

Attributes:

Source code in vllm/v1/worker/encoder_cudagraph_defs.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
@dataclass
class EncoderItemSpec:
    """Description of a single encoder input item.

    Returned by ``get_encoder_cudagraph_item_specs()`` to describe each
    image or video in a batch without the manager needing to understand
    model-specific input formats.
    """

    input_size: int
    """Number of input patches/rows for this item."""

    output_tokens: int
    """Number of output tokens after encoder processing (e.g. after
    spatial merge)."""

    path_output_tokens: dict[str, int] = field(default_factory=dict)
    """Per-path output token counts for multi-path encoders.

    Single-path encoders leave this empty and use ``output_tokens`` for the
    default path.
    """

    def get_path_output_tokens(self, path: str) -> int:
        if path == "default" and not self.path_output_tokens:
            return self.output_tokens
        return self.path_output_tokens.get(path, 0)

input_size instance-attribute

Number of input patches/rows for this item.

output_tokens instance-attribute

Number of output tokens after encoder processing (e.g. after spatial merge).

path_output_tokens = field(default_factory=dict) class-attribute instance-attribute

Per-path output token counts for multi-path encoders.

Single-path encoders leave this empty and use output_tokens for the default path.