iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[1/N] Unify multiple-path encoder cuda graph support (#49934)

Browse Source 
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Isotr0py
2026-08-02 10:33:23 +08:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent c67fe497a2
commit e2fa28594f
7 changed files with 193 additions and 388 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/cuda_graphs_multimodal.md
+16 -25
View File
@@ -41,13 +41,13 @@ class BudgetGraphMetadata:
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`.
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.
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.
### Greedy bin-packing at runtime
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.
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)).
For multi-path models, the same greedy packing loop constrains every configured path simultaneously (see [Multi-Path graph capture](#multi-path-graph-capture)).
For each graph replay:
@@ -56,38 +56,29 @@ For each graph replay:
3. Replay the CUDA Graph.
4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays).
### Dual-Path graph capture
### Multi-Path graph capture
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.
`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"]`.
**Budget generation.** Two separate budget lists are generated:
**Budget generation.** Each path gets a separate budget list. For DeepSeek-OCR:
* `global_token_budgets` — power-of-2 multiples of `global_token_per_image` (e.g., `[272, 544, 1088, 2176, 4352, 8704, 13824]` for DeepSeek-OCR).
* `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).
* 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).
* 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).
Both lists are capped at the same `max_budget`.
**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:
**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:
* Sort images by total output tokens (global + local), smallest first.
* 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`.
* Once either constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path.
* 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`.
* Once any path constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path.
* Repeat until all images are packed.
**Partial graph fallback.** After packing, each sub-batch falls into one of four execution scenarios:
| Global budget | Local budget | Execution |
| :---: | :---: | --- |
| Found | Found | Both paths use CUDA graph replay |
| Found | `None` | Global graph replay + local path skipped (no patches) |
| `None` | Found | Global eager fallback + local graph replay |
| `None` | `None` | Both paths fall back to eager execution |
Note that the `0`-budget graph is never actually replayed for local — it signals that local patch processing should be skipped entirely.
**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.
**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.
**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.
**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.
!!! note
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
* `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video).
* `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets.
* `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.
* `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.
* `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding.
* `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.
* `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()`.
* `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).
* `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.
* `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.
* `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.
!!! note
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
## Configuration
Three fields in `CompilationConfig` control encoder CUDA Graphs:
Four fields in `CompilationConfig` control encoder CUDA Graphs:
* `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.
* `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.
* `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`.
* `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).
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.
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.
## Usage guide
tests/v1/cudagraph/test_encoder_cudagraph.py
+2
View File
@@ -106,6 +106,7 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
"""
mgr = object.__new__(EncoderCudaGraphManager)
mgr.token_budgets = sorted(budgets)
mgr.path_token_budgets = {"default": mgr.token_budgets}
mgr.max_batch_size = 16
mgr.use_dp = False
mgr.config = EncoderCudaGraphConfig(
@@ -417,6 +418,7 @@ def _make_manager_for_gpu(
"""Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests."""
mgr = object.__new__(EncoderCudaGraphManager)
mgr.token_budgets = sorted(token_budgets)
mgr.path_token_budgets = {"default": mgr.token_budgets}
mgr.max_batch_size = max_batch_size
mgr.max_frames_per_batch = (
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
View File
@@ -65,6 +65,7 @@ from vllm.v1.sample.logits_processor import (
from vllm.v1.worker.encoder_cudagraph_defs import (
EncoderCudaGraphCaptureInputs,
EncoderCudaGraphConfig,
EncoderCudaGraphPathConfig,
EncoderCudaGraphReplayBuffers,
EncoderItemSpec,
)
@@ -697,9 +698,15 @@ class DeepseekOCRForCausalLM(
modalities=["image"],
buffer_keys=["pixel_values"],
out_hidden_size=self.projector_config.n_embed,
enable_dual_path_graph=True,
global_token_per_image=self.global_image_output_token,
local_token_per_patch=self.single_patch_output_token,
paths={
"global": EncoderCudaGraphPathConfig(
min_token_budget=self.global_image_output_token
),
"local": EncoderCudaGraphPathConfig(
min_token_budget=self.single_patch_output_token,
allow_zero_tokens=True,
),
},
)
def get_encoder_cudagraph_budget_range(
@@ -730,8 +737,10 @@ class DeepseekOCRForCausalLM(
EncoderItemSpec(
input_size=num_input_tokens,
output_tokens=num_output_tokens,
global_output_tokens=global_output_token,
local_output_tokens=local_output_token,
path_output_tokens={
"global": global_output_token,
"local": local_output_token,
},
)
)
return item_specs
@@ -917,32 +926,34 @@ class DeepseekOCRForCausalLM(
def postprocess_encoder_output(
self,
output: torch.Tensor,
outputs: dict[str, torch.Tensor],
indices: list[int],
per_item_out_tokens: list[int],
dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
clone: bool = False,
batch_mm_kwargs: dict[str, Any] | None = None,
local_output: torch.Tensor | None = None,
) -> None:
"""
Assemble per-image embeddings from global and local encoder outputs.
``output`` contains global-image features with newlines already
``output['global']`` contains global-image features with newlines already
inserted (from CUDA graph replay or eager fallback):
``[B * 272, n_embed]``.
``local_output`` contains local-patch features without
``output['local']`` contains local-patch features without
newlines (from CUDA graph replay or eager fallback):
``[P * 100, n_embed]``. May be ``None`` if no patches in batch.
This method:
1. Splits ``output`` into per-image global portions.
2. Splits ``local_output`` into per-image patch groups.
1. Splits ``output['global']`` into per-image global portions.
2. Splits ``output['local']`` into per-image patch groups.
3. For each image: assembles patch grid with newlines via
``_assemble_patch_grid``, then concatenates
``[local_tiled, global, view_seperator]``.
"""
output = outputs["global"]
local_output = outputs.get("local")
assert batch_mm_kwargs is not None
bsz = len(indices)
n_embed = output.shape[-1]
vllm/model_executor/models/interfaces.py
+4 -3
View File
@@ -1696,13 +1696,12 @@ class SupportsEncoderCudaGraph(Protocol):
def postprocess_encoder_output(
self,
output: torch.Tensor,
outputs: dict[str, torch.Tensor],
indices: list[int],
per_item_out_tokens: list[int],
dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
clone: bool = False,
batch_mm_kwargs: dict[str, Any] | None = None,
local_output: torch.Tensor | None = None,
) -> None:
"""
Post-process encoder output, directly call scatter_output_slices by default.
@@ -1714,7 +1713,9 @@ class SupportsEncoderCudaGraph(Protocol):
"""
from vllm.model_executor.models.utils import scatter_output_slices
scatter_output_slices(output, indices, per_item_out_tokens, dest, clone)
scatter_output_slices(
outputs["default"], indices, per_item_out_tokens, dest, clone
)
def prepare_encoder_cudagraph_capture_inputs(
self,
vllm/model_executor/models/step3_vl.py
+20 -9
View File
@@ -725,6 +725,7 @@ class Step3VLForConditionalGeneration(
def get_encoder_cudagraph_config(self):
from vllm.v1.worker.encoder_cudagraph_defs import (
EncoderCudaGraphConfig,
EncoderCudaGraphPathConfig,
)
return EncoderCudaGraphConfig(
@@ -734,9 +735,15 @@ class Step3VLForConditionalGeneration(
"patch_pixel_values",
],
out_hidden_size=self.config.hidden_size,
enable_dual_path_graph=True,
global_token_per_image=self.img_output_tokens,
local_token_per_patch=self.patch_output_tokens,
paths={
"global": EncoderCudaGraphPathConfig(
min_token_budget=self.img_output_tokens
),
"local": EncoderCudaGraphPathConfig(
min_token_budget=self.patch_output_tokens,
allow_zero_tokens=True,
),
},
)
def get_encoder_cudagraph_budget_range(
@@ -771,8 +778,10 @@ class Step3VLForConditionalGeneration(
output_tokens=(
self.img_output_tokens + num_patch * self.patch_output_tokens
),
global_output_tokens=self.img_output_tokens,
local_output_tokens=num_patch * self.patch_output_tokens,
path_output_tokens={
"global": self.img_output_tokens,
"local": num_patch * self.patch_output_tokens,
},
)
for num_patch in num_patches
]
@@ -876,19 +885,21 @@ class Step3VLForConditionalGeneration(
def postprocess_encoder_output(
self,
output: torch.Tensor,
outputs: dict[str, torch.Tensor],
indices: list[int],
per_item_out_tokens: list[int],
dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
clone: bool = False,
batch_mm_kwargs: dict[str, Any] | None = None,
local_output: torch.Tensor | None = None,
):
"""CPU-side per-item merge after dual-path graph replay.
``output`` contains global-image features and ``local_output``
``outputs['global']`` contains global-image features and ``outputs['local']``
contains local-patch features (or ``None`` when there are no patches).
"""
output = outputs["global"]
local_output = outputs.get("local")
assert batch_mm_kwargs is not None
num_patches = batch_mm_kwargs["num_patches"]
hidden = output.shape[-1]
bsz = len(indices)
@@ -899,7 +910,7 @@ class Step3VLForConditionalGeneration(
patch_tokens = total_patches * self.patch_output_tokens
global_part = output[:img_tokens].reshape(bsz, self.img_output_tokens, hidden)
if total_patches > 0:
if total_patches > 0 and local_output is not None:
patch_part = local_output[:patch_tokens].reshape(
-1, self.patch_output_tokens, hidden
)
vllm/v1/worker/encoder_cudagraph.py
+104 -323
View File
@@ -158,34 +158,33 @@ class EncoderCudaGraphManager:
self.graph_misses = 0
self.log_stats_interval = 100
if self.config.enable_dual_path_graph:
max_budget = self.token_budgets[-1]
self.global_token_budgets = self._generate_budgets(
self.config.global_token_per_image,
max_budget,
)
self.local_token_budgets = self._generate_budgets(
self.config.local_token_per_patch,
max_budget,
)
# When `image_width <= 640 and image_height <= 640`, the mm inputs
# will only contain global image, without generating local patches.
self.local_token_budgets.insert(0, 0)
logger.info(
"EncoderCudaGraphManager dual-path mode: "
"global_budgets=%s, local_budgets=%s",
self.global_token_budgets,
self.local_token_budgets,
)
else:
logger.info(
"EncoderCudaGraphManager initialized with "
"budgets=%s, max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",
self.token_budgets,
self.max_batch_size,
self.max_frames_per_batch,
self.use_dp,
)
max_budget = self.token_budgets[-1]
if not self.config.paths:
raise ValueError("Encoder CUDA graph config must define at least one path")
self.path_token_budgets: dict[str, list[int]] = {}
for path, path_config in self.config.paths.items():
min_path_budget = path_config.min_token_budget
if min_path_budget is None:
budgets = list(self.token_budgets)
else:
if min_path_budget <= 0 or min_path_budget > max_budget:
raise ValueError(
f"Invalid minimum budget {min_path_budget} for encoder "
f"CUDA graph path {path!r}; max budget is {max_budget}"
)
budgets = self._generate_budgets(min_path_budget, max_budget)
if path_config.allow_zero_tokens:
budgets.insert(0, 0)
self.path_token_budgets[path] = budgets
logger.info(
"EncoderCudaGraphManager initialized with paths=%s, "
"max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",
self.path_token_budgets,
self.max_batch_size,
self.max_frames_per_batch,
self.use_dp,
)
@staticmethod
def _generate_budgets(min_budget: int, max_budget: int) -> list[int]:
@@ -211,40 +210,29 @@ class EncoderCudaGraphManager:
self.graph_pool = None
def capture(self, graph_pool: Any):
"""Capture CUDA graphs for all token budgets."""
"""Capture CUDA graphs for every configured path and token budget."""
self.graph_pool = graph_pool
if self.config.enable_dual_path_graph:
for token_budget in sorted(self.global_token_budgets, reverse=True):
self._capture_budget_graph(token_budget, path="global")
for token_budget in sorted(self.local_token_budgets, reverse=True):
if token_budget == 0:
continue
self._capture_budget_graph(token_budget, path="local")
logger.info(
"Encoder CUDA graph capture complete. "
"Captured %d global + %d local budget graphs.",
len(self.budget_graphs["global"]),
len(self.budget_graphs["local"]),
)
return
for token_budget in sorted(self.token_budgets, reverse=True):
self._capture_budget_graph(token_budget)
for path, budgets in self.path_token_budgets.items():
for token_budget in sorted(budgets, reverse=True):
if token_budget > 0:
self._capture_budget_graph(token_budget, path=path)
logger.info(
"Encoder CUDA graph capture complete. Captured %d budget graphs.",
len(self.budget_graphs["default"]),
"Encoder CUDA graph capture complete. Captured %d graphs across %d paths.",
self.get_num_graphs_to_capture(),
len(self.path_token_budgets),
)
def get_num_graphs_to_capture(self) -> int:
if self.config.enable_dual_path_graph:
return len(self.global_token_budgets) + len(self.local_token_budgets)
return len(self.token_budgets)
return sum(
token_budget > 0
for budgets in self.path_token_budgets.values()
for token_budget in budgets
)
def _get_graph_set(self, path: str = "default") -> dict[int, BudgetGraphMetadata]:
# Lazy init global/local graph sets for dual-path models, or default graph
# set for single-path models.
"""Return the captured graphs for one encoder path."""
if path not in self.budget_graphs:
self.budget_graphs[path] = {}
return self.budget_graphs[path]
@@ -331,7 +319,7 @@ class EncoderCudaGraphManager:
Args:
mm_kwargs: Multimodal inputs for the batch.
token_budget: Token budget to use.
path: Path for the graph. Should be one of ["default", "global", "local"].
path: Configured encoder path.
Returns:
Encoder outputs, or None if graph not captured.
"""
@@ -351,9 +339,7 @@ class EncoderCudaGraphManager:
path,
)
# Copy replay buffers into graph input buffers. Iterate over the
# graph's own buffer keys (which may differ per path for dual-path
# models) rather than the global config.buffer_keys.
# Copy replay values into the buffers recorded for this path.
for key, buf in graph_meta.input_buffers.items():
src = replay.values.get(key)
if src is None:
@@ -375,296 +361,91 @@ class EncoderCudaGraphManager:
self,
mm_kwargs: dict[str, Any],
) -> list[torch.Tensor]:
"""Execute encoder on local inputs using greedy-packed CUDA graphs.
Sort images by output token count (smallest first), then greedily pack
as many images as possible into each batch while staying within
max_budget tokens and max_batch_size. Once a batch is finalised (next
image would overflow either constraint), find the smallest fitting
budget once for that batch.
For dual-path models (``enable_dual_path_graph=True``), two independent
graph sets are used: one for global images, one for local patches.
Budgets are found independently per path; if only one path fits, the
other falls back to eager via partial fallback.
By exchange argument, greedy smallest-first packing minimises eager
fallbacks -- any other ordering yields a higher token sum in some batch,
making that batch more likely to exceed the budget.
Stats note:
graph_hits -- counted inside _run_budget_graph after successful replay.
graph_misses -- counted here for single-image batches where the image
exceeds max_budget. Batches split due to max_batch_size
always satisfy total_tokens <= max_budget and therefore
always find a valid budget (no miss).
"""
if self.config.enable_dual_path_graph:
return self._execute_local_dual_path(mm_kwargs)
return self._execute_local_single_path(mm_kwargs)
def _execute_local_single_path(
self,
mm_kwargs: dict[str, Any],
) -> list[torch.Tensor]:
"""Single-path greedy-packing execution (original behaviour)."""
"""Execute locally using greedy packing across all configured paths."""
item_specs = self._get_item_specs(mm_kwargs)
num_items = len(item_specs)
max_budget = self.token_budgets[-1]
paths = tuple(self.path_token_budgets)
per_item_path_tokens = {
path: [spec.get_path_output_tokens(path) for spec in item_specs]
for path in paths
}
per_item_out_tokens = [spec.output_tokens for spec in item_specs]
max_path_budgets = {path: max(self.path_token_budgets[path]) for path in paths}
# Sort ascending by output token count (smallest first)
sorted_indices = sorted(range(num_items), key=lambda i: per_item_out_tokens[i])
# Greedy pack against max_budget and max_batch_size.
# _find_smallest_fitting_budget_given_tokens is called once per
# finalised batch, not per image.
batches: list[tuple[list[int], int | None]] = []
batches: list[tuple[list[int], dict[str, int | None]]] = []
current_batch: list[int] = []
current_batch_tokens = 0
current_tokens = dict.fromkeys(paths, 0)
def append_current_batch() -> None:
if not current_batch:
return
path_budgets = {
path: (
0
if current_tokens[path] == 0
else self._find_smallest_fitting_budget_given_tokens(
current_tokens[path], self.path_token_budgets[path]
)
)
for path in paths
}
batches.append((list(current_batch), path_budgets))
for orig_idx in sorted_indices:
item_tokens = per_item_out_tokens[orig_idx]
if (
current_batch_tokens + item_tokens <= max_budget
and len(current_batch) < self.max_batch_size
):
current_batch.append(orig_idx)
current_batch_tokens += item_tokens
else:
if current_batch:
batches.append(
(
current_batch,
self._find_smallest_fitting_budget_given_tokens(
current_batch_tokens
),
)
)
current_batch = [orig_idx]
current_batch_tokens = item_tokens
if current_batch:
batches.append(
(
current_batch,
self._find_smallest_fitting_budget_given_tokens(
current_batch_tokens
),
)
item_tokens = {path: per_item_path_tokens[path][orig_idx] for path in paths}
fits = len(current_batch) < self.max_batch_size and all(
current_tokens[path] + item_tokens[path] <= max_path_budgets[path]
for path in paths
)
if current_batch and not fits:
append_current_batch()
current_batch = []
current_tokens = dict.fromkeys(paths, 0)
# outputs_by_orig_idx maps each original image index to its output
# tensor. Needed because greedy packing reorders images; we restore
# the original order before returning.
outputs_by_orig_idx: dict[int, torch.Tensor] = {}
current_batch.append(orig_idx)
for path in paths:
current_tokens[path] += item_tokens[path]
for batch_orig_indices, token_budget in batches:
batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
mm_kwargs, batch_orig_indices
)
batch_out_tokens = sum(per_item_out_tokens[i] for i in batch_orig_indices)
if token_budget is None:
# Single oversized image: item_tokens > max_budget.
# graph_misses counted here for this eager fallback.
logger.debug(
"Encoder CUDA graph fallback to eager: no budget for "
"%d tokens from %d images",
batch_out_tokens,
len(batch_orig_indices),
)
self.graph_misses += len(batch_orig_indices)
with torch.inference_mode():
raw = self.model.encoder_eager_forward(batch_mm_kwargs)
scatter_output_slices(
raw,
batch_orig_indices,
per_item_out_tokens,
outputs_by_orig_idx,
)
else:
logger.debug(
"Encoder CUDA graph: batch_size=%d, tokens=%d, "
"budget=%d, waste=%.1f%%",
len(batch_orig_indices),
batch_out_tokens,
token_budget,
(token_budget - batch_out_tokens) / token_budget * 100,
)
# graph_hits counted inside _run_budget_graph after replay.
output = self._run_budget_graph(batch_mm_kwargs, token_budget)
assert output is not None
self.model.postprocess_encoder_output(
output,
batch_orig_indices,
per_item_out_tokens,
outputs_by_orig_idx,
clone=True,
batch_mm_kwargs=batch_mm_kwargs,
)
# Return in original batch order (caller maps outputs to token positions)
return [outputs_by_orig_idx[i] for i in range(num_items)]
def _execute_local_dual_path(
self,
mm_kwargs: dict[str, Any],
) -> list[torch.Tensor]:
"""Dual-path greedy-packing execution.
Each image contributes both global tokens (constant per image)
and local tokens (patches * patch_tokens). Greedy packing
respects both budgets independently, then selects the smallest
fitting budget per path with partial eager fallback.
"""
item_specs = self._get_item_specs(mm_kwargs)
num_items = len(item_specs)
max_global_budget = self.global_token_budgets[-1]
max_local_budget = self.local_token_budgets[-1]
per_item_global_tokens = [spec.global_output_tokens for spec in item_specs]
per_item_local_tokens = [spec.local_output_tokens for spec in item_specs]
per_item_total_tokens = [spec.output_tokens for spec in item_specs]
# Sort ascending by total output tokens
sorted_indices = sorted(
range(num_items), key=lambda i: per_item_total_tokens[i]
)
# Each batch is a tuple of (indices, global_budget, local_budget).
batches: list[tuple[list[int], int | None, int | None]] = []
current_batch: list[int] = []
current_global_tokens = 0
current_local_tokens = 0
for orig_idx in sorted_indices:
global_token = per_item_global_tokens[orig_idx]
local_token = per_item_local_tokens[orig_idx]
if (
current_global_tokens + global_token <= max_global_budget
and current_local_tokens + local_token <= max_local_budget
and len(current_batch) < self.max_batch_size
):
current_batch.append(orig_idx)
current_global_tokens += global_token
current_local_tokens += local_token
else:
if current_batch:
batches.append(
(
current_batch,
self._find_smallest_fitting_budget_given_tokens(
current_global_tokens, self.global_token_budgets
),
self._find_smallest_fitting_budget_given_tokens(
current_local_tokens, self.local_token_budgets
),
)
)
current_batch = [orig_idx]
current_global_tokens = global_token
current_local_tokens = local_token
if current_batch:
batches.append(
(
current_batch,
self._find_smallest_fitting_budget_given_tokens(
current_global_tokens, self.global_token_budgets
),
self._find_smallest_fitting_budget_given_tokens(
current_local_tokens, self.local_token_budgets
),
)
)
append_current_batch()
outputs_by_orig_idx: dict[int, torch.Tensor] = {}
for batch_orig_indices, global_budget, local_budget in batches:
for batch_indices, path_budgets in batches:
batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
mm_kwargs, batch_orig_indices
)
batch_global_tokens = sum(
per_item_global_tokens[i] for i in batch_orig_indices
)
batch_local_tokens = sum(
per_item_local_tokens[i] for i in batch_orig_indices
mm_kwargs, batch_indices
)
graph_outputs: dict[str, torch.Tensor] = {}
all_eager = True
both_eager = global_budget is None and local_budget is None
for path in paths:
path_tokens = sum(per_item_path_tokens[path][i] for i in batch_indices)
if path_tokens == 0:
continue
if both_eager:
logger.debug(
"Encoder CUDA graph dual-path full eager fallback: "
"%d global + %d local tokens from %d images",
batch_global_tokens,
batch_local_tokens,
len(batch_orig_indices),
)
self.graph_misses += len(batch_orig_indices)
with torch.inference_mode():
raw = self.model.encoder_eager_forward(batch_mm_kwargs)
per_item_total = [
per_item_global_tokens[i] + per_item_local_tokens[i] + 1
for i in batch_orig_indices
]
scatter_output_slices(
raw, batch_orig_indices, per_item_total, outputs_by_orig_idx
)
continue
logger.debug(
"Encoder CUDA graph dual-path: batch_size=%d, "
"global=%d (budget=%s), local=%d (budget=%s)",
len(batch_orig_indices),
batch_global_tokens,
global_budget,
batch_local_tokens,
local_budget,
)
# Execute global path: graph or eager fallback
if global_budget is not None:
global_output = self._run_budget_graph(
batch_mm_kwargs,
global_budget,
path="global",
)
assert global_output is not None
else:
with torch.inference_mode():
global_output = self.model.encoder_eager_forward(
batch_mm_kwargs, path="global"
token_budget = path_budgets[path]
if token_budget is None:
with torch.inference_mode():
output = self.model.encoder_eager_forward(
batch_mm_kwargs, path=path
)
else:
all_eager = False
output = self._run_budget_graph(
batch_mm_kwargs, token_budget, path=path
)
assert output is not None
graph_outputs[path] = output
# Execute local path: graph or eager fallback
if local_budget is not None and batch_local_tokens > 0:
local_output = self._run_budget_graph(
batch_mm_kwargs,
local_budget,
path="local",
)
assert local_output is not None
elif batch_local_tokens > 0:
with torch.inference_mode():
local_output = self.model.encoder_eager_forward(
batch_mm_kwargs, path="local"
)
else:
local_output = None
if all_eager:
self.graph_misses += len(batch_indices)
self.model.postprocess_encoder_output(
global_output,
batch_orig_indices,
per_item_global_tokens,
graph_outputs,
batch_indices,
per_item_out_tokens,
outputs_by_orig_idx,
clone=True,
batch_mm_kwargs=batch_mm_kwargs,
local_output=local_output,
)
return [outputs_by_orig_idx[i] for i in range(num_items)]
vllm/v1/worker/encoder_cudagraph_defs.py
+25 -17
View File
@@ -26,13 +26,28 @@ class EncoderItemSpec:
"""Number of output tokens after encoder processing (e.g. after
spatial merge)."""
global_output_tokens: int = 0
"""Number of output tokens from the global image path.
Only used when ``EncoderCudaGraphConfig.enable_dual_path_graph`` is True."""
path_output_tokens: dict[str, int] = field(default_factory=dict)
"""Per-path output token counts for multi-path encoders.
local_output_tokens: int = 0
"""Number of output tokens from the local patch path.
Only used when ``EncoderCudaGraphConfig.enable_dual_path_graph`` is True."""
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)
@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."""
@dataclass
@@ -68,17 +83,10 @@ class EncoderCudaGraphConfig:
Only relevant when "video" is in ``modalities``.
Image-only models can use the default of 1."""
enable_dual_path_graph: bool = False
"""If True, the manager captures two independent graph sets
(global + local) and runs dual-path graph selection during inference."""
global_token_per_image: int = 0
"""Tokens per global image (e.g. 272 for DeepSeek-OCR).
Only used when ``enable_dual_path_graph`` is True."""
local_token_per_patch: int = 0
"""Tokens per local patch (e.g. 100 for DeepSeek-OCR).
Only used when ``enable_dual_path_graph`` is True."""
paths: dict[str, EncoderCudaGraphPathConfig] = field(
default_factory=lambda: {"default": EncoderCudaGraphPathConfig()}
)
"""Independently captured encoder paths keyed by their forward name."""
@dataclass