@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 315 lines 20.8 KB
Rendered Raw
1 # Vision Encoder (ViT) CUDA Graphs
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
10 ## Compatibility Matrix
11
12 !!! note
13 The symbols used below have the following meanings:
14
15 - ✅ = Full compatibility
16 - 🟠 = Partial compatibility
17 - ❌ = No compatibility
18 - ❔ = Unknown or TBD
19
20 ### Model x Feature
21
22 | Architecture | Models | CG for Image | CG for Video | Multi-Path Graph |
23 | ------------ | ------ | ------------ | ------------ | --------------- |
24 | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
25 | `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ |
26 | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |
27 | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
28 | `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ |
29 | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
30 | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
31 | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
32 | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ |
33 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ |
34 | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ |
35 | `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ |
36 | `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ |
37 | `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ |
38 | `MiniCPMV` | `MiniCPMV2.5` | ✅︎ | ❌︎ | ❌︎ |
39 | `MiniCPMV` | `MiniCPMV2.6`, `MiniCPMV4.0` | ✅︎ | ✅︎ | ❌︎ |
40
41 ### Model x Hardware
42
43 | Architecture | NV Blackwell | NV Ampere | AMD MI300X | AMD MI350X / MI355X |
44 | ------------ | ---------------- | ------------- | -------------- | --------------------- |
45 | `DeepseekOCRForCausalLM` | ✅︎ | ✅︎ | ❔ | ✅︎ |
46 | `Ernie4_5_VLMoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
47 | `Gemma3ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
48 | `Glm4vForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
49 | `Gemma4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
50 | `InternVLChatModel` | ✅︎ | ✅︎ | ❔ | ✅︎ |
51 | `KimiVLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
52 | `Llama4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
53 | `Qwen2VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
54 | `Qwen2_5_VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
55 | `Qwen3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
56 | `Qwen3_5ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
57 | `Qwen3_5MoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
58 | `Step3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
59 | `MiniCPMV` | ✅︎ | ✅︎ | ❔ | ❔ |
60
61 !!! note
62 Encoder CUDA Graph has currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
63 For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested.
64 Encoder CUDA Graph has also been tested with AMD MI350X (gfx950) used `--mm-encoder-attn-backend=FLASH_ATTN` (the ROCm default).
65
66 ## Motivation
67
68 Vision encoder inference incurs CUDA kernel launch overhead on the host side. The overhead is more significant when the batch size is small or image size is small.
69
70 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.
71
72 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.
73
74 ## Design
75
76 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:
77
78 * [EncoderCudaGraphManager][vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager]: orchestrates capture, replay, greedy packing, and data-parallel execution for encoder CUDA Graphs.
79 * [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph]: a runtime-checkable protocol that models implement to opt-in to encoder CUDA Graphs.
80 * [EncoderItemSpec][vllm.v1.worker.encoder_cudagraph_defs.EncoderItemSpec]: describes a single encoder input item (image or video) with its input size and output token count.
81 * [BudgetGraphMetadata][vllm.v1.worker.encoder_cudagraph.BudgetGraphMetadata]: holds the captured CUDA Graph and its associated I/O buffers for a single token budget level.
82
83 ### Budget-based graph capture
84
85 Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g., `[2048, 4096, 8192, 13824]`). Each budget defines a fixed token capacity, and all budgets share the same maximum batch size (number of images). The `BudgetGraphMetadata` for each level stores the graph along with pre-allocated input, metadata, and output buffers:
86
87 ```python
88 @dataclass
89 class BudgetGraphMetadata:
90 token_budget: int
91 max_batch_size: int
92 max_frames_per_batch: int
93 graph: torch.cuda.CUDAGraph
94 input_buffers: dict[str, torch.Tensor] # e.g. pixel_values, embeddings, seq metadata
95 output_buffer: torch.Tensor # encoder hidden states
96 ```
97
98 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`.
99
100 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.
101
102 ### Greedy bin-packing at runtime
103
104 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.
105
106 For multi-path models, the same greedy packing loop constrains every configured path simultaneously (see [Multi-Path graph capture](#multi-path-graph-capture)).
107
108 For each graph replay:
109
110 1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs.
111 2. Zero the pre-allocated `input_buffers`, then slice-copy the replay values into them.
112 3. Replay the CUDA Graph.
113 4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays).
114
115 ### Multi-Path graph capture
116
117 `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"]`.
118
119 **Budget generation.** Each path gets a separate budget list. For DeepSeek-OCR:
120
121 * 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).
122 * 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).
123
124 Both lists are capped at the same `max_budget`.
125
126 **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:
127
128 * Sort images by total output tokens (global + local), smallest first.
129 * 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`.
130 * Once any path constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path.
131 * Repeat until all images are packed.
132
133 **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.
134
135 **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.
136
137 **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.
138
139 !!! note
140 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.
141
142 ### Data-parallel support
143
144 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`.
145
146 ### Video inference support
147
148 Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
149
150 !!! note
151 Video CUDA graphs are automatically disabled when video token pruning (EVS or VidCom2) is enabled, since pruning makes the token count data-dependent and incompatible with CUDA graph capture.
152
153 Mixed inputs (image+video) per prompt are also supported now.
154
155 ## Model integration via `SupportsEncoderCudaGraph`
156
157 Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods:
158
159 * `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video).
160 * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets.
161 * `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.
162 * `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding.
163 * `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.
164 * `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()`.
165 * `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).
166 * `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.
167 * `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.
168
169 !!! note
170 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.
171
172 ## Configuration
173
174 Four fields in `CompilationConfig` control encoder CUDA Graphs:
175
176 * `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.
177 * `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.
178 * `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`.
179 * `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).
180
181 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.
182
183 ## Usage guide
184
185 ### Image inference
186
187 Enable encoder CUDA Graphs via `compilation_config`:
188
189 ```bash
190 vllm serve Qwen/Qwen3-VL-32B \
191 --compilation-config '{"cudagraph_mm_encoder": true}'
192 ```
193
194 For `Llama 4` (image only):
195
196 ```bash
197 vllm serve meta-llama/Llama-4-Scout-17B-16E-Instruct \
198 --limit-mm-per-prompt '{"image": 1}' \
199 --compilation-config '{"cudagraph_mm_encoder": true}'
200 ```
201
202 With explicit budgets:
203
204 ```bash
205 vllm serve Qwen/Qwen3-VL-32B \
206 --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}'
207 ```
208
209 Python example:
210
211 ```python
212 import vllm
213
214 compilation_config = {
215 "cudagraph_mm_encoder": True,
216 # Optional: override auto-inferred budgets
217 # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
218 # "encoder_cudagraph_max_vision_items_per_batch": 8,
219 }
220
221 model = vllm.LLM(
222 model="Qwen/Qwen3-VL-32B",
223 compilation_config=compilation_config,
224 )
225 ```
226
227 The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets).
228
229 ### Video inference
230
231 Enable encoder CUDA Graphs via `compilation_config`:
232
233 ```bash
234 vllm serve Qwen/Qwen3-VL-32B \
235 --compilation-config '{"cudagraph_mm_encoder": true}'
236 ```
237
238 With explicit budgets:
239
240 ```bash
241 vllm serve Qwen/Qwen3-VL-32B \
242 --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
243 ```
244
245 Python example:
246
247 ```python
248 import vllm
249
250 compilation_config = {
251 "cudagraph_mm_encoder": True,
252 # Optional: override auto-inferred budgets
253 # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
254 # "encoder_cudagraph_max_vision_items_per_batch": 8,
255 # "encoder_cudagraph_max_frames_per_batch": 64,
256 }
257
258 model = vllm.LLM(
259 model="Qwen/Qwen3-VL-32B",
260 compilation_config=compilation_config,
261 )
262 ```
263
264 ## Benchmark Results
265
266 The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details.
267
268 ### Single GPU (1x GB200)
269
270 Model: `Qwen/Qwen3-VL-30B-A3B-Instruct`, dataset: `lmarena-ai/VisionArena-Chat` (3000 prompts, 300 warmup), `max_model_len=32768`.
271
272 | Backend | Mean latency improvement | P99 latency improvement |
273 | :------ | :----------------------- | :---------------------- |
274 | FLASH_ATTN | +11.8% (5.13→4.52ms) | +31.6% (9.16→6.26ms) |
275 | FLASHINFER | +19.6% (5.42→4.36ms) | +40.3% (10.87→6.49ms) |
276
277 To reproduce:
278
279 ```bash
280 vllm bench mm-processor \
281 --model Qwen/Qwen3-VL-30B-A3B-Instruct \
282 --dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
283 --num-prompts 3000 --num-warmups 300 \
284 --max-model-len 32768 --seed 42 \
285 --mm-encoder-attn-backend FLASH_ATTN \
286 --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
287 ```
288
289 ### Multi-GPU (4x GB200, TP=4, DP=4)
290
291 Model: `Qwen/Qwen3-VL-32B-Instruct`, dataset: `random-mm` (1000 prompts, 200 warmup, 20 images/request at 336x336), `max_model_len=8192`.
292
293 | Backend | Mean latency improvement | P99 latency improvement |
294 | :------ | :----------------------- | :---------------------- |
295 | FLASH_ATTN | +18.4% (28.39→23.16ms) | +14.0% (238.78→205.28ms) |
296 | FLASHINFER | +44.4% (23.24→12.91ms) | +84.9% (172.41→26.05ms) |
297
298 To reproduce:
299
300 ```bash
301 vllm bench mm-processor \
302 --model Qwen/Qwen3-VL-32B-Instruct \
303 --dataset-name random-mm \
304 --random-mm-base-items-per-request 20 \
305 --random-mm-num-mm-items-range-ratio 0.0 \
306 --random-mm-bucket-config '{"(336,336,1)": 1.0}' \
307 --num-prompts 1000 --num-warmups 200 \
308 --max-model-len 8192 --seed 42 \
309 --mm-encoder-attn-backend FLASHINFER \
310 --tensor-parallel-size 4 --mm-encoder-tp-mode data \
311 --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
312 ```
313
314 !!! note
315 Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061).