diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 04fee73af6..e7a847f4c7 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -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 diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 45c7a503a8..f2a54c55c8 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -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 diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index b811afafb0..58cedbb0f9 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -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] diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 668c958cef..efdab8c4d1 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -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, diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index 7b3bb93ad1..943549c46e 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -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 ) diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index d3c5481289..0b3543da1e 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -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)] diff --git a/vllm/v1/worker/encoder_cudagraph_defs.py b/vllm/v1/worker/encoder_cudagraph_defs.py index ae79024302..f86f43db8e 100644 --- a/vllm/v1/worker/encoder_cudagraph_defs.py +++ b/vllm/v1/worker/encoder_cudagraph_defs.py @@ -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