@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ad32608e24
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[MM][Perf][CG] Support dual-path ViT full CUDA graph for DeepSeek-OCR (#43586)

Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Shanshan Shen committed Jun 16, 2026 at 19:35 UTC ad32608e24c91b5a21a22eeaa7b94dc3882b3854
16 files changed +809 -69
docs/design/cuda_graphs_multimodal.md
+63 -16
@@ -2,6 +2,8 @@
2
3 The [CUDA Graphs](cuda_graphs.md) infrastructure in vLLM primarily targets the **decoder** (language model) forward pass. vLLM also supports capturing the **encoder** (vision transformer) forward pass as CUDA Graphs, independently from the decoder. This is based on <https://github.com/vllm-project/vllm/pull/35963>.
4
5 +For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tiling), a **dual-path graph** mode captures two independent sets of CUDA graphs — one for the global image path and one for the local patch path — enabling independent budget selection and partial eager fallback per path. This is based on <https://github.com/vllm-project/vllm/pull/43586>.
6 +
7 !!! note
8 Encoder CUDA Graphs are orthogonal to decoder CUDA Graphs — both can be enabled simultaneously. Encoder graphs capture the vision encoder execution (e.g., ViT in Qwen3-VL), while decoder graphs capture the language model execution as described in the [CUDA Graphs design document](cuda_graphs.md).
9
@@ -11,6 +13,8 @@ Vision encoder inference incurs CUDA kernel launch overhead on the host side. Th
13
14 Encoder CUDA Graphs eliminate this overhead by pre-capturing the full encoder forward pass at multiple token budget levels during model initialization, then replaying the appropriate graph at runtime.
15
16 +For two-tower vision encoders such as DeepSeek-OCR (SAM + CLIP with dynamic tiling), the global image path and local patch path have independent token profiles (272 tokens per global image vs. 100 tokens per local patch). Capturing a single monolithic graph for both paths would significantly reduce packing efficiency. The dual-path graph mode captures each path as a separate set of budgets, allowing the manager to pack and replay each path independently.
17 +
18 ## Design
19
20 The encoder CUDA Graph system uses a **budget-based capture/replay** strategy, managed by [EncoderCudaGraphManager][vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager]. The system contains the following core components:
@@ -37,10 +41,14 @@ 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.
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)).
51 +
52 For each graph replay:
53
54 1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs.
@@ -48,6 +56,42 @@ 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
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.
62 +
63 +**Budget generation.** Two separate budget lists are generated:
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).
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:
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.
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.
87 +
88 +**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.
89 +
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.
91 +
92 +!!! note
93 + 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.
94 +
95 ### Data-parallel support
96
97 When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
@@ -67,29 +111,30 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
111
112 * `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video).
113 * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets.
70 -* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size and output token count. Replaces the former three separate methods (`get_num_items`, `get_per_item_output_tokens`, `get_per_item_input_sizes`).
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.
115 * `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding.
72 -* `prepare_encoder_cudagraph_capture_inputs(...)` — creates dummy inputs for graph capture. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph.
73 -* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch)` — computes buffer values from actual batch inputs. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match `buffer_keys` in the config.
74 -* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor])` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `pixel_values` tensor is included in `inputs` alongside metadata buffers.
75 -* `encoder_eager_forward(mm_kwargs)` — fallback eager forward when no graph fits.
76 -* `postprocess_encoder_output(...)` — post-process encoder output, delegates to `scatter_output_slices` by default.
116 +* `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.
117 +* `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()`.
118 +* `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).
119 +* `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.
121
122 !!! note
123 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.
124
125 **Supported models:**
126
83 -| Architecture | Models | CG for Image | CG for Video |
84 -| ------------ | ------ | ------------ | ------------ |
85 -| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | - |
86 -| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
87 -| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
88 -| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
89 -| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
90 -| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
91 -| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |
92 -| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ |
127 +| Architecture | Models | CG for Image | CG for Video | Dual-Path Graph |
128 +| ------------ | ------ | ------------ | ------------ | --------------- |
129 +| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
130 +| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
131 +| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
132 +| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
133 +| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ |
134 +| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ |
135 +| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ |
136 +| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ | ❌︎ |
137 +| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ❌︎ |
138
139 !!! note
140 Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
@@ -104,6 +149,8 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
149 * `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`.
150 * `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).
151
152 +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.
153 +
154 ## Usage guide
155
156 ### Image inference
examples/generate/multimodal/vision_language_offline.py
+3 -2
@@ -2533,15 +2533,16 @@ MODELS_NEED_VIDEO_METADATA = [
2533
2534 MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2535 "llama4",
2536 - "internvl_chat",
2536 + "qwen2_vl",
2537 "qwen2_5_vl",
2538 "qwen3_vl",
2539 "qwen3_vl_moe",
2540 - "qwen2_vl",
2540 "qwen3_5",
2541 "qwen3_5_moe",
2542 + "internvl_chat",
2543 "stepvl",
2544 "glm4_1v",
2545 + "deepseek_ocr",
2546 ]
2547
2548
tests/models/multimodal/generation/test_vit_cudagraph.py
+41 -15
@@ -29,6 +29,7 @@ class VitCudagraphTestConfig:
29 vllm_runner_kwargs: dict = field(default_factory=dict)
30 compilation_config_overrides: dict = field(default_factory=dict)
31 marks: list = field(default_factory=list)
32 + skip: bool = False
33
34
35 def params_with_marks(
@@ -75,15 +76,16 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
76 },
77 marks=[pytest.mark.core_model],
78 ),
78 - "internvl": VitCudagraphTestConfig(
79 - model="OpenGVLab/InternVL3-1B",
80 - num_video_frames=8,
81 - image_prompt=internvl_chat_template("<image>\nWhat is in this image?"),
82 - video_prompt=internvl_chat_template(
83 - "<video>\nDescribe this video in one sentence."
79 + "qwen2_vl": VitCudagraphTestConfig(
80 + model="Qwen/Qwen2-VL-2B-Instruct",
81 + image_prompt=qwen_vl_chat_template(
82 + "<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
83 + ),
84 + video_prompt=qwen_vl_chat_template(
85 + "<|vision_start|><|video_pad|><|vision_end|>"
86 + "Describe this video in one sentence."
87 ),
88 needs_video_metadata=False,
86 - vllm_runner_kwargs={"trust_remote_code": True},
89 marks=[pytest.mark.core_model],
90 ),
91 "qwen2_5_vl": VitCudagraphTestConfig(
@@ -122,16 +124,15 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
124 needs_video_metadata=True,
125 marks=[pytest.mark.core_model],
126 ),
125 - "qwen2_vl": VitCudagraphTestConfig(
126 - model="Qwen/Qwen2-VL-2B-Instruct",
127 - image_prompt=qwen_vl_chat_template(
128 - "<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
129 - ),
130 - video_prompt=qwen_vl_chat_template(
131 - "<|vision_start|><|video_pad|><|vision_end|>"
132 - "Describe this video in one sentence."
127 + "internvl": VitCudagraphTestConfig(
128 + model="OpenGVLab/InternVL3-1B",
129 + num_video_frames=8,
130 + image_prompt=internvl_chat_template("<image>\nWhat is in this image?"),
131 + video_prompt=internvl_chat_template(
132 + "<video>\nDescribe this video in one sentence."
133 ),
134 needs_video_metadata=False,
135 + vllm_runner_kwargs={"trust_remote_code": True},
136 marks=[pytest.mark.core_model],
137 ),
138 "step3_vl": VitCudagraphTestConfig(
@@ -179,6 +180,25 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
180 ),
181 },
182 ),
183 + "deepseek_ocr": VitCudagraphTestConfig(
184 + model="deepseek-ai/DeepSeek-OCR",
185 + modalities=["image"],
186 + image_prompt="<image>\nWhat is in this image?",
187 + marks=[pytest.mark.core_model],
188 + compilation_config_overrides={
189 + "encoder_cudagraph_token_budgets": [272],
190 + "mode": 0,
191 + "cudagraph_mode": 2,
192 + },
193 + vllm_runner_kwargs={
194 + "load_format": "dummy",
195 + "hf_overrides": partial(
196 + dummy_hf_overrides,
197 + model_arch="DeepseekOCRForCausalLM",
198 + ),
199 + },
200 + skip=True, # TODO: Re-enable this once OOM issues are resolved on CI.
201 + ),
202 }
203
204
@@ -201,6 +221,9 @@ def get_compilation_config(config: VitCudagraphTestConfig):
221 def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):
222 config = MODEL_CONFIGS[model_id]
223
224 + if config.skip:
225 + pytest.skip(f"{model_id} is marked to be skipped.")
226 +
227 if "image" not in config.modalities:
228 pytest.skip(f"{model_id} does not support the image modality.")
229
@@ -242,6 +265,9 @@ def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):
265 def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):
266 config = MODEL_CONFIGS[model_id]
267
268 + if config.skip:
269 + pytest.skip(f"{model_id} is marked to be skipped.")
270 +
271 if "video" not in config.modalities:
272 pytest.skip(f"{model_id} does not support the video modality")
273
tests/v1/cudagraph/test_encoder_cudagraph.py
+15 -7
@@ -108,7 +108,7 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
108 mgr.token_budgets = sorted(budgets)
109 mgr.max_batch_size = 16
110 mgr.use_dp = False
111 - mgr.budget_graphs = {}
111 + mgr.budget_graphs = {"default": {}}
112 mgr.graph_pool = None
113 mgr.graph_hits = 0
114 mgr.graph_misses = 0
@@ -341,6 +341,7 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
341 max_frames_per_batch: int,
342 device: torch.device,
343 dtype: torch.dtype,
344 + path: str = "default",
345 ) -> EncoderCudaGraphCaptureInputs:
346 per_image_output = token_budget // max_batch_size
347 grid_config = [
@@ -365,6 +366,7 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
366 mm_kwargs: dict[str, Any],
367 max_batch_size: int,
368 max_frames_per_batch: int,
369 + path: str = "default",
370 ) -> EncoderCudaGraphReplayBuffers:
371 grid_thw = mm_kwargs["image_grid_thw"]
372 n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)
@@ -380,12 +382,14 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
382 def encoder_cudagraph_forward(
383 self,
384 values: dict[str, torch.Tensor],
385 + path: str = "default",
386 ) -> torch.Tensor:
387 return self._forward(values["pixel_values"])
388
389 def encoder_eager_forward(
390 self,
391 mm_kwargs: dict[str, Any],
392 + path: str = "default",
393 ) -> torch.Tensor:
394 return self._forward(mm_kwargs["pixel_values"])
395
@@ -413,7 +417,7 @@ def _make_manager_for_gpu(
417 max_frames_per_batch if max_frames_per_batch is not None else max_batch_size * 2
418 )
419 mgr.use_dp = False
416 - mgr.budget_graphs = {}
420 + mgr.budget_graphs = {"default": {}}
421 mgr.graph_pool = None
422 mgr.graph_hits = 0
423 mgr.graph_misses = 0
@@ -479,15 +483,15 @@ class TestEncoderCudaGraphCaptureReplay:
483 # --- capture ---
484
485 def test_capture_creates_one_graph_per_budget(self):
482 - assert len(self.mgr.budget_graphs) == len(_BUDGETS)
483 - assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)
486 + assert len(self.mgr.budget_graphs["default"]) == len(_BUDGETS)
487 + assert set(self.mgr.budget_graphs["default"].keys()) == set(_BUDGETS)
488
489 def test_capture_uses_supplied_graph_pool(self):
490 assert self.mgr.graph_pool is self.graph_pool
491
492 def test_clear_releases_graphs_and_pool(self):
493 self.mgr.clear()
490 - assert self.mgr.budget_graphs == {}
494 + assert self.mgr.budget_graphs == {"default": {}}
495 assert self.mgr.graph_pool is None
496
497 # --- output shape ---
@@ -642,6 +646,7 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
646 max_frames_per_batch: int,
647 device: torch.device,
648 dtype: torch.dtype,
649 + path: str = "default",
650 ) -> EncoderCudaGraphCaptureInputs:
651 per_item_output = token_budget // max_batch_size
652 frames_per_item = max_frames_per_batch // max_batch_size
@@ -678,6 +683,7 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
683 mm_kwargs: dict[str, Any],
684 max_batch_size: int,
685 max_frames_per_batch: int,
686 + path: str = "default",
687 ) -> EncoderCudaGraphReplayBuffers:
688 n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE)
689 p = next(self.parameters())
@@ -692,12 +698,14 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
698 def encoder_cudagraph_forward(
699 self,
700 values: dict[str, torch.Tensor],
701 + path: str = "default",
702 ) -> torch.Tensor:
703 return self._forward(values["pixel_values"])
704
705 def encoder_eager_forward(
706 self,
707 mm_kwargs: dict[str, Any],
708 + path: str = "default",
709 ) -> torch.Tensor:
710 return self._forward(self._get_pixel_values(mm_kwargs))
711
@@ -763,8 +771,8 @@ class TestEncoderCudaGraphVideoReplay:
771 # --- capture ---
772
773 def test_capture_creates_one_graph_per_budget(self):
766 - assert len(self.mgr.budget_graphs) == len(_BUDGETS)
767 - assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)
774 + assert len(self.mgr.budget_graphs["default"]) == len(_BUDGETS)
775 + assert set(self.mgr.budget_graphs["default"].keys()) == set(_BUDGETS)
776
777 # --- output shape ---
778
vllm/model_executor/models/deepseek_ocr.py
+375 -5
@@ -4,7 +4,7 @@
4
5 import math
6 from collections.abc import Iterable, Mapping, Sequence
7 -from typing import Annotated, Literal
7 +from typing import Annotated, Any, Literal
8
9 import torch
10 import torch.nn as nn
@@ -15,6 +15,7 @@ from vllm.config.multimodal import BaseDummyOptions
15 from vllm.inputs import MultiModalDataDict
16 from vllm.model_executor.models.interfaces import (
17 MultiModalEmbeddings,
18 + SupportsEncoderCudaGraph,
19 SupportsLoRA,
20 SupportsMultiModal,
21 SupportsPP,
@@ -52,6 +53,7 @@ from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
53 from vllm.transformers_utils.processors.deepseek_ocr import (
54 BASE_SIZE,
55 CROP_MODE,
56 + IMAGE_SIZE,
57 DeepseekOCRProcessor,
58 count_tiles,
59 )
@@ -60,12 +62,17 @@ from vllm.v1.sample.logits_processor import (
62 AdapterLogitsProcessor,
63 RequestLogitsProcessor,
64 )
65 +from vllm.v1.worker.encoder_cudagraph_defs import (
66 + EncoderCudaGraphCaptureInputs,
67 + EncoderCudaGraphConfig,
68 + EncoderCudaGraphReplayBuffers,
69 + EncoderItemSpec,
70 +)
71
72 from .deepencoder import DeepCLIPVisionTransformer, build_sam_vit_b
73 from .deepseek_vl2 import MlpProjector
74
75 # The image token id may be various
68 -IMAGE_SIZE = 640
76 _IMAGE_TOKEN = "<image>"
77
78
@@ -355,7 +362,9 @@ class DeepseekOCRMultiModalProcessor(
362 info=DeepseekOCRProcessingInfo,
363 dummy_inputs=DeepseekOCRDummyInputsBuilder,
364 )
358 -class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
365 +class DeepseekOCRForCausalLM(
366 + nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsEncoderCudaGraph
367 +):
368 hf_to_vllm_mapper = WeightsMapper(
369 orig_to_new_prefix={
370 # map prefix for language backbone
@@ -383,6 +392,7 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
392 multimodal_config = vllm_config.model_config.multimodal_config
393
394 self.config = config
395 + self.model_config = vllm_config.model_config
396 self.multimodal_config = multimodal_config
397
398 self.vision_config = config.vision_config
@@ -504,9 +514,14 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
514 )
515 features = self.projector(features)
516
517 + return self._assemble_patch_grid(features, crop_shape)
518 +
519 + def _assemble_patch_grid(
520 + self, features: torch.Tensor, crop_shape: torch.Tensor
521 + ) -> torch.Tensor:
522 + """Assemble projected patches into a 2-D tile grid with newline columns."""
523 _, hw, dim = features.shape
524 patch_side = int(hw**0.5)
509 -
525 width_tiles = int(crop_shape[0].item())
526 height_tiles = int(crop_shape[1].item())
527
@@ -519,7 +534,6 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
534 height_tiles * patch_side, 1, dim
535 )
536 features = torch.cat([features, newline], dim=1)
522 -
537 return features.view(-1, dim)
538
539 def _pixel_values_to_embedding(
@@ -614,3 +628,359 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
628 connector="projector",
629 tower_model=["sam_model", "vision_model"],
630 )
631 +
632 + # -- Fixed spatial constants (computed from BASE_SIZE / IMAGE_SIZE) --
633 +
634 + @property
635 + def image_side(self) -> int:
636 + """Number of output grid cells per spatial dim for a global image."""
637 + return math.ceil((BASE_SIZE // 16) / 4) # 16
638 +
639 + @property
640 + def global_image_output_token(self) -> int:
641 + """Tokens per global image (grid + one newline per row)."""
642 + return self.image_side * (self.image_side + 1) # 272
643 +
644 + @property
645 + def patch_side(self) -> int:
646 + """Number of output grid cells per spatial dim for a local patch."""
647 + return math.ceil((IMAGE_SIZE // 16) / 4) # 10
648 +
649 + @property
650 + def single_patch_output_token(self) -> int:
651 + """Tokens per local patch (square grid, no newlines)."""
652 + return self.patch_side * self.patch_side # 100
653 +
654 + # -- SupportsEncoderCudaGraph protocol methods --
655 +
656 + def _get_num_input_output_tokens(
657 + self,
658 + image_spatial_crop: torch.Tensor | None = None,
659 + ) -> tuple[int, int, int, int]:
660 + """
661 + Return (num_input_tokens, num_output_tokens, global_output_token,
662 + local_output_token) for a single image described by
663 + ``image_spatial_crop``.
664 + """
665 + is_tiled = False
666 + if image_spatial_crop is not None:
667 + is_tiled = image_spatial_crop[0] > 1 or image_spatial_crop[1] > 1
668 +
669 + # Compute input size:
670 + global_input_side = BASE_SIZE // 16 # 64
671 + local_input_side = IMAGE_SIZE // 16 # 40
672 + num_input_tokens = global_input_side**2
673 +
674 + if is_tiled:
675 + num_patches = image_spatial_crop.prod(dim=-1)
676 + num_input_tokens += num_patches * (local_input_side**2)
677 +
678 + global_output_token = self.global_image_output_token
679 + num_output_tokens = global_output_token
680 +
681 + local_output_token = 0
682 + if is_tiled:
683 + local_output_token = num_patches * self.single_patch_output_token
684 + num_output_tokens += local_output_token
685 +
686 + return (
687 + num_input_tokens,
688 + num_output_tokens,
689 + global_output_token,
690 + local_output_token,
691 + )
692 +
693 + def get_encoder_cudagraph_config(self):
694 + return EncoderCudaGraphConfig(
695 + modalities=["image"],
696 + buffer_keys=["pixel_values"],
697 + out_hidden_size=self.projector_config.n_embed,
698 + enable_dual_path_graph=True,
699 + global_token_per_image=self.global_image_output_token,
700 + local_token_per_patch=self.single_patch_output_token,
701 + )
702 +
703 + def get_encoder_cudagraph_budget_range(
704 + self,
705 + vllm_config,
706 + ) -> tuple[int, int]:
707 + # Min budget: at least one global image with newline tokens (without patches).
708 + min_budget = self.global_image_output_token
709 + max_budget = min(
710 + vllm_config.scheduler_config.max_num_batched_tokens,
711 + self.model_config.max_model_len,
712 + )
713 + return (min_budget, max_budget)
714 +
715 + def get_encoder_cudagraph_item_specs(
716 + self,
717 + mm_kwargs: dict[str, Any],
718 + ) -> list[EncoderItemSpec]:
719 + item_specs = []
720 + for image_spatial_crop in mm_kwargs["images_spatial_crop"]:
721 + (
722 + num_input_tokens,
723 + num_output_tokens,
724 + global_output_token,
725 + local_output_token,
726 + ) = self._get_num_input_output_tokens(image_spatial_crop)
727 + item_specs.append(
728 + EncoderItemSpec(
729 + input_size=num_input_tokens,
730 + output_tokens=num_output_tokens,
731 + global_output_tokens=global_output_token,
732 + local_output_tokens=local_output_token,
733 + )
734 + )
735 + return item_specs
736 +
737 + def select_encoder_cudagraph_items(
738 + self,
739 + mm_kwargs: dict[str, Any],
740 + indices: list[int],
741 + ) -> dict[str, Any]:
742 + pixel_values = mm_kwargs["pixel_values"]
743 + images_crop = mm_kwargs["images_crop"]
744 + images_spatial_crop = mm_kwargs["images_spatial_crop"]
745 +
746 + if len(indices) == 0:
747 + return {
748 + "pixel_values": pixel_values[:0],
749 + "images_crop": images_crop[:0],
750 + "images_spatial_crop": images_spatial_crop[:0],
751 + }
752 +
753 + is_tiled = (images_spatial_crop[:, 0] > 1) | (images_spatial_crop[:, 1] > 1)
754 + patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
755 + cum_patches = [0]
756 + for num_patches in patches_per_image:
757 + cum_patches.append(cum_patches[-1] + int(num_patches))
758 +
759 + selected_pv = pixel_values[indices]
760 + selected_ic = torch.cat(
761 + [images_crop[cum_patches[i] : cum_patches[i + 1]] for i in indices]
762 + )
763 + selected_sp = images_spatial_crop[indices]
764 +
765 + return {
766 + "pixel_values": selected_pv,
767 + "images_crop": selected_ic,
768 + "images_spatial_crop": selected_sp,
769 + }
770 +
771 + def prepare_encoder_cudagraph_capture_inputs(
772 + self,
773 + token_budget: int,
774 + max_batch_size: int,
775 + max_frames_per_batch: int,
776 + device: torch.device,
777 + dtype: torch.dtype,
778 + path: str = "default",
779 + ):
780 + assert path in ("global", "local")
781 +
782 + if path == "global":
783 + max_num_images = token_budget // self.global_image_output_token
784 + max_batch_size = min(max_batch_size, max_num_images)
785 + dummy_pixel_values = torch.randn(
786 + max_batch_size,
787 + 3,
788 + BASE_SIZE,
789 + BASE_SIZE,
790 + device=device,
791 + dtype=dtype,
792 + )
793 + values = {"pixel_values": dummy_pixel_values}
794 + else:
795 + max_num_patches = token_budget // self.single_patch_output_token
796 + dummy_images_crop = torch.randn(
797 + max_num_patches,
798 + 3,
799 + IMAGE_SIZE,
800 + IMAGE_SIZE,
801 + device=device,
802 + dtype=dtype,
803 + )
804 + values = {"images_crop": dummy_images_crop}
805 +
806 + return EncoderCudaGraphCaptureInputs(values=values)
807 +
808 + def prepare_encoder_cudagraph_replay_buffers(
809 + self,
810 + mm_kwargs: dict[str, Any],
811 + max_batch_size: int,
812 + max_frames_per_batch: int,
813 + path: str = "default",
814 + ):
815 + assert path in ("global", "local")
816 +
817 + if path == "global":
818 + values = {"pixel_values": mm_kwargs["pixel_values"]}
819 + else:
820 + values = {"images_crop": mm_kwargs["images_crop"]}
821 +
822 + return EncoderCudaGraphReplayBuffers(values=values)
823 +
824 + def _batched_encoder_forward_global_path(
825 + self,
826 + pixel_values: torch.Tensor,
827 + ) -> torch.Tensor:
828 + """
829 + Encode batched global images with newline tokens inserted.
830 + Output shape: ``[B * 272, n_embed]``.
831 + """
832 + bsz = pixel_values.shape[0]
833 + global_features_1 = self.sam_model(pixel_values)
834 + global_features_2 = self.vision_model(pixel_values, global_features_1)
835 + features = torch.cat(
836 + (
837 + global_features_2[:, 1:],
838 + global_features_1.flatten(2).permute(0, 2, 1),
839 + ),
840 + dim=-1,
841 + )
842 + features = self.projector(features)
843 + side = self.image_side
844 + dim = features.shape[-1]
845 + features = features.view(bsz, side, side, dim)
846 + newline = self.image_newline.view(1, 1, 1, dim).expand(bsz, side, 1, dim)
847 + features = torch.cat([features, newline], dim=2)
848 + return features.view(-1, dim)
849 +
850 + def _batched_encoder_forward_local_path(
851 + self,
852 + images_crop: torch.Tensor,
853 + ) -> torch.Tensor:
854 + """
855 + Encode local patches without newline insertion (newlines are added later
856 + in ``postprocess_encoder_output`` via ``_assemble_patch_grid``).
857 + Output shape: ``[P * 100, n_embed]``.
858 + """
859 + features_1 = self.sam_model(images_crop)
860 + features_2 = self.vision_model(images_crop, features_1)
861 + features = torch.cat(
862 + (
863 + features_2[:, 1:],
864 + features_1.flatten(2).permute(0, 2, 1),
865 + ),
866 + dim=-1,
867 + )
868 + features = self.projector(features)
869 + return features.view(-1, features.shape[-1])
870 +
871 + def encoder_cudagraph_forward(
872 + self,
873 + values: dict[str, torch.Tensor],
874 + path: str = "default",
875 + ) -> torch.Tensor:
876 + assert path in ("global", "local")
877 +
878 + if path == "global":
879 + pixel_values = values["pixel_values"]
880 + return self._batched_encoder_forward_global_path(pixel_values)
881 + else:
882 + images_crop = values["images_crop"]
883 + return self._batched_encoder_forward_local_path(images_crop)
884 +
885 + def encoder_eager_forward(
886 + self,
887 + mm_kwargs: dict[str, Any],
888 + path: str = "default",
889 + ) -> torch.Tensor:
890 + """Eager encoder forward with optional per-path execution.
891 +
892 + ``path="default"``: full forward (global + local + assembly).
893 + ``path="global"``: global-only batched forward with newlines.
894 + ``path="local"``: local-only batched forward without newlines.
895 + """
896 + if path == "default":
897 + # Original eager implementation: process each image one by one
898 + # (with both global and local paths) and concatenate results.
899 + image_input = DeepseekOCRImagePixelInputs(
900 + type="pixel_values",
901 + data=mm_kwargs["pixel_values"],
902 + images_crop=mm_kwargs["images_crop"],
903 + images_spatial_crop=mm_kwargs["images_spatial_crop"],
904 + )
905 + vision_embeddings = self._process_image_input(image_input)
906 + return torch.cat(vision_embeddings, dim=0)
907 +
908 + assert path in ("global", "local")
909 + if path == "global":
910 + pixel_values = mm_kwargs["pixel_values"]
911 + return self._batched_encoder_forward_global_path(pixel_values)
912 + else:
913 + images_crop = mm_kwargs["images_crop"]
914 + return self._batched_encoder_forward_local_path(images_crop)
915 +
916 + def postprocess_encoder_output(
917 + self,
918 + output: torch.Tensor,
919 + indices: list[int],
920 + per_item_out_tokens: list[int],
921 + dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
922 + clone: bool = False,
923 + batch_mm_kwargs: dict[str, Any] | None = None,
924 + local_output: torch.Tensor | None = None,
925 + ) -> None:
926 + """
927 + Assemble per-image embeddings from global and local encoder outputs.
928 +
929 + ``output`` contains global-image features with newlines already
930 + inserted (from CUDA graph replay or eager fallback):
931 + ``[B * 272, n_embed]``.
932 +
933 + ``local_output`` contains local-patch features without
934 + newlines (from CUDA graph replay or eager fallback):
935 + ``[P * 100, n_embed]``. May be ``None`` if no patches in batch.
936 +
937 + This method:
938 + 1. Splits ``output`` into per-image global portions.
939 + 2. Splits ``local_output`` into per-image patch groups.
940 + 3. For each image: assembles patch grid with newlines via
941 + ``_assemble_patch_grid``, then concatenates
942 + ``[local_tiled, global, view_seperator]``.
943 + """
944 + bsz = len(indices)
945 + n_embed = output.shape[-1]
946 +
947 + images_spatial_crop = batch_mm_kwargs["images_spatial_crop"]
948 + is_tiled = (images_spatial_crop[:, 0] > 1) | (images_spatial_crop[:, 1] > 1)
949 + num_patches = [
950 + int(np) for np in torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
951 + ]
952 + total_patches = sum(num_patches)
953 +
954 + global_part = output[: bsz * self.global_image_output_token].reshape(
955 + bsz, self.global_image_output_token, n_embed
956 + )
957 +
958 + # Split local output into per-patch groups.
959 + local_flat = None
960 + if total_patches > 0 and local_output is not None:
961 + local_flat = local_output[: total_patches * self.single_patch_output_token]
962 + local_flat = local_flat.reshape(
963 + total_patches, self.single_patch_output_token, n_embed
964 + )
965 +
966 + cur_patch = 0
967 + for i, idx in enumerate(indices):
968 + num_patch = num_patches[i]
969 + single_image_output: list[torch.Tensor] = []
970 +
971 + # 1. Process local patches: assemble tile grid, add 1 newline per row.
972 + if num_patch > 0 and local_flat is not None:
973 + patches = local_flat[cur_patch : cur_patch + num_patch]
974 + cur_patch += num_patch
975 + single_image_output.append(
976 + self._assemble_patch_grid(patches, images_spatial_crop[i])
977 + )
978 +
979 + # 2. Global image: newlines already inserted.
980 + single_image_output.append(global_part[i])
981 +
982 + # 3. Add view separator for each image.
983 + single_image_output.append(self.view_seperator[None, :])
984 +
985 + # 4. Save final outputs for each image.
986 + dest[idx] = torch.cat(single_image_output, dim=0)
vllm/model_executor/models/glm4_1v.py
+4
@@ -2068,6 +2068,7 @@ class Glm4vForConditionalGeneration(
2068 max_frames_per_batch: int,
2069 device: torch.device,
2070 dtype: torch.dtype,
2071 + path: str = "default",
2072 ):
2073 from vllm.v1.worker.encoder_cudagraph_defs import (
2074 EncoderCudaGraphCaptureInputs,
@@ -2145,6 +2146,7 @@ class Glm4vForConditionalGeneration(
2146 mm_kwargs: dict[str, Any],
2147 max_batch_size: int,
2148 max_frames_per_batch: int,
2149 + path: str = "default",
2150 ):
2151 modality = self.get_input_modality(mm_kwargs)
2152 grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
@@ -2170,6 +2172,7 @@ class Glm4vForConditionalGeneration(
2172 def encoder_cudagraph_forward(
2173 self,
2174 values: dict[str, torch.Tensor],
2175 + path: str = "default",
2176 ) -> torch.Tensor:
2177 pixel_values = values.pop("pixel_values")
2178 metadata = values
@@ -2178,6 +2181,7 @@ class Glm4vForConditionalGeneration(
2181 def encoder_eager_forward(
2182 self,
2183 mm_kwargs: dict[str, Any],
2184 + path: str = "default",
2185 ) -> torch.Tensor:
2186 pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
2187 grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
vllm/model_executor/models/interfaces.py
+5
@@ -1623,6 +1623,7 @@ class SupportsEncoderCudaGraph(Protocol):
1623 dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
1624 clone: bool = False,
1625 batch_mm_kwargs: dict[str, Any] | None = None,
1626 + local_output: torch.Tensor | None = None,
1627 ) -> None:
1628 """
1629 Post-process encoder output, directly call scatter_output_slices by default.
@@ -1643,6 +1644,7 @@ class SupportsEncoderCudaGraph(Protocol):
1644 max_frames_per_batch: int,
1645 device: torch.device,
1646 dtype: torch.dtype,
1647 + path: str = "default",
1648 ) -> "EncoderCudaGraphCaptureInputs":
1649 """Create dummy inputs and buffers for CUDA graph capture."""
1650 ...
@@ -1652,6 +1654,7 @@ class SupportsEncoderCudaGraph(Protocol):
1654 mm_kwargs: dict[str, Any],
1655 max_batch_size: int,
1656 max_frames_per_batch: int,
1657 + path: str = "default",
1658 ) -> "EncoderCudaGraphReplayBuffers":
1659 """Compute buffer values from actual batch inputs for replay."""
1660 ...
@@ -1659,6 +1662,7 @@ class SupportsEncoderCudaGraph(Protocol):
1662 def encoder_cudagraph_forward(
1663 self,
1664 inputs: dict[str, torch.Tensor],
1665 + path: str = "default",
1666 ) -> torch.Tensor:
1667 """Run the encoder forward pass with precomputed buffers.
1668
@@ -1669,6 +1673,7 @@ class SupportsEncoderCudaGraph(Protocol):
1673 def encoder_eager_forward(
1674 self,
1675 mm_kwargs: dict[str, Any],
1676 + path: str = "default",
1677 ) -> torch.Tensor:
1678 """Run the encoder forward pass without precomputed buffers.
1679
vllm/model_executor/models/internvl.py
+4
@@ -1004,6 +1004,7 @@ class InternVLChatModel(
1004 max_frames_per_batch: int,
1005 device: torch.device,
1006 dtype: torch.dtype,
1007 + path: str = "default",
1008 ):
1009 from vllm.v1.worker.encoder_cudagraph_defs import (
1010 EncoderCudaGraphCaptureInputs,
@@ -1026,6 +1027,7 @@ class InternVLChatModel(
1027 mm_kwargs: dict[str, Any],
1028 max_batch_size: int,
1029 max_frames_per_batch: int,
1030 + path: str = "default",
1031 ):
1032 from vllm.v1.worker.encoder_cudagraph_defs import (
1033 EncoderCudaGraphReplayBuffers,
@@ -1042,6 +1044,7 @@ class InternVLChatModel(
1044 def encoder_cudagraph_forward(
1045 self,
1046 values: dict[str, torch.Tensor],
1047 + path: str = "default",
1048 ) -> torch.Tensor:
1049 # The graph is always captured with pixel_values_flat as the input
1050 # buffer. During video replay the manager copies video tiles into
@@ -1054,6 +1057,7 @@ class InternVLChatModel(
1057 def encoder_eager_forward(
1058 self,
1059 mm_kwargs: dict[str, Any],
1060 + path: str = "default",
1061 ) -> torch.Tensor:
1062 if self.get_input_modality(mm_kwargs) == "image":
1063 pixel_values = mm_kwargs["pixel_values_flat"]
vllm/model_executor/models/lfm2_vl.py
+4
@@ -1122,6 +1122,7 @@ class Lfm2VLForConditionalGeneration(
1122 max_frames_per_batch: int,
1123 device: torch.device,
1124 dtype: torch.dtype,
1125 + path: str = "default",
1126 ):
1127 from vllm.v1.worker.encoder_cudagraph_defs import (
1128 EncoderCudaGraphCaptureInputs,
@@ -1167,6 +1168,7 @@ class Lfm2VLForConditionalGeneration(
1168 mm_kwargs: dict[str, Any],
1169 max_batch_size: int,
1170 max_frames_per_batch: int,
1171 + path: str = "default",
1172 ):
1173 from vllm.v1.worker.encoder_cudagraph_defs import (
1174 EncoderCudaGraphReplayBuffers,
@@ -1181,6 +1183,7 @@ class Lfm2VLForConditionalGeneration(
1183 def encoder_cudagraph_forward(
1184 self,
1185 values: dict[str, torch.Tensor],
1186 + path: str = "default",
1187 ) -> torch.Tensor:
1188 embeddings = self.vision_tower.vision_model.embeddings
1189 pixel_values = values["pixel_values_packed"].to(
@@ -1208,6 +1211,7 @@ class Lfm2VLForConditionalGeneration(
1211 def encoder_eager_forward(
1212 self,
1213 mm_kwargs: dict[str, Any],
1214 + path: str = "default",
1215 ) -> torch.Tensor:
1216 image_input = LFM2VLImageInputs(
1217 type="pixel_values",
vllm/model_executor/models/mllama4.py
+4
@@ -926,6 +926,7 @@ class Llama4ForConditionalGeneration(
926 max_frames_per_batch: int,
927 device: torch.device,
928 dtype: torch.dtype,
929 + path: str = "default",
930 ):
931 from vllm.v1.worker.encoder_cudagraph_defs import (
932 EncoderCudaGraphCaptureInputs,
@@ -954,6 +955,7 @@ class Llama4ForConditionalGeneration(
955 mm_kwargs: dict[str, Any],
956 max_batch_size: int,
957 max_frames_per_batch: int,
958 + path: str = "default",
959 ):
960 from vllm.v1.worker.encoder_cudagraph_defs import (
961 EncoderCudaGraphReplayBuffers,
@@ -966,6 +968,7 @@ class Llama4ForConditionalGeneration(
968 def encoder_cudagraph_forward(
969 self,
970 inputs: dict[str, torch.Tensor],
971 + path: str = "default",
972 ) -> torch.Tensor:
973 return self.encode_image_chunks(
974 inputs["pixel_values"],
@@ -975,6 +978,7 @@ class Llama4ForConditionalGeneration(
978 def encoder_eager_forward(
979 self,
980 mm_kwargs: dict[str, Any],
981 + path: str = "default",
982 ) -> torch.Tensor:
983 return self.encode_image_chunks(
984 mm_kwargs["pixel_values"],
vllm/model_executor/models/qwen2_5_vl.py
+4
@@ -1856,6 +1856,7 @@ class Qwen2_5_VLForConditionalGeneration(
1856 max_frames_per_batch: int,
1857 device: torch.device,
1858 dtype: torch.dtype,
1859 + path: str = "default",
1860 ):
1861 from vllm.v1.worker.encoder_cudagraph_defs import (
1862 EncoderCudaGraphCaptureInputs,
@@ -1952,6 +1953,7 @@ class Qwen2_5_VLForConditionalGeneration(
1953 mm_kwargs: dict[str, Any],
1954 max_batch_size: int,
1955 max_frames_per_batch: int,
1956 + path: str = "default",
1957 ):
1958 modality = self.get_input_modality(mm_kwargs)
1959 grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
@@ -1983,6 +1985,7 @@ class Qwen2_5_VLForConditionalGeneration(
1985 def encoder_cudagraph_forward(
1986 self,
1987 values: dict[str, torch.Tensor],
1988 + path: str = "default",
1989 ) -> torch.Tensor:
1990 pixel_values = values.pop("pixel_values")
1991 metadata = values
@@ -1991,6 +1994,7 @@ class Qwen2_5_VLForConditionalGeneration(
1994 def encoder_eager_forward(
1995 self,
1996 mm_kwargs: dict[str, Any],
1997 + path: str = "default",
1998 ) -> torch.Tensor:
1999 pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
2000 grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
vllm/model_executor/models/qwen2_vl.py
+4
@@ -1579,6 +1579,7 @@ class Qwen2VLForConditionalGeneration(
1579 max_frames_per_batch: int,
1580 device: torch.device,
1581 dtype: torch.dtype,
1582 + path: str = "default",
1583 ):
1584 from vllm.v1.worker.encoder_cudagraph_defs import (
1585 EncoderCudaGraphCaptureInputs,
@@ -1644,6 +1645,7 @@ class Qwen2VLForConditionalGeneration(
1645 mm_kwargs: dict[str, Any],
1646 max_batch_size: int,
1647 max_frames_per_batch: int,
1648 + path: str = "default",
1649 ) -> EncoderCudaGraphReplayBuffers:
1650 modality = self.get_input_modality(mm_kwargs)
1651 grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
@@ -1667,6 +1669,7 @@ class Qwen2VLForConditionalGeneration(
1669 def encoder_cudagraph_forward(
1670 self,
1671 values: dict[str, torch.Tensor],
1672 + path: str = "default",
1673 ) -> torch.Tensor:
1674 pixel_values = values.pop("pixel_values")
1675 metadata = values
@@ -1675,6 +1678,7 @@ class Qwen2VLForConditionalGeneration(
1678 def encoder_eager_forward(
1679 self,
1680 mm_kwargs: dict[str, Any],
1681 + path: str = "default",
1682 ) -> torch.Tensor:
1683 pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
1684 grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
vllm/model_executor/models/qwen3_vl.py
+4
@@ -1918,6 +1918,7 @@ class Qwen3VLForConditionalGeneration(
1918 max_frames_per_batch: int,
1919 device: torch.device,
1920 dtype: torch.dtype,
1921 + path: str = "default",
1922 ):
1923 from vllm.v1.worker.encoder_cudagraph_defs import (
1924 EncoderCudaGraphCaptureInputs,
@@ -1998,6 +1999,7 @@ class Qwen3VLForConditionalGeneration(
1999 mm_kwargs: dict[str, Any],
2000 max_batch_size: int,
2001 max_frames_per_batch: int,
2002 + path: str = "default",
2003 ):
2004 modality = self.get_input_modality(mm_kwargs)
2005 grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
@@ -2023,6 +2025,7 @@ class Qwen3VLForConditionalGeneration(
2025 def encoder_cudagraph_forward(
2026 self,
2027 values: dict[str, torch.Tensor],
2028 + path: str = "default",
2029 ) -> torch.Tensor:
2030 pixel_values = values.pop("pixel_values")
2031 metadata = values
@@ -2031,6 +2034,7 @@ class Qwen3VLForConditionalGeneration(
2034 def encoder_eager_forward(
2035 self,
2036 mm_kwargs: dict[str, Any],
2037 + path: str = "default",
2038 ) -> torch.Tensor:
2039 pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
2040 grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
vllm/model_executor/models/step3_vl.py
+5
@@ -804,6 +804,7 @@ class Step3VLForConditionalGeneration(
804 max_frames_per_batch: int,
805 device: torch.device,
806 dtype: torch.dtype,
807 + path: str = "default",
808 ):
809 from vllm.v1.worker.encoder_cudagraph_defs import (
810 EncoderCudaGraphCaptureInputs,
@@ -857,6 +858,7 @@ class Step3VLForConditionalGeneration(
858 def encoder_cudagraph_forward(
859 self,
860 values: dict[str, torch.Tensor],
861 + path: str = "default",
862 ) -> torch.Tensor:
863 # Graph captures only the compute (vision model + conv projector).
864 # Per-item merge happens CPU-side in finalize_encoder_cudagraph_output
@@ -884,6 +886,7 @@ class Step3VLForConditionalGeneration(
886 def encoder_eager_forward(
887 self,
888 mm_kwargs: dict[str, Any],
889 + path: str = "default",
890 ) -> torch.Tensor:
891 image_input = Step3VLImagePixelInputs(
892 type="pixel_values",
@@ -902,6 +905,7 @@ class Step3VLForConditionalGeneration(
905 dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
906 clone: bool = False,
907 batch_mm_kwargs: dict[str, Any] | None = None,
908 + local_output: torch.Tensor | None = None,
909 ):
910 """CPU-side per-item merge after graph replay.
911
@@ -959,6 +963,7 @@ class Step3VLForConditionalGeneration(
963 mm_kwargs: dict[str, Any],
964 max_batch_size: int,
965 max_frames_per_batch: int,
966 + path: str = "default",
967 ):
968 from vllm.v1.worker.encoder_cudagraph_defs import (
969 EncoderCudaGraphReplayBuffers,
vllm/v1/worker/encoder_cudagraph.py
+254 -24
@@ -152,20 +152,40 @@ class EncoderCudaGraphManager:
152 and vllm_config.parallel_config.tensor_parallel_size > 1
153 )
154
155 - self.budget_graphs: dict[int, BudgetGraphMetadata] = {}
155 + self.budget_graphs: dict[str, dict[int, BudgetGraphMetadata]] = {}
156 self.graph_pool: Any | None = None
157 self.graph_hits = 0
158 self.graph_misses = 0
159 self.log_stats_interval = 100
160
161 - logger.info(
162 - "EncoderCudaGraphManager initialized with "
163 - "budgets=%s, max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",
164 - self.token_budgets,
165 - self.max_batch_size,
166 - self.max_frames_per_batch,
167 - self.use_dp,
168 - )
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 + )
189
190 @staticmethod
191 def _generate_budgets(min_budget: int, max_budget: int) -> list[int]:
@@ -186,25 +206,50 @@ class EncoderCudaGraphManager:
206
207 def clear(self) -> None:
208 """Release captured encoder CUDA graphs and the manager-local pool."""
189 - self.budget_graphs.clear()
209 + for graph_set in self.budget_graphs.values():
210 + graph_set.clear()
211 self.graph_pool = None
212
213 def capture(self, graph_pool: Any):
214 """Capture CUDA graphs for all token budgets."""
215 self.graph_pool = graph_pool
216
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)
234
235 logger.info(
236 "Encoder CUDA graph capture complete. Captured %d budget graphs.",
201 - len(self.budget_graphs),
237 + len(self.budget_graphs["default"]),
238 )
239
240 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)
244
207 - def _capture_budget_graph(self, token_budget: int):
245 + 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.
248 + if path not in self.budget_graphs:
249 + self.budget_graphs[path] = {}
250 + return self.budget_graphs[path]
251 +
252 + def _capture_budget_graph(self, token_budget: int, path: str = "default"):
253 """Capture CUDA graph for a single token budget."""
254 logger.debug(
255 "Capturing encoder cudagraph for budget=%d, max_batch_size=%d, "
@@ -214,26 +259,29 @@ class EncoderCudaGraphManager:
259 self.max_frames_per_batch,
260 )
261
262 + graph_set = self._get_graph_set(path)
263 +
264 capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs(
265 token_budget,
266 self.max_batch_size,
267 self.max_frames_per_batch,
268 self.device,
269 self.dtype,
270 + path,
271 )
272
273 values = capture_inputs.values
274
275 with torch.inference_mode():
228 - output = self.model.encoder_cudagraph_forward({**values})
276 + output = self.model.encoder_cudagraph_forward({**values}, path=path)
277 output_buffer = torch.empty_like(output)
278
279 graph = torch.cuda.CUDAGraph()
280 with torch.inference_mode(), torch.cuda.graph(graph, pool=self.graph_pool):
233 - output = self.model.encoder_cudagraph_forward({**values})
281 + output = self.model.encoder_cudagraph_forward({**values}, path=path)
282 output_buffer.copy_(output)
283
236 - self.budget_graphs[token_budget] = BudgetGraphMetadata(
284 + graph_set[token_budget] = BudgetGraphMetadata(
285 token_budget=token_budget,
286 max_batch_size=self.max_batch_size,
287 max_frames_per_batch=self.max_frames_per_batch,
@@ -243,14 +291,15 @@ class EncoderCudaGraphManager:
291 )
292
293 def _find_smallest_fitting_budget_given_tokens(
246 - self, total_tokens: int
294 + self, total_tokens: int, budgets: list[int] | None = None
295 ) -> int | None:
296 """Find smallest budget >= total_tokens.
297
298 Returns:
299 Token budget if found, None if no fitting budget.
300 """
253 - for budget in self.token_budgets:
301 + budgets = budgets if budgets is not None else self.token_budgets
302 + for budget in budgets:
303 if budget >= total_tokens:
304 return budget
305 return None
@@ -275,35 +324,40 @@ class EncoderCudaGraphManager:
324 self,
325 mm_kwargs: dict[str, Any],
326 token_budget: int,
327 + path: str = "default",
328 ) -> torch.Tensor | None:
329 """Execute budget graph.
330
331 Args:
332 mm_kwargs: Multimodal inputs for the batch.
333 token_budget: Token budget to use.
284 -
334 + path: Path for the graph. Should be one of ["default", "global", "local"].
335 Returns:
336 Encoder outputs, or None if graph not captured.
337 """
338 + graph_set = self._get_graph_set(path)
339 num_items = len(self._get_item_specs(mm_kwargs))
289 - if token_budget not in self.budget_graphs:
340 +
341 + if token_budget not in graph_set:
342 self.graph_misses += num_items
343 return None
344
293 - graph_meta = self.budget_graphs[token_budget]
345 + graph_meta = graph_set[token_budget]
346
347 replay = self.model.prepare_encoder_cudagraph_replay_buffers(
348 mm_kwargs,
349 self.max_batch_size,
350 self.max_frames_per_batch,
351 + path,
352 )
353
301 - # Copy metadata buffers using keys from config.buffer_keys.
302 - for key in self.config.buffer_keys:
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.
357 + for key, buf in graph_meta.input_buffers.items():
358 src = replay.values.get(key)
359 if src is None:
360 continue
306 - buf = graph_meta.input_buffers[key]
361 if src.ndim == 0:
362 buf.copy_(src)
363 else:
@@ -329,6 +383,11 @@ class EncoderCudaGraphManager:
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.
@@ -340,6 +399,15 @@ class EncoderCudaGraphManager:
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)."""
411 item_specs = self._get_item_specs(mm_kwargs)
412 num_items = len(item_specs)
413 max_budget = self.token_budgets[-1]
@@ -441,6 +509,166 @@ class EncoderCudaGraphManager:
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)]
511
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)
525 +
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 + )
584 +
585 + outputs_by_orig_idx: dict[int, torch.Tensor] = {}
586 +
587 + for batch_orig_indices, global_budget, local_budget in batches:
588 + batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
589 + mm_kwargs, batch_orig_indices
590 + )
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
619 +
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 + )
629 +
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"
642 + )
643 +
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
659 +
660 + self.model.postprocess_encoder_output(
661 + global_output,
662 + batch_orig_indices,
663 + per_item_global_tokens,
664 + outputs_by_orig_idx,
665 + clone=True,
666 + batch_mm_kwargs=batch_mm_kwargs,
667 + local_output=local_output,
668 + )
669 +
670 + return [outputs_by_orig_idx[i] for i in range(num_items)]
671 +
672 def _dp_shard(
673 self,
674 mm_kwargs: dict[str, Any],
@@ -626,10 +854,12 @@ class EncoderCudaGraphManager:
854 total_requests = self.graph_hits + self.graph_misses
855 hit_rate = self.graph_hits / total_requests if total_requests > 0 else 0.0
856
857 + num_budgets = sum(len(g) for g in self.budget_graphs.values())
858 +
859 return {
860 "graph_hits": self.graph_hits,
861 "graph_misses": self.graph_misses,
862 "hit_rate": hit_rate,
633 - "num_budgets": len(self.budget_graphs),
863 + "num_budgets": num_budgets,
864 "token_budgets": self.token_budgets,
865 }
vllm/v1/worker/encoder_cudagraph_defs.py
+20
@@ -26,6 +26,14 @@ 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."""
32 +
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."""
36 +
37
38 @dataclass
39 class EncoderCudaGraphConfig:
@@ -60,6 +68,18 @@ class EncoderCudaGraphConfig:
68 Only relevant when "video" is in ``modalities``.
69 Image-only models can use the default of 1."""
70
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."""
82 +
83
84 @dataclass
85 class EncoderCudaGraphCaptureInputs: