[1/N] Unify multiple-path encoder cuda graph support (#49934)
Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Isotr0py committed
Aug 2, 2026 at 10:33 UTC
e2fa28594f7baad142a426b0b6a2cfe2c79201c7
7 files changed
+192
-387
docs/design/cuda_graphs_multimodal.md
+16
-25
@@ -41,13 +41,13 @@ class BudgetGraphMetadata:
41
42
Budgets are auto-generated as power-of-2 levels from a model-provided range via `get_encoder_cudagraph_budget_range()`, with the maximum budget always included even if it does not fall on a power-of-2 boundary. Budgets can also be explicitly specified by the user via `encoder_cudagraph_token_budgets` in `CompilationConfig`.
43
44
-When `EncoderCudaGraphConfig.enable_dual_path_graph` is `True`, the manager generates two independent budget lists — `global_token_budgets` (multiples of `global_token_per_image`) and `local_token_budgets` (multiples of `local_token_per_patch`) — and stores captured graphs under `budget_graphs["global"]` and `budget_graphs["local"]` respectively.
44
+Each entry in `EncoderCudaGraphConfig.paths` defines an independently captured encoder path. A path can provide its own minimum token budget and opt into zero-token batches; the manager generates and stores a separate budget graph set for every configured path.
45
46
### Greedy bin-packing at runtime
47
48
When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution.
49
50
-For dual-path models, the manager routes to `_execute_local_dual_path()`, which constrains both global and local token budgets simultaneously during packing (see [Dual-Path graph capture](#dual-path-graph-capture)).
50
+For multi-path models, the same greedy packing loop constrains every configured path simultaneously (see [Multi-Path graph capture](#multi-path-graph-capture)).
51
52
For each graph replay:
53
@@ -56,38 +56,29 @@ For each graph replay:
56
3. Replay the CUDA Graph.
57
4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays).
58
59
-### Dual-Path graph capture
59
+### Multi-Path graph capture
60
61
-For two-tower vision encoders (e.g., DeepSeek-OCR), the `EncoderCudaGraphConfig` sets `enable_dual_path_graph=True` and provides `global_token_per_image` / `local_token_per_patch`. The manager captures two independent sets of CUDA graphs — one for the **global** image path and one for the **local** patch path — stored under `budget_graphs["global"]` and `budget_graphs["local"]` respectively.
61
+`EncoderCudaGraphConfig.paths` maps path names to `EncoderCudaGraphPathConfig` capture policies. For example, DeepSeek-OCR configures a **global** image path and a **local** patch path, which are captured independently under `budget_graphs["global"]` and `budget_graphs["local"]`.
62
63
-**Budget generation.** Two separate budget lists are generated:
63
+**Budget generation.** Each path gets a separate budget list. For DeepSeek-OCR:
64
65
-* `global_token_budgets` — power-of-2 multiples of `global_token_per_image` (e.g., `[272, 544, 1088, 2176, 4352, 8704, 13824]` for DeepSeek-OCR).
66
-* `local_token_budgets` — power-of-2 multiples of `local_token_per_patch` (e.g., `[0, 100, 200, 400, 800, 1600, 3200, 6400, 12800]` for DeepSeek-OCR). A budget of `0` is always included to handle images with no local patches (images ≤ 640×640 that produce only global features).
65
+* the `global` path — power-of-2 budgets starting at the global path minimum (e.g., `[272, 544, 1088, 2176, 4352, 8704, 13824]` for DeepSeek-OCR).
66
+* the `local` path — power-of-2 budgets starting at the local path minimum (e.g., `[0, 100, 200, 400, 800, 1600, 3200, 6400, 12800, 13824]` for DeepSeek-OCR). A budget of `0` is included when `allow_zero_tokens=True` to handle images with no local patches (images ≤ 640×640 that produce only global features).
67
68
Both lists are capped at the same `max_budget`.
69
70
-**Dual-path greedy packing.** Each `EncoderItemSpec` provides both `global_output_tokens` (constant per image) and `local_output_tokens` (proportional to the patch count). The dual-path packing algorithm constrains both budgets simultaneously:
70
+**Multi-path greedy packing.** Each `EncoderItemSpec` provides `path_output_tokens`, mapping each path to its contribution for that item. The packing algorithm constrains every path simultaneously:
71
72
* Sort images by total output tokens (global + local), smallest first.
73
-* Greedily pack images: an image is added to the current sub-batch only if both the accumulated global tokens ≤ `max_global_budget` **and** the accumulated local tokens ≤ `max_local_budget`, with the image count ≤ `max_batch_size`.
74
-* Once either constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path.
73
+* Greedily pack images: an image is added to the current sub-batch only if every accumulated path token count is within that path's maximum budget, with the image count ≤ `max_batch_size`.
74
+* Once any path constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path.
75
* Repeat until all images are packed.
76
77
-**Partial graph fallback.** After packing, each sub-batch falls into one of four execution scenarios:
78
-
79
-| Global budget | Local budget | Execution |
80
-| :---: | :---: | --- |
81
-| Found | Found | Both paths use CUDA graph replay |
82
-| Found | `None` | Global graph replay + local path skipped (no patches) |
83
-| `None` | Found | Global eager fallback + local graph replay |
84
-| `None` | `None` | Both paths fall back to eager execution |
85
-
86
-Note that the `0`-budget graph is never actually replayed for local — it signals that local patch processing should be skipped entirely.
77
+**Partial graph fallback.** For each non-empty path, the manager replays the smallest fitting graph or runs only that path eagerly when no graph fits. Paths with zero tokens are skipped; a `0`-budget graph is never captured or replayed.
78
79
**Buffer keys per path.** Global and local paths use different buffer keys. For DeepSeek-OCR, the global path uses `pixel_values` (full images, shape `[B, 3, 1280, 1280]`) while the local path uses `images_crop` (patches, shape `[P, 3, 1024, 1024]`). The manager iterates over each captured graph's own `input_buffers.keys()` rather than a shared `buffer_keys` list, so both paths can use different buffers.
80
90
-**Post-processing.** The `postprocess_encoder_output` method receives a `local_output` parameter (a tensor or `None`) containing the local-path encoder output. The model is responsible for assembling global and local features into the final per-image embedding. For DeepSeek-OCR, this means reshaping the global output into `[B, 272, n_embed]`, the local output into `[P, 100, n_embed]`, assembling patch grids with newline tokens, and concatenating `[patches_grid, global, view_separator]` for each image.
81
+**Post-processing.** The `postprocess_encoder_output` method receives an `outputs` dictionary keyed by path name. The model is responsible for assembling global and local features into the final per-image embedding. For DeepSeek-OCR, it reshapes the global output into `[B, 272, n_embed]`, reshapes the local output into `[P, 100, n_embed]`, assembles patch grids with newline tokens, and concatenating `[patches_grid, global, view_separator]` for each image.
82
83
!!! note
84
The dual-path design enables partial CUDA graph coverage — one path can hit while the other falls back to eager. This avoids wasted compute on zero-padded patch buffers for untiled images and avoids graph invalidation caused by variable `crop_shape` per image.
@@ -111,13 +102,13 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
102
103
* `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video).
104
* `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets.
114
-* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size, total output token count (`output_tokens`), and optionally per-path token counts (`global_output_tokens`, `local_output_tokens`) for dual-path models.
105
+* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size, total output token count (`output_tokens`), and per-path token counts (`path_output_tokens`) for multi-path models.
106
* `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding.
107
* `prepare_encoder_cudagraph_capture_inputs(..., path="default")` — creates dummy inputs for graph capture. The `path` parameter (`"global"` or `"local"`) tells the model which path to generate dummy inputs for. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph.
108
* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch, path="default")` — computes buffer values from actual batch inputs. The `path` parameter selects which modality keys to extract from `mm_kwargs`. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match the captured graph's `input_buffers.keys()`.
109
* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor], path="default")` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `path` parameter dispatches to the correct encoder sub-module (e.g., global vs. local path for DeepSeek-OCR).
110
* `encoder_eager_forward(mm_kwargs, path="default")` — fallback eager forward when no graph fits. When `path` is `"global"` or `"local"`, runs only that encoder path without graph capture.
120
-* `postprocess_encoder_output(..., local_output=None)` — post-process encoder output. The `local_output` parameter receives the local-path encoder output tensor (or `None`), enabling dual-path models to assemble global and local features into the final per-image embedding.
111
+* `postprocess_encoder_output(outputs, ...)` — post-process encoder outputs keyed by path name, enabling multi-path models to assemble their path-specific features into the final per-item embedding.
112
113
!!! note
114
The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
@@ -146,14 +137,14 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
137
138
## Configuration
139
149
-Three fields in `CompilationConfig` control encoder CUDA Graphs:
140
+Four fields in `CompilationConfig` control encoder CUDA Graphs:
141
142
* `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
143
* `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
144
* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
145
* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value from `EncoderCudaGraphConfig`, computed by `get_max_frames_per_video()` on the model). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).
146
156
-Dual-path mode is configured at the model level via `EncoderCudaGraphConfig` fields (`enable_dual_path_graph`, `global_token_per_image`, `local_token_per_patch`) — no additional user configuration is required. The manager automatically generates separate budget lists and routes to dual-path execution when the model opts in.
147
+Multi-path mode is configured at the model level through `EncoderCudaGraphConfig.paths`. Each `EncoderCudaGraphPathConfig` can set a path-specific minimum budget and whether zero-token batches are allowed. The manager automatically generates separate budget lists and uses the same execution loop for single- and multi-path models.
148
149
## Usage guide
150
tests/v1/cudagraph/test_encoder_cudagraph.py
+2
@@ -106,6 +106,7 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
106
"""
107
mgr = object.__new__(EncoderCudaGraphManager)
108
mgr.token_budgets = sorted(budgets)
109
+ mgr.path_token_budgets = {"default": mgr.token_budgets}
110
mgr.max_batch_size = 16
111
mgr.use_dp = False
112
mgr.config = EncoderCudaGraphConfig(
@@ -417,6 +418,7 @@ def _make_manager_for_gpu(
418
"""Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests."""
419
mgr = object.__new__(EncoderCudaGraphManager)
420
mgr.token_budgets = sorted(token_budgets)
421
+ mgr.path_token_budgets = {"default": mgr.token_budgets}
422
mgr.max_batch_size = max_batch_size
423
mgr.max_frames_per_batch = (
424
max_frames_per_batch if max_frames_per_batch is not None else max_batch_size * 2
vllm/model_executor/models/deepseek_ocr.py
+22
-11
@@ -65,6 +65,7 @@ from vllm.v1.sample.logits_processor import (
65
from vllm.v1.worker.encoder_cudagraph_defs import (
66
EncoderCudaGraphCaptureInputs,
67
EncoderCudaGraphConfig,
68
+ EncoderCudaGraphPathConfig,
69
EncoderCudaGraphReplayBuffers,
70
EncoderItemSpec,
71
)
@@ -697,9 +698,15 @@ class DeepseekOCRForCausalLM(
698
modalities=["image"],
699
buffer_keys=["pixel_values"],
700
out_hidden_size=self.projector_config.n_embed,
700
- enable_dual_path_graph=True,
701
- global_token_per_image=self.global_image_output_token,
702
- local_token_per_patch=self.single_patch_output_token,
701
+ paths={
702
+ "global": EncoderCudaGraphPathConfig(
703
+ min_token_budget=self.global_image_output_token
704
+ ),
705
+ "local": EncoderCudaGraphPathConfig(
706
+ min_token_budget=self.single_patch_output_token,
707
+ allow_zero_tokens=True,
708
+ ),
709
+ },
710
)
711
712
def get_encoder_cudagraph_budget_range(
@@ -730,8 +737,10 @@ class DeepseekOCRForCausalLM(
737
EncoderItemSpec(
738
input_size=num_input_tokens,
739
output_tokens=num_output_tokens,
733
- global_output_tokens=global_output_token,
734
- local_output_tokens=local_output_token,
740
+ path_output_tokens={
741
+ "global": global_output_token,
742
+ "local": local_output_token,
743
+ },
744
)
745
)
746
return item_specs
@@ -917,32 +926,34 @@ class DeepseekOCRForCausalLM(
926
927
def postprocess_encoder_output(
928
self,
920
- output: torch.Tensor,
929
+ outputs: dict[str, torch.Tensor],
930
indices: list[int],
931
per_item_out_tokens: list[int],
932
dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
933
clone: bool = False,
934
batch_mm_kwargs: dict[str, Any] | None = None,
926
- local_output: torch.Tensor | None = None,
935
) -> None:
936
"""
937
Assemble per-image embeddings from global and local encoder outputs.
938
931
- ``output`` contains global-image features with newlines already
939
+ ``output['global']`` contains global-image features with newlines already
940
inserted (from CUDA graph replay or eager fallback):
941
``[B * 272, n_embed]``.
942
935
- ``local_output`` contains local-patch features without
943
+ ``output['local']`` contains local-patch features without
944
newlines (from CUDA graph replay or eager fallback):
945
``[P * 100, n_embed]``. May be ``None`` if no patches in batch.
946
947
This method:
940
- 1. Splits ``output`` into per-image global portions.
941
- 2. Splits ``local_output`` into per-image patch groups.
948
+ 1. Splits ``output['global']`` into per-image global portions.
949
+ 2. Splits ``output['local']`` into per-image patch groups.
950
3. For each image: assembles patch grid with newlines via
951
``_assemble_patch_grid``, then concatenates
952
``[local_tiled, global, view_seperator]``.
953
"""
954
+ output = outputs["global"]
955
+ local_output = outputs.get("local")
956
+ assert batch_mm_kwargs is not None
957
bsz = len(indices)
958
n_embed = output.shape[-1]
959
vllm/model_executor/models/interfaces.py
+4
-3
@@ -1696,13 +1696,12 @@ class SupportsEncoderCudaGraph(Protocol):
1696
1697
def postprocess_encoder_output(
1698
self,
1699
- output: torch.Tensor,
1699
+ outputs: dict[str, torch.Tensor],
1700
indices: list[int],
1701
per_item_out_tokens: list[int],
1702
dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
1703
clone: bool = False,
1704
batch_mm_kwargs: dict[str, Any] | None = None,
1705
- local_output: torch.Tensor | None = None,
1705
) -> None:
1706
"""
1707
Post-process encoder output, directly call scatter_output_slices by default.
@@ -1714,7 +1713,9 @@ class SupportsEncoderCudaGraph(Protocol):
1713
"""
1714
from vllm.model_executor.models.utils import scatter_output_slices
1715
1717
- scatter_output_slices(output, indices, per_item_out_tokens, dest, clone)
1716
+ scatter_output_slices(
1717
+ outputs["default"], indices, per_item_out_tokens, dest, clone
1718
+ )
1719
1720
def prepare_encoder_cudagraph_capture_inputs(
1721
self,
vllm/model_executor/models/step3_vl.py
+20
-9
@@ -725,6 +725,7 @@ class Step3VLForConditionalGeneration(
725
def get_encoder_cudagraph_config(self):
726
from vllm.v1.worker.encoder_cudagraph_defs import (
727
EncoderCudaGraphConfig,
728
+ EncoderCudaGraphPathConfig,
729
)
730
731
return EncoderCudaGraphConfig(
@@ -734,9 +735,15 @@ class Step3VLForConditionalGeneration(
735
"patch_pixel_values",
736
],
737
out_hidden_size=self.config.hidden_size,
737
- enable_dual_path_graph=True,
738
- global_token_per_image=self.img_output_tokens,
739
- local_token_per_patch=self.patch_output_tokens,
738
+ paths={
739
+ "global": EncoderCudaGraphPathConfig(
740
+ min_token_budget=self.img_output_tokens
741
+ ),
742
+ "local": EncoderCudaGraphPathConfig(
743
+ min_token_budget=self.patch_output_tokens,
744
+ allow_zero_tokens=True,
745
+ ),
746
+ },
747
)
748
749
def get_encoder_cudagraph_budget_range(
@@ -771,8 +778,10 @@ class Step3VLForConditionalGeneration(
778
output_tokens=(
779
self.img_output_tokens + num_patch * self.patch_output_tokens
780
),
774
- global_output_tokens=self.img_output_tokens,
775
- local_output_tokens=num_patch * self.patch_output_tokens,
781
+ path_output_tokens={
782
+ "global": self.img_output_tokens,
783
+ "local": num_patch * self.patch_output_tokens,
784
+ },
785
)
786
for num_patch in num_patches
787
]
@@ -876,19 +885,21 @@ class Step3VLForConditionalGeneration(
885
886
def postprocess_encoder_output(
887
self,
879
- output: torch.Tensor,
888
+ outputs: dict[str, torch.Tensor],
889
indices: list[int],
890
per_item_out_tokens: list[int],
891
dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
892
clone: bool = False,
893
batch_mm_kwargs: dict[str, Any] | None = None,
885
- local_output: torch.Tensor | None = None,
894
):
895
"""CPU-side per-item merge after dual-path graph replay.
896
889
- ``output`` contains global-image features and ``local_output``
897
+ ``outputs['global']`` contains global-image features and ``outputs['local']``
898
contains local-patch features (or ``None`` when there are no patches).
899
"""
900
+ output = outputs["global"]
901
+ local_output = outputs.get("local")
902
+ assert batch_mm_kwargs is not None
903
num_patches = batch_mm_kwargs["num_patches"]
904
hidden = output.shape[-1]
905
bsz = len(indices)
@@ -899,7 +910,7 @@ class Step3VLForConditionalGeneration(
910
patch_tokens = total_patches * self.patch_output_tokens
911
912
global_part = output[:img_tokens].reshape(bsz, self.img_output_tokens, hidden)
902
- if total_patches > 0:
913
+ if total_patches > 0 and local_output is not None:
914
patch_part = local_output[:patch_tokens].reshape(
915
-1, self.patch_output_tokens, hidden
916
)
vllm/v1/worker/encoder_cudagraph.py
+103
-322
@@ -158,34 +158,33 @@ class EncoderCudaGraphManager:
158
self.graph_misses = 0
159
self.log_stats_interval = 100
160
161
- if self.config.enable_dual_path_graph:
162
- max_budget = self.token_budgets[-1]
163
- self.global_token_budgets = self._generate_budgets(
164
- self.config.global_token_per_image,
165
- max_budget,
166
- )
167
- self.local_token_budgets = self._generate_budgets(
168
- self.config.local_token_per_patch,
169
- max_budget,
170
- )
171
- # When `image_width <= 640 and image_height <= 640`, the mm inputs
172
- # will only contain global image, without generating local patches.
173
- self.local_token_budgets.insert(0, 0)
174
- logger.info(
175
- "EncoderCudaGraphManager dual-path mode: "
176
- "global_budgets=%s, local_budgets=%s",
177
- self.global_token_budgets,
178
- self.local_token_budgets,
179
- )
180
- else:
181
- logger.info(
182
- "EncoderCudaGraphManager initialized with "
183
- "budgets=%s, max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",
184
- self.token_budgets,
185
- self.max_batch_size,
186
- self.max_frames_per_batch,
187
- self.use_dp,
188
- )
161
+ max_budget = self.token_budgets[-1]
162
+ if not self.config.paths:
163
+ raise ValueError("Encoder CUDA graph config must define at least one path")
164
+ self.path_token_budgets: dict[str, list[int]] = {}
165
+ for path, path_config in self.config.paths.items():
166
+ min_path_budget = path_config.min_token_budget
167
+ if min_path_budget is None:
168
+ budgets = list(self.token_budgets)
169
+ else:
170
+ if min_path_budget <= 0 or min_path_budget > max_budget:
171
+ raise ValueError(
172
+ f"Invalid minimum budget {min_path_budget} for encoder "
173
+ f"CUDA graph path {path!r}; max budget is {max_budget}"
174
+ )
175
+ budgets = self._generate_budgets(min_path_budget, max_budget)
176
+ if path_config.allow_zero_tokens:
177
+ budgets.insert(0, 0)
178
+ self.path_token_budgets[path] = budgets
179
+
180
+ logger.info(
181
+ "EncoderCudaGraphManager initialized with paths=%s, "
182
+ "max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",
183
+ self.path_token_budgets,
184
+ self.max_batch_size,
185
+ self.max_frames_per_batch,
186
+ self.use_dp,
187
+ )
188
189
@staticmethod
190
def _generate_budgets(min_budget: int, max_budget: int) -> list[int]:
@@ -211,40 +210,29 @@ class EncoderCudaGraphManager:
210
self.graph_pool = None
211
212
def capture(self, graph_pool: Any):
214
- """Capture CUDA graphs for all token budgets."""
213
+ """Capture CUDA graphs for every configured path and token budget."""
214
self.graph_pool = graph_pool
215
217
- if self.config.enable_dual_path_graph:
218
- for token_budget in sorted(self.global_token_budgets, reverse=True):
219
- self._capture_budget_graph(token_budget, path="global")
220
- for token_budget in sorted(self.local_token_budgets, reverse=True):
221
- if token_budget == 0:
222
- continue
223
- self._capture_budget_graph(token_budget, path="local")
224
- logger.info(
225
- "Encoder CUDA graph capture complete. "
226
- "Captured %d global + %d local budget graphs.",
227
- len(self.budget_graphs["global"]),
228
- len(self.budget_graphs["local"]),
229
- )
230
- return
231
-
232
- for token_budget in sorted(self.token_budgets, reverse=True):
233
- self._capture_budget_graph(token_budget)
216
+ for path, budgets in self.path_token_budgets.items():
217
+ for token_budget in sorted(budgets, reverse=True):
218
+ if token_budget > 0:
219
+ self._capture_budget_graph(token_budget, path=path)
220
221
logger.info(
236
- "Encoder CUDA graph capture complete. Captured %d budget graphs.",
237
- len(self.budget_graphs["default"]),
222
+ "Encoder CUDA graph capture complete. Captured %d graphs across %d paths.",
223
+ self.get_num_graphs_to_capture(),
224
+ len(self.path_token_budgets),
225
)
226
227
def get_num_graphs_to_capture(self) -> int:
241
- if self.config.enable_dual_path_graph:
242
- return len(self.global_token_budgets) + len(self.local_token_budgets)
243
- return len(self.token_budgets)
228
+ return sum(
229
+ token_budget > 0
230
+ for budgets in self.path_token_budgets.values()
231
+ for token_budget in budgets
232
+ )
233
234
def _get_graph_set(self, path: str = "default") -> dict[int, BudgetGraphMetadata]:
246
- # Lazy init global/local graph sets for dual-path models, or default graph
247
- # set for single-path models.
235
+ """Return the captured graphs for one encoder path."""
236
if path not in self.budget_graphs:
237
self.budget_graphs[path] = {}
238
return self.budget_graphs[path]
@@ -331,7 +319,7 @@ class EncoderCudaGraphManager:
319
Args:
320
mm_kwargs: Multimodal inputs for the batch.
321
token_budget: Token budget to use.
334
- path: Path for the graph. Should be one of ["default", "global", "local"].
322
+ path: Configured encoder path.
323
Returns:
324
Encoder outputs, or None if graph not captured.
325
"""
@@ -351,9 +339,7 @@ class EncoderCudaGraphManager:
339
path,
340
)
341
354
- # Copy replay buffers into graph input buffers. Iterate over the
355
- # graph's own buffer keys (which may differ per path for dual-path
356
- # models) rather than the global config.buffer_keys.
342
+ # Copy replay values into the buffers recorded for this path.
343
for key, buf in graph_meta.input_buffers.items():
344
src = replay.values.get(key)
345
if src is None:
@@ -375,296 +361,91 @@ class EncoderCudaGraphManager:
361
self,
362
mm_kwargs: dict[str, Any],
363
) -> list[torch.Tensor]:
378
- """Execute encoder on local inputs using greedy-packed CUDA graphs.
379
-
380
- Sort images by output token count (smallest first), then greedily pack
381
- as many images as possible into each batch while staying within
382
- max_budget tokens and max_batch_size. Once a batch is finalised (next
383
- image would overflow either constraint), find the smallest fitting
384
- budget once for that batch.
385
-
386
- For dual-path models (``enable_dual_path_graph=True``), two independent
387
- graph sets are used: one for global images, one for local patches.
388
- Budgets are found independently per path; if only one path fits, the
389
- other falls back to eager via partial fallback.
390
-
391
- By exchange argument, greedy smallest-first packing minimises eager
392
- fallbacks -- any other ordering yields a higher token sum in some batch,
393
- making that batch more likely to exceed the budget.
394
-
395
- Stats note:
396
- graph_hits -- counted inside _run_budget_graph after successful replay.
397
- graph_misses -- counted here for single-image batches where the image
398
- exceeds max_budget. Batches split due to max_batch_size
399
- always satisfy total_tokens <= max_budget and therefore
400
- always find a valid budget (no miss).
401
- """
402
- if self.config.enable_dual_path_graph:
403
- return self._execute_local_dual_path(mm_kwargs)
404
- return self._execute_local_single_path(mm_kwargs)
405
-
406
- def _execute_local_single_path(
407
- self,
408
- mm_kwargs: dict[str, Any],
409
- ) -> list[torch.Tensor]:
410
- """Single-path greedy-packing execution (original behaviour)."""
364
+ """Execute locally using greedy packing across all configured paths."""
365
item_specs = self._get_item_specs(mm_kwargs)
366
num_items = len(item_specs)
413
- max_budget = self.token_budgets[-1]
414
-
367
+ paths = tuple(self.path_token_budgets)
368
+ per_item_path_tokens = {
369
+ path: [spec.get_path_output_tokens(path) for spec in item_specs]
370
+ for path in paths
371
+ }
372
per_item_out_tokens = [spec.output_tokens for spec in item_specs]
373
+ max_path_budgets = {path: max(self.path_token_budgets[path]) for path in paths}
374
417
- # Sort ascending by output token count (smallest first)
375
sorted_indices = sorted(range(num_items), key=lambda i: per_item_out_tokens[i])
419
-
420
- # Greedy pack against max_budget and max_batch_size.
421
- # _find_smallest_fitting_budget_given_tokens is called once per
422
- # finalised batch, not per image.
423
- batches: list[tuple[list[int], int | None]] = []
376
+ batches: list[tuple[list[int], dict[str, int | None]]] = []
377
current_batch: list[int] = []
425
- current_batch_tokens = 0
426
-
427
- for orig_idx in sorted_indices:
428
- item_tokens = per_item_out_tokens[orig_idx]
429
- if (
430
- current_batch_tokens + item_tokens <= max_budget
431
- and len(current_batch) < self.max_batch_size
432
- ):
433
- current_batch.append(orig_idx)
434
- current_batch_tokens += item_tokens
435
- else:
436
- if current_batch:
437
- batches.append(
438
- (
439
- current_batch,
440
- self._find_smallest_fitting_budget_given_tokens(
441
- current_batch_tokens
442
- ),
443
- )
378
+ current_tokens = dict.fromkeys(paths, 0)
379
+
380
+ def append_current_batch() -> None:
381
+ if not current_batch:
382
+ return
383
+ path_budgets = {
384
+ path: (
385
+ 0
386
+ if current_tokens[path] == 0
387
+ else self._find_smallest_fitting_budget_given_tokens(
388
+ current_tokens[path], self.path_token_budgets[path]
389
)
445
- current_batch = [orig_idx]
446
- current_batch_tokens = item_tokens
447
-
448
- if current_batch:
449
- batches.append(
450
- (
451
- current_batch,
452
- self._find_smallest_fitting_budget_given_tokens(
453
- current_batch_tokens
454
- ),
390
)
456
- )
391
+ for path in paths
392
+ }
393
+ batches.append((list(current_batch), path_budgets))
394
458
- # outputs_by_orig_idx maps each original image index to its output
459
- # tensor. Needed because greedy packing reorders images; we restore
460
- # the original order before returning.
461
- outputs_by_orig_idx: dict[int, torch.Tensor] = {}
462
-
463
- for batch_orig_indices, token_budget in batches:
464
- batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
465
- mm_kwargs, batch_orig_indices
395
+ for orig_idx in sorted_indices:
396
+ item_tokens = {path: per_item_path_tokens[path][orig_idx] for path in paths}
397
+ fits = len(current_batch) < self.max_batch_size and all(
398
+ current_tokens[path] + item_tokens[path] <= max_path_budgets[path]
399
+ for path in paths
400
)
467
- batch_out_tokens = sum(per_item_out_tokens[i] for i in batch_orig_indices)
468
-
469
- if token_budget is None:
470
- # Single oversized image: item_tokens > max_budget.
471
- # graph_misses counted here for this eager fallback.
472
- logger.debug(
473
- "Encoder CUDA graph fallback to eager: no budget for "
474
- "%d tokens from %d images",
475
- batch_out_tokens,
476
- len(batch_orig_indices),
477
- )
478
- self.graph_misses += len(batch_orig_indices)
479
- with torch.inference_mode():
480
- raw = self.model.encoder_eager_forward(batch_mm_kwargs)
481
- scatter_output_slices(
482
- raw,
483
- batch_orig_indices,
484
- per_item_out_tokens,
485
- outputs_by_orig_idx,
486
- )
487
- else:
488
- logger.debug(
489
- "Encoder CUDA graph: batch_size=%d, tokens=%d, "
490
- "budget=%d, waste=%.1f%%",
491
- len(batch_orig_indices),
492
- batch_out_tokens,
493
- token_budget,
494
- (token_budget - batch_out_tokens) / token_budget * 100,
495
- )
496
-
497
- # graph_hits counted inside _run_budget_graph after replay.
498
- output = self._run_budget_graph(batch_mm_kwargs, token_budget)
499
- assert output is not None
500
- self.model.postprocess_encoder_output(
501
- output,
502
- batch_orig_indices,
503
- per_item_out_tokens,
504
- outputs_by_orig_idx,
505
- clone=True,
506
- batch_mm_kwargs=batch_mm_kwargs,
507
- )
508
-
509
- # Return in original batch order (caller maps outputs to token positions)
510
- return [outputs_by_orig_idx[i] for i in range(num_items)]
401
+ if current_batch and not fits:
402
+ append_current_batch()
403
+ current_batch = []
404
+ current_tokens = dict.fromkeys(paths, 0)
405
512
- def _execute_local_dual_path(
513
- self,
514
- mm_kwargs: dict[str, Any],
515
- ) -> list[torch.Tensor]:
516
- """Dual-path greedy-packing execution.
517
-
518
- Each image contributes both global tokens (constant per image)
519
- and local tokens (patches * patch_tokens). Greedy packing
520
- respects both budgets independently, then selects the smallest
521
- fitting budget per path with partial eager fallback.
522
- """
523
- item_specs = self._get_item_specs(mm_kwargs)
524
- num_items = len(item_specs)
406
+ current_batch.append(orig_idx)
407
+ for path in paths:
408
+ current_tokens[path] += item_tokens[path]
409
526
- max_global_budget = self.global_token_budgets[-1]
527
- max_local_budget = self.local_token_budgets[-1]
528
-
529
- per_item_global_tokens = [spec.global_output_tokens for spec in item_specs]
530
- per_item_local_tokens = [spec.local_output_tokens for spec in item_specs]
531
- per_item_total_tokens = [spec.output_tokens for spec in item_specs]
532
-
533
- # Sort ascending by total output tokens
534
- sorted_indices = sorted(
535
- range(num_items), key=lambda i: per_item_total_tokens[i]
536
- )
537
-
538
- # Each batch is a tuple of (indices, global_budget, local_budget).
539
- batches: list[tuple[list[int], int | None, int | None]] = []
540
- current_batch: list[int] = []
541
- current_global_tokens = 0
542
- current_local_tokens = 0
543
-
544
- for orig_idx in sorted_indices:
545
- global_token = per_item_global_tokens[orig_idx]
546
- local_token = per_item_local_tokens[orig_idx]
547
- if (
548
- current_global_tokens + global_token <= max_global_budget
549
- and current_local_tokens + local_token <= max_local_budget
550
- and len(current_batch) < self.max_batch_size
551
- ):
552
- current_batch.append(orig_idx)
553
- current_global_tokens += global_token
554
- current_local_tokens += local_token
555
- else:
556
- if current_batch:
557
- batches.append(
558
- (
559
- current_batch,
560
- self._find_smallest_fitting_budget_given_tokens(
561
- current_global_tokens, self.global_token_budgets
562
- ),
563
- self._find_smallest_fitting_budget_given_tokens(
564
- current_local_tokens, self.local_token_budgets
565
- ),
566
- )
567
- )
568
- current_batch = [orig_idx]
569
- current_global_tokens = global_token
570
- current_local_tokens = local_token
571
-
572
- if current_batch:
573
- batches.append(
574
- (
575
- current_batch,
576
- self._find_smallest_fitting_budget_given_tokens(
577
- current_global_tokens, self.global_token_budgets
578
- ),
579
- self._find_smallest_fitting_budget_given_tokens(
580
- current_local_tokens, self.local_token_budgets
581
- ),
582
- )
583
- )
410
+ append_current_batch()
411
412
outputs_by_orig_idx: dict[int, torch.Tensor] = {}
586
-
587
- for batch_orig_indices, global_budget, local_budget in batches:
413
+ for batch_indices, path_budgets in batches:
414
batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
589
- mm_kwargs, batch_orig_indices
415
+ mm_kwargs, batch_indices
416
)
591
- batch_global_tokens = sum(
592
- per_item_global_tokens[i] for i in batch_orig_indices
593
- )
594
- batch_local_tokens = sum(
595
- per_item_local_tokens[i] for i in batch_orig_indices
596
- )
597
-
598
- both_eager = global_budget is None and local_budget is None
599
-
600
- if both_eager:
601
- logger.debug(
602
- "Encoder CUDA graph dual-path full eager fallback: "
603
- "%d global + %d local tokens from %d images",
604
- batch_global_tokens,
605
- batch_local_tokens,
606
- len(batch_orig_indices),
607
- )
608
- self.graph_misses += len(batch_orig_indices)
609
- with torch.inference_mode():
610
- raw = self.model.encoder_eager_forward(batch_mm_kwargs)
611
- per_item_total = [
612
- per_item_global_tokens[i] + per_item_local_tokens[i] + 1
613
- for i in batch_orig_indices
614
- ]
615
- scatter_output_slices(
616
- raw, batch_orig_indices, per_item_total, outputs_by_orig_idx
617
- )
618
- continue
417
+ graph_outputs: dict[str, torch.Tensor] = {}
418
+ all_eager = True
419
620
- logger.debug(
621
- "Encoder CUDA graph dual-path: batch_size=%d, "
622
- "global=%d (budget=%s), local=%d (budget=%s)",
623
- len(batch_orig_indices),
624
- batch_global_tokens,
625
- global_budget,
626
- batch_local_tokens,
627
- local_budget,
628
- )
420
+ for path in paths:
421
+ path_tokens = sum(per_item_path_tokens[path][i] for i in batch_indices)
422
+ if path_tokens == 0:
423
+ continue
424
630
- # Execute global path: graph or eager fallback
631
- if global_budget is not None:
632
- global_output = self._run_budget_graph(
633
- batch_mm_kwargs,
634
- global_budget,
635
- path="global",
636
- )
637
- assert global_output is not None
638
- else:
639
- with torch.inference_mode():
640
- global_output = self.model.encoder_eager_forward(
641
- batch_mm_kwargs, path="global"
425
+ token_budget = path_budgets[path]
426
+ if token_budget is None:
427
+ with torch.inference_mode():
428
+ output = self.model.encoder_eager_forward(
429
+ batch_mm_kwargs, path=path
430
+ )
431
+ else:
432
+ all_eager = False
433
+ output = self._run_budget_graph(
434
+ batch_mm_kwargs, token_budget, path=path
435
)
436
+ assert output is not None
437
+ graph_outputs[path] = output
438
644
- # Execute local path: graph or eager fallback
645
- if local_budget is not None and batch_local_tokens > 0:
646
- local_output = self._run_budget_graph(
647
- batch_mm_kwargs,
648
- local_budget,
649
- path="local",
650
- )
651
- assert local_output is not None
652
- elif batch_local_tokens > 0:
653
- with torch.inference_mode():
654
- local_output = self.model.encoder_eager_forward(
655
- batch_mm_kwargs, path="local"
656
- )
657
- else:
658
- local_output = None
439
+ if all_eager:
440
+ self.graph_misses += len(batch_indices)
441
442
self.model.postprocess_encoder_output(
661
- global_output,
662
- batch_orig_indices,
663
- per_item_global_tokens,
443
+ graph_outputs,
444
+ batch_indices,
445
+ per_item_out_tokens,
446
outputs_by_orig_idx,
447
clone=True,
448
batch_mm_kwargs=batch_mm_kwargs,
667
- local_output=local_output,
449
)
450
451
return [outputs_by_orig_idx[i] for i in range(num_items)]
vllm/v1/worker/encoder_cudagraph_defs.py
+25
-17
@@ -26,13 +26,28 @@ class EncoderItemSpec:
26
"""Number of output tokens after encoder processing (e.g. after
27
spatial merge)."""
28
29
- global_output_tokens: int = 0
30
- """Number of output tokens from the global image path.
31
- Only used when ``EncoderCudaGraphConfig.enable_dual_path_graph`` is True."""
29
+ path_output_tokens: dict[str, int] = field(default_factory=dict)
30
+ """Per-path output token counts for multi-path encoders.
31
33
- local_output_tokens: int = 0
34
- """Number of output tokens from the local patch path.
35
- Only used when ``EncoderCudaGraphConfig.enable_dual_path_graph`` is True."""
32
+ Single-path encoders leave this empty and use ``output_tokens`` for the
33
+ default path.
34
+ """
35
+
36
+ def get_path_output_tokens(self, path: str) -> int:
37
+ if path == "default" and not self.path_output_tokens:
38
+ return self.output_tokens
39
+ return self.path_output_tokens.get(path, 0)
40
+
41
+
42
+@dataclass(frozen=True)
43
+class EncoderCudaGraphPathConfig:
44
+ """Capture policy for one independently replayable encoder path."""
45
+
46
+ min_token_budget: int | None = None
47
+ """Smallest capture budget, or the model default minimum when unset."""
48
+
49
+ allow_zero_tokens: bool = False
50
+ """Whether a batch may omit this path entirely."""
51
52
53
@dataclass
@@ -68,17 +83,10 @@ class EncoderCudaGraphConfig:
83
Only relevant when "video" is in ``modalities``.
84
Image-only models can use the default of 1."""
85
71
- enable_dual_path_graph: bool = False
72
- """If True, the manager captures two independent graph sets
73
- (global + local) and runs dual-path graph selection during inference."""
74
-
75
- global_token_per_image: int = 0
76
- """Tokens per global image (e.g. 272 for DeepSeek-OCR).
77
- Only used when ``enable_dual_path_graph`` is True."""
78
-
79
- local_token_per_patch: int = 0
80
- """Tokens per local patch (e.g. 100 for DeepSeek-OCR).
81
- Only used when ``enable_dual_path_graph`` is True."""
86
+ paths: dict[str, EncoderCudaGraphPathConfig] = field(
87
+ default_factory=lambda: {"default": EncoderCudaGraphPathConfig()}
88
+ )
89
+ """Independently captured encoder paths keyed by their forward name."""
90
91
92
@dataclass