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

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 05:45 UTC

[Feature] ViT Full CUDA Graph (#35963)

Signed-off-by: Baorun Mu <bmu@nvidia.com>

Baorun (Lauren) Mu committed Mar 23, 2026 at 01:01 UTC f85e479e66167532feddc9aafd56562402def9d9
7 files changed +1584 -31
tests/v1/cudagraph/test_encoder_cudagraph.py new
+451
@@ -0,0 +1,451 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for EncoderCudaGraphManager.
4 +
5 +Test organization:
6 + No GPU required:
7 + - TestFindBudgetGraph — greedy budget selection logic
8 + - TestGetCumulativeStats — hit/miss rate statistics
9 + GPU required:
10 + - TestEncoderCudaGraphCaptureReplay — capture, replay, fallback, counters, chunking
11 +"""
12 +
13 +from typing import Any
14 +
15 +import pytest
16 +import torch
17 +
18 +from vllm.platforms import current_platform
19 +from vllm.v1.worker.gpu.mm.encoder_cudagraph import (
20 + EncoderCudaGraphManager,
21 +)
22 +from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
23 + EncoderCudaGraphCaptureInputs,
24 + EncoderCudaGraphConfig,
25 + EncoderCudaGraphReplayBuffers,
26 +)
27 +
28 +# ---------------------------------------------------------------------------
29 +# Helpers
30 +# ---------------------------------------------------------------------------
31 +
32 +
33 +def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
34 + """Create a minimal EncoderCudaGraphManager with only token_budgets set.
35 +
36 + Skips the parts of __init__ that require a real VllmConfig / model
37 + by patching the attributes directly after construction.
38 + """
39 + mgr = object.__new__(EncoderCudaGraphManager)
40 + mgr.token_budgets = sorted(budgets)
41 + mgr.max_batch_size = 16
42 + mgr.use_dp = False
43 + mgr.budget_graphs = {}
44 + mgr.graph_hits = 0
45 + mgr.graph_misses = 0
46 + mgr.log_stats_interval = 100
47 + return mgr
48 +
49 +
50 +# ---------------------------------------------------------------------------
51 +# _generate_budgets
52 +# ---------------------------------------------------------------------------
53 +
54 +
55 +class TestGenerateBudgets:
56 + """Auto-generate power-of-2 budgets from min to max."""
57 +
58 + def test_exact_powers_of_2(self):
59 + result = EncoderCudaGraphManager._generate_budgets(64, 1024)
60 + assert result == [64, 128, 256, 512, 1024]
61 +
62 + def test_max_not_power_of_2(self):
63 + result = EncoderCudaGraphManager._generate_budgets(64, 800)
64 + assert result == [64, 128, 256, 512, 800]
65 +
66 + def test_min_equals_max(self):
67 + result = EncoderCudaGraphManager._generate_budgets(64, 64)
68 + assert result == [64]
69 +
70 + def test_large_range(self):
71 + result = EncoderCudaGraphManager._generate_budgets(64, 8192)
72 + assert result == [64, 128, 256, 512, 1024, 2048, 4096, 8192]
73 +
74 +
75 +# ---------------------------------------------------------------------------
76 +# _find_smallest_fitting_budget_given_tokens
77 +# ---------------------------------------------------------------------------
78 +
79 +
80 +class TestFindBudgetGraph:
81 + """Budget greedy selection: smallest budget >= total_tokens."""
82 +
83 + @pytest.mark.parametrize(
84 + "total_tokens,budgets,expected",
85 + [
86 + # Exact match
87 + (2048, [2048, 4096, 8192], 2048),
88 + # Below smallest budget — picks smallest
89 + (100, [2048, 4096, 8192], 2048),
90 + # Zero tokens — picks smallest
91 + (0, [2048, 4096, 8192], 2048),
92 + # Between budgets — picks next one up
93 + (2049, [2048, 4096, 8192], 4096),
94 + (4097, [2048, 4096, 8192], 8192),
95 + # Exceeds all budgets — returns None (eager fallback)
96 + (9000, [2048, 4096, 8192], None),
97 + # Single budget, fits
98 + (1000, [2048], 2048),
99 + # Single budget, does not fit
100 + (3000, [2048], None),
101 + ],
102 + )
103 + def test_find_budget(self, total_tokens, budgets, expected):
104 + mgr = _make_manager_with_budgets(budgets)
105 + result = mgr._find_smallest_fitting_budget_given_tokens(total_tokens)
106 + assert result == expected
107 +
108 + def test_budgets_are_sorted(self):
109 + """Manager always sorts budgets ascending at init."""
110 + mgr = _make_manager_with_budgets([8192, 2048, 4096])
111 + assert mgr.token_budgets == [2048, 4096, 8192]
112 + # Budget selection still works correctly after sorting
113 + assert mgr._find_smallest_fitting_budget_given_tokens(3000) == 4096
114 +
115 +
116 +# ---------------------------------------------------------------------------
117 +# get_cumulative_stats
118 +# ---------------------------------------------------------------------------
119 +
120 +
121 +class TestGetCumulativeStats:
122 + """Statistics tracking and reporting."""
123 +
124 + def test_initial_stats_are_zero(self):
125 + mgr = _make_manager_with_budgets([2048])
126 + stats = mgr.get_cumulative_stats()
127 + assert stats["graph_hits"] == 0
128 + assert stats["graph_misses"] == 0
129 + assert stats["hit_rate"] == 0.0
130 +
131 + def test_hit_rate_calculation(self):
132 + mgr = _make_manager_with_budgets([2048])
133 + mgr.graph_hits = 75
134 + mgr.graph_misses = 25
135 + stats = mgr.get_cumulative_stats()
136 + assert stats["graph_hits"] == 75
137 + assert stats["graph_misses"] == 25
138 + assert stats["hit_rate"] == pytest.approx(0.75)
139 +
140 + def test_all_hits(self):
141 + mgr = _make_manager_with_budgets([2048])
142 + mgr.graph_hits = 100
143 + mgr.graph_misses = 0
144 + assert mgr.get_cumulative_stats()["hit_rate"] == pytest.approx(1.0)
145 +
146 + def test_all_misses(self):
147 + mgr = _make_manager_with_budgets([2048])
148 + mgr.graph_hits = 0
149 + mgr.graph_misses = 50
150 + assert mgr.get_cumulative_stats()["hit_rate"] == pytest.approx(0.0)
151 +
152 + def test_stats_report_budget_info(self):
153 + budgets = [2048, 4096, 8192]
154 + mgr = _make_manager_with_budgets(budgets)
155 + stats = mgr.get_cumulative_stats()
156 + assert stats["num_budgets"] == 0 # no graphs captured yet
157 + assert stats["token_budgets"] == budgets
158 +
159 +
160 +# ---------------------------------------------------------------------------
161 +# GPU fixtures and helpers
162 +# ---------------------------------------------------------------------------
163 +
164 +# Mock encoder parameters (kept small for fast capture)
165 +_SPATIAL_MERGE = 2
166 +_HIDDEN = 32
167 +_PATCH_SIZE = 4 # H/W per patch in grid_thw units
168 +_TEMPORAL_PATCH = 1
169 +_IN_CHANNELS = 3
170 +# flattened_patch_size = in_channels * temporal_patch * patch_size^2
171 +_FLAT = _IN_CHANNELS * _TEMPORAL_PATCH * _PATCH_SIZE * _PATCH_SIZE # 48
172 +
173 +# Test budgets: small to keep capture fast
174 +_BUDGETS = [16, 64]
175 +_MAX_BATCH = 4
176 +
177 +
178 +def _count_input_patches(grid_thw_list: list[list[int]]) -> int:
179 + return sum(t * h * w for t, h, w in grid_thw_list)
180 +
181 +
182 +def _count_output_tokens(
183 + grid_thw_list: list[list[int]], spatial_merge_size: int
184 +) -> int:
185 + m = spatial_merge_size
186 + return sum(t * (h // m) * (w // m) for t, h, w in grid_thw_list)
187 +
188 +
189 +class SimpleMockViTModel(torch.nn.Module):
190 + """Minimal ViT model for CUDA graph tests.
191 +
192 + Implements the SupportsEncoderCudaGraph protocol by providing
193 + all required methods. The forward pass projects patches and
194 + simulates spatial merge by averaging groups of m^2 patches.
195 + """
196 +
197 + supports_encoder_cudagraph = True
198 +
199 + def __init__(self):
200 + super().__init__()
201 + self.proj = torch.nn.Linear(_FLAT, _HIDDEN)
202 + self.spatial_merge_size = _SPATIAL_MERGE
203 + self.out_hidden_size = _HIDDEN
204 +
205 + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
206 + return EncoderCudaGraphConfig(
207 + modalities=["image"],
208 + input_key="pixel_values",
209 + buffer_keys=["dummy_buf"],
210 + out_hidden_size=_HIDDEN,
211 + )
212 +
213 + def get_encoder_cudagraph_budget_range(
214 + self,
215 + vllm_config,
216 + ) -> tuple[int, int]:
217 + # For tests: min=4, max=128 (small values for fast capture)
218 + return (4, 128)
219 +
220 + def get_encoder_cudagraph_num_items(
221 + self,
222 + mm_kwargs: dict[str, Any],
223 + ) -> int:
224 + return len(mm_kwargs["image_grid_thw"])
225 +
226 + def get_encoder_cudagraph_per_item_output_tokens(
227 + self,
228 + mm_kwargs: dict[str, Any],
229 + ) -> list[int]:
230 + m = _SPATIAL_MERGE
231 + return [t * (h // m) * (w // m) for t, h, w in mm_kwargs["image_grid_thw"]]
232 +
233 + def get_encoder_cudagraph_per_item_input_sizes(
234 + self,
235 + mm_kwargs: dict[str, Any],
236 + ) -> list[int]:
237 + return [t * h * w for t, h, w in mm_kwargs["image_grid_thw"]]
238 +
239 + def select_encoder_cudagraph_items(
240 + self,
241 + mm_kwargs: dict[str, Any],
242 + indices: list[int],
243 + ) -> dict[str, Any]:
244 + grid_thw = mm_kwargs["image_grid_thw"]
245 + pixel_values = mm_kwargs["pixel_values"]
246 +
247 + if len(indices) == 0:
248 + return {
249 + "pixel_values": pixel_values[:0],
250 + "image_grid_thw": [],
251 + }
252 +
253 + patches_per_item = [t * h * w for t, h, w in grid_thw]
254 + cum_patches = [0]
255 + for p in patches_per_item:
256 + cum_patches.append(cum_patches[-1] + p)
257 +
258 + selected_pv = torch.cat(
259 + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
260 + )
261 + selected_grid = [grid_thw[i] for i in indices]
262 + return {
263 + "pixel_values": selected_pv,
264 + "image_grid_thw": selected_grid,
265 + }
266 +
267 + def prepare_encoder_cudagraph_capture_inputs(
268 + self,
269 + token_budget: int,
270 + max_batch_size: int,
271 + device: torch.device,
272 + dtype: torch.dtype,
273 + ) -> EncoderCudaGraphCaptureInputs:
274 + per_image_output = token_budget // max_batch_size
275 + grid_config = [
276 + [1, _SPATIAL_MERGE, per_image_output * _SPATIAL_MERGE]
277 + for _ in range(max_batch_size)
278 + ]
279 + total_patches = _count_input_patches(grid_config)
280 + dummy_pixel_values = torch.randn(
281 + total_patches, _FLAT, device=device, dtype=dtype
282 + )
283 + n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
284 + dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
285 + return EncoderCudaGraphCaptureInputs(
286 + mm_kwargs={
287 + "pixel_values": dummy_pixel_values,
288 + "image_grid_thw": grid_config,
289 + },
290 + buffers={"dummy_buf": dummy_buf},
291 + )
292 +
293 + def prepare_encoder_cudagraph_replay_buffers(
294 + self,
295 + mm_kwargs: dict[str, Any],
296 + max_batch_size: int,
297 + ) -> EncoderCudaGraphReplayBuffers:
298 + grid_thw = mm_kwargs["image_grid_thw"]
299 + n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)
300 + p = next(self.parameters())
301 + dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
302 + return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
303 +
304 + def encoder_cudagraph_forward(
305 + self,
306 + mm_kwargs: dict[str, Any],
307 + buffers: dict[str, torch.Tensor],
308 + ) -> torch.Tensor:
309 + return self._forward(mm_kwargs["pixel_values"])
310 +
311 + def encoder_eager_forward(
312 + self,
313 + mm_kwargs: dict[str, Any],
314 + ) -> torch.Tensor:
315 + return self._forward(mm_kwargs["pixel_values"])
316 +
317 + def _forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
318 + m2 = _SPATIAL_MERGE**2
319 + out = self.proj(pixel_values)
320 + n_out = out.shape[0] // m2
321 + return out[: n_out * m2].view(n_out, m2, _HIDDEN).mean(dim=1)
322 +
323 +
324 +def _make_manager_for_gpu(
325 + model: SimpleMockViTModel,
326 + token_budgets: list[int],
327 + max_batch_size: int,
328 + device: torch.device,
329 + dtype: torch.dtype,
330 +) -> EncoderCudaGraphManager:
331 + """Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests."""
332 + mgr = object.__new__(EncoderCudaGraphManager)
333 + mgr.token_budgets = sorted(token_budgets)
334 + mgr.max_batch_size = max_batch_size
335 + mgr.use_dp = False
336 + mgr.budget_graphs = {}
337 + mgr.graph_hits = 0
338 + mgr.graph_misses = 0
339 + mgr.log_stats_interval = 100
340 + mgr.model = model
341 + mgr.config = model.get_encoder_cudagraph_config()
342 + mgr.device = device
343 + mgr.dtype = dtype
344 + return mgr
345 +
346 +
347 +def _make_pixel_values(
348 + grid_thw_list: list[list[int]],
349 + device: torch.device,
350 + dtype: torch.dtype,
351 +) -> torch.Tensor:
352 + """Random pixel_values matching the total input patch count."""
353 + n = _count_input_patches(grid_thw_list)
354 + return torch.randn(n, _FLAT, device=device, dtype=dtype)
355 +
356 +
357 +def _make_mm_kwargs(
358 + grid_thw_list: list[list[int]],
359 + device: torch.device,
360 + dtype: torch.dtype,
361 +) -> dict[str, Any]:
362 + """Create mm_kwargs for testing."""
363 + return {
364 + "pixel_values": _make_pixel_values(grid_thw_list, device, dtype),
365 + "image_grid_thw": grid_thw_list,
366 + }
367 +
368 +
369 +# ---------------------------------------------------------------------------
370 +# GPU tests — capture, replay, fallback, counters, chunking
371 +# ---------------------------------------------------------------------------
372 +
373 +
374 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
375 +class TestEncoderCudaGraphCaptureReplay:
376 + def setup_method(self):
377 + self.device = torch.device("cuda:0")
378 + self.dtype = torch.float16
379 + self.model = SimpleMockViTModel().to(self.device).half()
380 + self.mgr = _make_manager_for_gpu(
381 + self.model, _BUDGETS, _MAX_BATCH, self.device, self.dtype
382 + )
383 + self.mgr.capture()
384 +
385 + # --- capture ---
386 +
387 + def test_capture_creates_one_graph_per_budget(self):
388 + assert len(self.mgr.budget_graphs) == len(_BUDGETS)
389 + assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)
390 +
391 + # --- output shape ---
392 +
393 + def test_execute_returns_one_tensor_per_image(self):
394 + grid_thw = [[1, 4, 4], [1, 4, 4]]
395 + mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
396 + result = self.mgr.execute(mm_kwargs)
397 + assert result is not None
398 + assert len(result) == 2
399 +
400 + def test_execute_output_tokens_per_image(self):
401 + # [1,4,4] → 1*(4//2)*(4//2) = 4 tokens; [1,8,8] → 16 tokens
402 + grid_thw = [[1, 4, 4], [1, 8, 8]]
403 + mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
404 + result = self.mgr.execute(mm_kwargs)
405 + assert result is not None
406 + assert result[0].shape == (4, _HIDDEN)
407 + assert result[1].shape == (16, _HIDDEN)
408 +
409 + # --- budget fallback ---
410 +
411 + def test_eager_fallback_when_tokens_exceed_all_budgets(self):
412 + # [1,18,18] → 1*(18//2)*(18//2) = 81 tokens > max budget 64.
413 + # Greedy packing handles the fallback internally: the oversized image
414 + # gets an eager forward pass and is returned as part of the output list
415 + # (execute() no longer returns None for individual image misses).
416 + grid_thw = [[1, 18, 18]]
417 + mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
418 + result = self.mgr.execute(mm_kwargs)
419 + assert result is not None
420 + assert len(result) == 1
421 + # Eager output: SimpleMockViTModel produces n_out = 81 tokens
422 + assert result[0].shape == (81, _HIDDEN)
423 + assert self.mgr.graph_misses == 1
424 +
425 + # --- counters ---
426 +
427 + def test_hit_counter_increments_by_num_images(self):
428 + grid_thw = [[1, 4, 4], [1, 4, 4]]
429 + mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
430 + self.mgr.execute(mm_kwargs)
431 + assert self.mgr.graph_hits == 2
432 +
433 + def test_miss_counter_increments_by_num_images(self):
434 + grid_thw = [[1, 18, 18]] # 81 tokens > 64
435 + mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
436 + self.mgr.execute(mm_kwargs)
437 + assert self.mgr.graph_misses == 1
438 +
439 + # --- chunking ---
440 +
441 + def test_chunking_when_images_exceed_max_batch(self):
442 + # 8 images > max_batch_size=4 → 2 chunks of 4
443 + # each chunk: 4 * 4 = 16 tokens → fits budget 16
444 + n_images = _MAX_BATCH * 2
445 + grid_thw = [[1, 4, 4]] * n_images
446 + mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
447 + result = self.mgr.execute(mm_kwargs)
448 + assert result is not None
449 + assert len(result) == n_images
450 + for out in result:
451 + assert out.shape == (4, _HIDDEN)
vllm/config/compilation.py
+32
@@ -489,6 +489,28 @@ class CompilationConfig:
489 on selected platforms. Disabled by default until more models
490 are supported/tested to work."""
491
492 + # Vision encoder CUDA graph
493 + cudagraph_mm_encoder: bool = False
494 + """Enable CUDA graph capture for multimodal encoder (ViT).
495 + When enabled, captures full encoder forward as CUDA graph
496 + for each token budget level."""
497 +
498 + encoder_cudagraph_token_budgets: list[int] = field(default_factory=list)
499 + """Token budget levels for encoder CUDA graph capture.
500 + Each budget defines a fixed token capacity. At runtime, images are greedy-packed
501 + into the smallest fitting budget and the corresponding CUDA graph is replayed.
502 + If empty (default), auto-inferred from model architecture as power-of-2
503 + levels from the model's estimated min budget to max budget.
504 + User-provided values override auto-inference.
505 + Example: [2048, 4096, 8192, 13824]"""
506 +
507 + encoder_cudagraph_max_images_per_batch: int = 0
508 + """Maximum number of images per batch for encoder CUDA graph capture.
509 + Determines the fixed batch size used during graph capture.
510 + If 0 (default), auto-inferred as max_budget // min_budget from the
511 + model's budget range. User-provided positive value overrides
512 + auto-inference."""
513 +
514 # Inductor capture
515 compile_sizes: list[int | str] | None = None
516 """Sizes to compile for inductor. In addition
@@ -906,6 +928,16 @@ class CompilationConfig:
928 f"Invalid backend for piecewise compilation: {self.backend}"
929 )
930
931 + # Validate encoder CUDA graph configuration
932 + if (
933 + self.cudagraph_mm_encoder
934 + and self.encoder_cudagraph_max_images_per_batch < 0
935 + ):
936 + raise ValueError(
937 + "encoder_cudagraph_max_images_per_batch must be "
938 + "non-negative (0 = auto-infer)"
939 + )
940 +
941 if self.backend == "":
942 self.backend = current_platform.get_compile_backend()
943
vllm/model_executor/models/interfaces.py
+141
@@ -13,6 +13,7 @@ from collections.abc import (
13 from contextlib import ExitStack, contextmanager, nullcontext
14 from typing import (
15 TYPE_CHECKING,
16 + Any,
17 ClassVar,
18 Literal,
19 Protocol,
@@ -46,6 +47,11 @@ if TYPE_CHECKING:
47 from vllm.multimodal.inputs import MultiModalFeatureSpec
48 from vllm.multimodal.registry import _ProcessorFactories
49 from vllm.sequence import IntermediateTensors
50 + from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
51 + EncoderCudaGraphCaptureInputs,
52 + EncoderCudaGraphConfig,
53 + EncoderCudaGraphReplayBuffers,
54 + )
55 else:
56 VllmConfig = object
57 WeightsMapper = object
@@ -1494,3 +1500,138 @@ def supports_xdrope(
1500 model: type[object] | object,
1501 ) -> TypeIs[type[SupportsXDRoPE]] | TypeIs[SupportsXDRoPE]:
1502 return isinstance(model, SupportsXDRoPE)
1503 +
1504 +
1505 +@runtime_checkable
1506 +class SupportsEncoderCudaGraph(Protocol):
1507 + """Interface for models whose vision encoder supports CUDA graph
1508 + capture/replay.
1509 +
1510 + Models implement these methods to provide the
1511 + :class:`EncoderCudaGraphManager` with all model-specific logic
1512 + (input handling, metadata computation, forward pass) without the
1513 + manager needing to know model internals.
1514 + """
1515 +
1516 + supports_encoder_cudagraph: ClassVar[Literal[True]] = True
1517 +
1518 + def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": ...
1519 +
1520 + def get_encoder_cudagraph_budget_range(
1521 + self,
1522 + vllm_config: "VllmConfig",
1523 + ) -> tuple[int, int]:
1524 + """Return (min_token_budget, max_token_budget) for auto-inference.
1525 +
1526 + - min_token_budget: estimated smallest possible encoder input
1527 + (e.g. 64 for a 224x224 image)
1528 + - max_token_budget: estimated largest budget worth capturing
1529 + (e.g. max_num_batched_tokens)
1530 +
1531 + Used when ``encoder_cudagraph_token_budgets`` and/or
1532 + ``encoder_cudagraph_max_images_per_batch`` are not explicitly
1533 + specified by the user.
1534 + """
1535 + ...
1536 +
1537 + def get_encoder_cudagraph_num_items(
1538 + self,
1539 + mm_kwargs: dict[str, Any],
1540 + ) -> int:
1541 + """Return the number of items (e.g. images) in the batch."""
1542 + ...
1543 +
1544 + def get_encoder_cudagraph_per_item_output_tokens(
1545 + self,
1546 + mm_kwargs: dict[str, Any],
1547 + ) -> list[int]:
1548 + """Return output token count for each item.
1549 +
1550 + Used for greedy packing and DP load balancing.
1551 + """
1552 + ...
1553 +
1554 + def get_encoder_cudagraph_per_item_input_sizes(
1555 + self,
1556 + mm_kwargs: dict[str, Any],
1557 + ) -> list[int]:
1558 + """Return input size (e.g. patch count) for each item.
1559 +
1560 + Used for input tensor slicing offsets.
1561 + """
1562 + ...
1563 +
1564 + def select_encoder_cudagraph_items(
1565 + self,
1566 + mm_kwargs: dict[str, Any],
1567 + indices: list[int],
1568 + ) -> dict[str, Any]:
1569 + """Select a subset of items and return mm_kwargs for the sub-batch.
1570 +
1571 + Called by the manager during greedy packing and DP sharding to
1572 + extract inputs for a specific set of items (e.g. images at
1573 + indices [0, 3, 5]). The implementation is model-specific
1574 + because input formats differ:
1575 +
1576 + - Qwen-family: slice concatenated pixel_values by cumulative
1577 + patch offsets, subset grid_thw by indices.
1578 + - Batched models (CLIP): index pixel_values along dim 0.
1579 + """
1580 + ...
1581 +
1582 + def prepare_encoder_cudagraph_capture_inputs(
1583 + self,
1584 + token_budget: int,
1585 + max_batch_size: int,
1586 + device: torch.device,
1587 + dtype: torch.dtype,
1588 + ) -> "EncoderCudaGraphCaptureInputs":
1589 + """Create dummy inputs and buffers for CUDA graph capture."""
1590 + ...
1591 +
1592 + def prepare_encoder_cudagraph_replay_buffers(
1593 + self,
1594 + mm_kwargs: dict[str, Any],
1595 + max_batch_size: int,
1596 + ) -> "EncoderCudaGraphReplayBuffers":
1597 + """Compute buffer values from actual batch inputs for replay."""
1598 + ...
1599 +
1600 + def encoder_cudagraph_forward(
1601 + self,
1602 + mm_kwargs: dict[str, Any],
1603 + buffers: dict[str, torch.Tensor],
1604 + ) -> torch.Tensor:
1605 + """Run the encoder forward pass with precomputed buffers.
1606 +
1607 + Used during both CUDA graph capture and replay.
1608 + """
1609 + ...
1610 +
1611 + def encoder_eager_forward(
1612 + self,
1613 + mm_kwargs: dict[str, Any],
1614 + ) -> torch.Tensor:
1615 + """Run the encoder forward pass without precomputed buffers.
1616 +
1617 + Used as eager fallback when inputs exceed all budgets.
1618 + """
1619 + ...
1620 +
1621 +
1622 +@overload
1623 +def supports_encoder_cudagraph(
1624 + model: type[object],
1625 +) -> TypeIs[type[SupportsEncoderCudaGraph]]: ...
1626 +
1627 +
1628 +@overload
1629 +def supports_encoder_cudagraph(
1630 + model: object,
1631 +) -> TypeIs[SupportsEncoderCudaGraph]: ...
1632 +
1633 +
1634 +def supports_encoder_cudagraph(
1635 + model: type[object] | object,
1636 +) -> TypeIs[type[SupportsEncoderCudaGraph]] | TypeIs[SupportsEncoderCudaGraph]:
1637 + return isinstance(model, SupportsEncoderCudaGraph)
vllm/model_executor/models/qwen3_vl.py
+270 -30
@@ -103,6 +103,7 @@ from .interfaces import (
103 MultiModalEmbeddings,
104 SupportsEagle,
105 SupportsEagle3,
106 + SupportsEncoderCudaGraph,
107 SupportsLoRA,
108 SupportsMRoPE,
109 SupportsMultiModal,
@@ -528,54 +529,120 @@ class Qwen3_VisionTransformer(nn.Module):
529
530 return torch.cat(outputs, dim=0)
531
531 - def forward(
532 + def prepare_encoder_metadata(
533 self,
533 - x: torch.Tensor,
534 - grid_thw: torch.Tensor | list[list[int]],
535 - ) -> torch.Tensor:
536 - hidden_states = x.to(device=self.device, dtype=self.dtype, non_blocking=True)
537 - hidden_states = self.patch_embed(hidden_states)
534 + grid_thw_list: list[list[int]],
535 + *,
536 + max_batch_size: int | None = None,
537 + max_seqlen_override: int | None = None,
538 + device: torch.device | None = None,
539 + ) -> dict[str, torch.Tensor | None]:
540 + """Compute encoder metadata from grid_thw_list.
541
539 - if isinstance(grid_thw, list):
540 - grid_thw_list = grid_thw
541 - grid_thw = np.array(grid_thw, dtype=np.int32)
542 - else:
543 - grid_thw_list = grid_thw.tolist()
544 - grid_thw = grid_thw.numpy()
542 + Shared by the eager forward path, CUDA graph capture, and
543 + CUDA graph replay to avoid duplicated implementation.
544
546 - pos_embeds = self.fast_pos_embed_interpolate(grid_thw_list)
547 - hidden_states = hidden_states + pos_embeds
548 - rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
545 + Args:
546 + grid_thw_list: Grid configurations as list of [t, h, w].
547 + max_batch_size: If set, pad cu_seqlens to this size
548 + (needed for CUDA graph capture/replay).
549 + max_seqlen_override: If set, use this value for max_seqlen
550 + instead of computing from cu_seqlens (needed for CUDA
551 + graph capture to cover worst-case replay scenarios).
552 + device: Device to place tensors on. Defaults to self.device.
553 + """
554 + if device is None:
555 + device = self.device
556
550 - cu_seqlens = np.repeat(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
551 - axis=0, dtype=np.int32
557 + metadata: dict[str, torch.Tensor | None] = {}
558 +
559 + # Positional embeddings
560 + metadata["pos_embeds"] = self.fast_pos_embed_interpolate(grid_thw_list)
561 + rotary_cos, rotary_sin = self.rot_pos_emb(grid_thw_list)
562 + metadata["rotary_pos_emb_cos"] = rotary_cos
563 + metadata["rotary_pos_emb_sin"] = rotary_sin
564 +
565 + # cu_seqlens from grid_thw
566 + grid_thw_np = np.array(grid_thw_list, dtype=np.int32)
567 + patches_per_frame = grid_thw_np[:, 1] * grid_thw_np[:, 2]
568 + cu_seqlens = np.repeat(patches_per_frame, grid_thw_np[:, 0]).cumsum(
569 + dtype=np.int32
570 )
571 cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
554 - sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
555 - self.attn_backend, cu_seqlens, self.device
556 - )
557 - max_seqlen = torch.tensor(
558 - MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens),
559 - dtype=torch.int32,
572 +
573 + # Pad cu_seqlens if max_batch_size specified
574 + if max_batch_size is not None:
575 + num_seqs = len(cu_seqlens) - 1
576 + if num_seqs < max_batch_size:
577 + cu_seqlens = np.concatenate(
578 + [
579 + cu_seqlens,
580 + np.full(
581 + max_batch_size - num_seqs,
582 + cu_seqlens[-1],
583 + dtype=np.int32,
584 + ),
585 + ]
586 + )
587 +
588 + # sequence_lengths (backend-specific)
589 + metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(
590 + self.attn_backend, cu_seqlens, device
591 )
561 - cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
592 +
593 + # max_seqlen
594 + if max_seqlen_override is not None:
595 + max_seqlen_val = max_seqlen_override
596 + else:
597 + max_seqlen_val = MMEncoderAttention.compute_max_seqlen(
598 + self.attn_backend, cu_seqlens
599 + )
600 + # Keep max_seqlen on CPU: attention wrappers call .item() on it,
601 + # and having it on GPU would capture a wasteful D2H copy in CUDA
602 + # graphs without changing behavior (the scalar is baked at capture).
603 + metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)
604 +
605 + # Recompute cu_seqlens (backend-specific transformation)
606 + metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(
607 self.attn_backend,
608 cu_seqlens,
609 self.hidden_size,
610 self.tp_size,
566 - self.device,
611 + device,
612 )
613 +
614 + return metadata
615 +
616 + def forward(
617 + self,
618 + x: torch.Tensor,
619 + grid_thw: torch.Tensor | list[list[int]],
620 + *,
621 + encoder_metadata: dict[str, torch.Tensor] | None = None,
622 + ) -> torch.Tensor:
623 + hidden_states = x.to(device=self.device, dtype=self.dtype, non_blocking=True)
624 + hidden_states = self.patch_embed(hidden_states)
625 +
626 + if encoder_metadata is None:
627 + if isinstance(grid_thw, list):
628 + grid_thw_list = grid_thw
629 + else:
630 + grid_thw_list = grid_thw.tolist()
631 + encoder_metadata = self.prepare_encoder_metadata(grid_thw_list)
632 +
633 + pos_embeds = encoder_metadata["pos_embeds"]
634 + hidden_states = hidden_states + pos_embeds
635 hidden_states = hidden_states.unsqueeze(1)
636
637 deepstack_feature_lists = []
638 for layer_num, blk in enumerate(self.blocks):
639 hidden_states = blk(
640 hidden_states,
574 - cu_seqlens=cu_seqlens,
575 - rotary_pos_emb_cos=rotary_pos_emb_cos,
576 - rotary_pos_emb_sin=rotary_pos_emb_sin,
577 - max_seqlen=max_seqlen,
578 - sequence_lengths=sequence_lengths,
641 + cu_seqlens=encoder_metadata["cu_seqlens"],
642 + rotary_pos_emb_cos=encoder_metadata["rotary_pos_emb_cos"],
643 + rotary_pos_emb_sin=encoder_metadata["rotary_pos_emb_sin"],
644 + max_seqlen=encoder_metadata["max_seqlen"],
645 + sequence_lengths=encoder_metadata.get("sequence_lengths"),
646 )
647 if layer_num in self.deepstack_visual_indexes:
648 deepstack_merger_idx = self.deepstack_visual_indexes.index(layer_num)
@@ -1358,6 +1425,7 @@ class Qwen3LLMForCausalLM(Qwen3ForCausalLM):
1425 class Qwen3VLForConditionalGeneration(
1426 nn.Module,
1427 SupportsMultiModal,
1428 + SupportsEncoderCudaGraph,
1429 SupportsLoRA,
1430 SupportsPP,
1431 SupportsMRoPE,
@@ -1507,6 +1575,178 @@ class Qwen3VLForConditionalGeneration(
1575 for idx in range(self.deepstack_num_level):
1576 self.deepstack_input_embeds[idx][:num_tokens].zero_()
1577
1578 + # -- SupportsEncoderCudaGraph protocol methods --
1579 +
1580 + def get_encoder_cudagraph_config(self):
1581 + from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
1582 + EncoderCudaGraphConfig,
1583 + )
1584 +
1585 + return EncoderCudaGraphConfig(
1586 + modalities=["image"],
1587 + input_key="pixel_values",
1588 + buffer_keys=[
1589 + "pos_embeds",
1590 + "rotary_pos_emb_cos",
1591 + "rotary_pos_emb_sin",
1592 + "cu_seqlens",
1593 + "max_seqlen",
1594 + "sequence_lengths",
1595 + ],
1596 + out_hidden_size=self.visual.out_hidden_size,
1597 + )
1598 +
1599 + def get_encoder_cudagraph_budget_range(
1600 + self,
1601 + vllm_config,
1602 + ) -> tuple[int, int]:
1603 + # Min: estimated smallest possible encoder input.
1604 + # 224x224 image → 16x16 patches, spatial_merge_size=2 → 8x8 = 64 tokens
1605 + min_budget = 64
1606 + # Max: capped by max_num_batched_tokens
1607 + max_budget = vllm_config.scheduler_config.max_num_batched_tokens
1608 + return (min_budget, max_budget)
1609 +
1610 + def get_encoder_cudagraph_num_items(
1611 + self,
1612 + mm_kwargs: dict[str, Any],
1613 + ) -> int:
1614 + return len(mm_kwargs["image_grid_thw"])
1615 +
1616 + def get_encoder_cudagraph_per_item_output_tokens(
1617 + self,
1618 + mm_kwargs: dict[str, Any],
1619 + ) -> list[int]:
1620 + m = self.visual.spatial_merge_size
1621 + return [t * (h // m) * (w // m) for t, h, w in mm_kwargs["image_grid_thw"]]
1622 +
1623 + def get_encoder_cudagraph_per_item_input_sizes(
1624 + self,
1625 + mm_kwargs: dict[str, Any],
1626 + ) -> list[int]:
1627 + return [t * h * w for t, h, w in mm_kwargs["image_grid_thw"]]
1628 +
1629 + def select_encoder_cudagraph_items(
1630 + self,
1631 + mm_kwargs: dict[str, Any],
1632 + indices: list[int],
1633 + ) -> dict[str, Any]:
1634 + grid_thw = mm_kwargs["image_grid_thw"]
1635 + pixel_values = mm_kwargs["pixel_values"]
1636 +
1637 + if len(indices) == 0:
1638 + return {
1639 + "pixel_values": pixel_values[:0],
1640 + "image_grid_thw": [],
1641 + }
1642 +
1643 + # Compute cumulative patch offsets for slicing pixel_values
1644 + patches_per_item = [t * h * w for t, h, w in grid_thw]
1645 + cum_patches = [0]
1646 + for p in patches_per_item:
1647 + cum_patches.append(cum_patches[-1] + p)
1648 +
1649 + selected_pv = torch.cat(
1650 + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
1651 + )
1652 + selected_grid = [grid_thw[i] for i in indices]
1653 +
1654 + return {
1655 + "pixel_values": selected_pv,
1656 + "image_grid_thw": selected_grid,
1657 + }
1658 +
1659 + def prepare_encoder_cudagraph_capture_inputs(
1660 + self,
1661 + token_budget: int,
1662 + max_batch_size: int,
1663 + device: torch.device,
1664 + dtype: torch.dtype,
1665 + ):
1666 + from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
1667 + EncoderCudaGraphCaptureInputs,
1668 + )
1669 +
1670 + spatial_merge_size = self.visual.spatial_merge_size
1671 + per_image_output = token_budget // max_batch_size
1672 +
1673 + # Synthetic rectangular grid: [1, merge, per_image_output * merge]
1674 + # produces exactly per_image_output tokens per image.
1675 + grid_config = [
1676 + [1, spatial_merge_size, per_image_output * spatial_merge_size]
1677 + for _ in range(max_batch_size)
1678 + ]
1679 +
1680 + # Create dummy pixel_values
1681 + patch_embed = self.visual.patch_embed
1682 + in_channels = patch_embed.proj.in_channels
1683 + patch_size = patch_embed.patch_size
1684 + temporal_patch_size = patch_embed.temporal_patch_size
1685 + total_patches = sum(t * h * w for t, h, w in grid_config)
1686 + flattened_patch_size = (
1687 + in_channels * temporal_patch_size * patch_size * patch_size
1688 + )
1689 + dummy_pixel_values = torch.randn(
1690 + total_patches, flattened_patch_size, device=device, dtype=dtype
1691 + )
1692 +
1693 + # Override max_seqlen with a safe upper bound for capture.
1694 + # max_seqlen.item() gets baked into the CUDA graph (not replayed),
1695 + # so the capture value must cover any replay scenario.
1696 + # Worst case: 1 image consuming the full budget ->
1697 + # seq_len = token_budget * spatial_merge_size^2.
1698 + buffers = self.visual.prepare_encoder_metadata(
1699 + grid_config,
1700 + max_batch_size=max_batch_size,
1701 + max_seqlen_override=token_budget * (spatial_merge_size**2),
1702 + device=device,
1703 + )
1704 +
1705 + mm_kwargs = {
1706 + "pixel_values": dummy_pixel_values,
1707 + "image_grid_thw": grid_config,
1708 + }
1709 +
1710 + return EncoderCudaGraphCaptureInputs(
1711 + mm_kwargs=mm_kwargs,
1712 + buffers=buffers,
1713 + )
1714 +
1715 + def prepare_encoder_cudagraph_replay_buffers(
1716 + self,
1717 + mm_kwargs: dict[str, Any],
1718 + max_batch_size: int,
1719 + ):
1720 + from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
1721 + EncoderCudaGraphReplayBuffers,
1722 + )
1723 +
1724 + grid_thw_list = mm_kwargs["image_grid_thw"]
1725 +
1726 + buffers = self.visual.prepare_encoder_metadata(
1727 + grid_thw_list,
1728 + max_batch_size=max_batch_size,
1729 + )
1730 +
1731 + return EncoderCudaGraphReplayBuffers(buffers=buffers)
1732 +
1733 + def encoder_cudagraph_forward(
1734 + self,
1735 + mm_kwargs: dict[str, Any],
1736 + buffers: dict[str, torch.Tensor],
1737 + ) -> torch.Tensor:
1738 + pixel_values = mm_kwargs["pixel_values"]
1739 + grid_thw = mm_kwargs["image_grid_thw"]
1740 + return self.visual(pixel_values, grid_thw, encoder_metadata=buffers)
1741 +
1742 + def encoder_eager_forward(
1743 + self,
1744 + mm_kwargs: dict[str, Any],
1745 + ) -> torch.Tensor:
1746 + pixel_values = mm_kwargs["pixel_values"]
1747 + grid_thw = mm_kwargs["image_grid_thw"]
1748 + return self.visual(pixel_values, grid_thw)
1749 +
1750 def _parse_and_validate_image_input(
1751 self, **kwargs: object
1752 ) -> Qwen2_5_VLImageInputs | None:
vllm/v1/worker/gpu/mm/encoder_cudagraph.py new
+576
@@ -0,0 +1,576 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CUDA graph manager for vision encoder budget-batch execution."""
4 +
5 +from dataclasses import dataclass
6 +from typing import Any
7 +
8 +import torch
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.distributed import (
12 + get_tensor_model_parallel_rank,
13 + get_tensor_model_parallel_world_size,
14 + tensor_model_parallel_all_gather,
15 +)
16 +from vllm.logger import init_logger
17 +from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
18 +from vllm.model_executor.models.vision import get_load_balance_assignment
19 +from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
20 + EncoderCudaGraphConfig,
21 +)
22 +
23 +logger = init_logger(__name__)
24 +
25 +
26 +@dataclass
27 +class BudgetGraphMetadata:
28 + """Metadata for a single budget graph.
29 +
30 + CUDA graph replay pattern:
31 + 1. Copy new batch data into input_buffer (e.g. pixel_values)
32 + 2. Copy precomputed values into metadata_buffers
33 + 3. Replay graph
34 + 4. Read encoder outputs from output_buffer
35 + """
36 +
37 + token_budget: int
38 + max_batch_size: int # Max number of images/videos per batch
39 + graph: torch.cuda.CUDAGraph
40 + # The input tensor updated before replay (e.g. pixel_values)
41 + input_buffer: torch.Tensor
42 + # Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata).
43 + # Before replay the manager zeros then slice-copies new data into these.
44 + metadata_buffers: dict[str, torch.Tensor]
45 + # Output written by graph, read after replay
46 + output_buffer: torch.Tensor
47 +
48 +
49 +class EncoderCudaGraphManager:
50 + """Budget-based CUDA graph capture/replay for vision encoders."""
51 +
52 + def __init__(
53 + self,
54 + vllm_config: VllmConfig,
55 + device: torch.device,
56 + dtype: torch.dtype,
57 + model: SupportsEncoderCudaGraph,
58 + ):
59 + """Initialize CUDA graph manager with provided token budgets
60 + and max batch size."""
61 + self.vllm_config = vllm_config
62 + self.device = device
63 + self.dtype = dtype
64 + self.model = model
65 + self.config: EncoderCudaGraphConfig = model.get_encoder_cudagraph_config()
66 +
67 + comp_config = vllm_config.compilation_config
68 + user_budgets = comp_config.encoder_cudagraph_token_budgets
69 + user_max_images = comp_config.encoder_cudagraph_max_images_per_batch
70 +
71 + if user_budgets and user_max_images > 0:
72 + # Fully user-specified
73 + self.token_budgets = sorted(user_budgets)
74 + self.max_batch_size = user_max_images
75 + else:
76 + # Auto-infer missing values from model
77 + min_budget, max_budget = model.get_encoder_cudagraph_budget_range(
78 + vllm_config
79 + )
80 + self.token_budgets = (
81 + sorted(user_budgets)
82 + if user_budgets
83 + else self._generate_budgets(min_budget, max_budget)
84 + )
85 + self.max_batch_size = (
86 + user_max_images if user_max_images > 0 else max_budget // min_budget
87 + )
88 +
89 + mm_config = vllm_config.model_config.multimodal_config
90 + self.use_dp = (
91 + mm_config is not None
92 + and mm_config.mm_encoder_tp_mode == "data"
93 + and vllm_config.parallel_config.tensor_parallel_size > 1
94 + )
95 +
96 + self.budget_graphs: dict[int, BudgetGraphMetadata] = {}
97 + self.graph_hits = 0
98 + self.graph_misses = 0
99 + self.log_stats_interval = 100
100 +
101 + logger.info(
102 + "EncoderCudaGraphManager initialized with "
103 + "budgets=%s, max_batch_size=%d, use_dp=%s",
104 + self.token_budgets,
105 + self.max_batch_size,
106 + self.use_dp,
107 + )
108 +
109 + @staticmethod
110 + def _generate_budgets(min_budget: int, max_budget: int) -> list[int]:
111 + """Generate power-of-2 token budgets from min_budget to max_budget."""
112 + budgets: list[int] = []
113 + b = min_budget
114 + while b <= max_budget:
115 + budgets.append(b)
116 + b *= 2
117 + # Always include max_budget if it's not already a power-of-2 boundary
118 + if not budgets or budgets[-1] < max_budget:
119 + budgets.append(max_budget)
120 + return budgets
121 +
122 + def supports_modality(self, modality: str) -> bool:
123 + """Check if a modality is supported by this manager."""
124 + return modality in self.config.modalities
125 +
126 + def capture(self):
127 + """Capture CUDA graphs for all token budgets."""
128 + for token_budget in self.token_budgets:
129 + self._capture_budget_graph(token_budget)
130 +
131 + logger.info(
132 + "Encoder CUDA graph capture complete. Captured %d budget graphs.",
133 + len(self.budget_graphs),
134 + )
135 +
136 + def _capture_budget_graph(self, token_budget: int):
137 + """Capture CUDA graph for a single token budget."""
138 + logger.debug(
139 + "Capturing encoder cudagraph for budget=%d, max_batch_size=%d",
140 + token_budget,
141 + self.max_batch_size,
142 + )
143 +
144 + capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs(
145 + token_budget, self.max_batch_size, self.device, self.dtype
146 + )
147 +
148 + mm_kwargs = capture_inputs.mm_kwargs
149 + buffers = capture_inputs.buffers
150 +
151 + with torch.inference_mode():
152 + output = self.model.encoder_cudagraph_forward(mm_kwargs, buffers)
153 + output_buffer = torch.empty_like(output)
154 +
155 + graph = torch.cuda.CUDAGraph()
156 + with torch.inference_mode(), torch.cuda.graph(graph):
157 + output = self.model.encoder_cudagraph_forward(mm_kwargs, buffers)
158 + output_buffer.copy_(output)
159 +
160 + input_key = self.config.input_key
161 + self.budget_graphs[token_budget] = BudgetGraphMetadata(
162 + token_budget=token_budget,
163 + max_batch_size=self.max_batch_size,
164 + graph=graph,
165 + input_buffer=mm_kwargs[input_key],
166 + metadata_buffers=buffers,
167 + output_buffer=output_buffer,
168 + )
169 +
170 + def _find_smallest_fitting_budget_given_tokens(
171 + self, total_tokens: int
172 + ) -> int | None:
173 + """Find smallest budget >= total_tokens.
174 +
175 + Returns:
176 + Token budget if found, None if no fitting budget.
177 + """
178 + for budget in self.token_budgets:
179 + if budget >= total_tokens:
180 + return budget
181 + return None
182 +
183 + def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]:
184 + """Get per-item output token counts as plain ints."""
185 + return [
186 + int(t)
187 + for t in self.model.get_encoder_cudagraph_per_item_output_tokens(mm_kwargs)
188 + ]
189 +
190 + @staticmethod
191 + def _scatter_output_slices(
192 + output: torch.Tensor,
193 + indices: list[int],
194 + per_item_out_tokens: list[int],
195 + dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
196 + clone: bool = False,
197 + ) -> None:
198 + """Slice a concatenated output tensor and scatter into dest by index."""
199 + offset = 0
200 + for idx in indices:
201 + n_tok = per_item_out_tokens[idx]
202 + sliced = output[offset : offset + n_tok]
203 + dest[idx] = sliced.clone() if clone else sliced
204 + offset += n_tok
205 +
206 + def _run_budget_graph(
207 + self,
208 + mm_kwargs: dict[str, Any],
209 + token_budget: int,
210 + replay_buffers: dict[str, torch.Tensor | None],
211 + ) -> torch.Tensor | None:
212 + """Execute budget graph.
213 +
214 + Args:
215 + mm_kwargs: Multimodal inputs for the batch.
216 + token_budget: Token budget to use.
217 + replay_buffers: Buffer values to copy into captured buffers.
218 + None values leave the corresponding buffer unchanged.
219 +
220 + Returns:
221 + Encoder outputs, or None if graph not captured.
222 + """
223 + num_items = self.model.get_encoder_cudagraph_num_items(mm_kwargs)
224 + if token_budget not in self.budget_graphs:
225 + self.graph_misses += num_items
226 + return None
227 +
228 + graph_meta = self.budget_graphs[token_budget]
229 +
230 + # Copy the input tensor. Buffers are sized for the full budget;
231 + # actual inputs may be smaller. Zero then slice-copy so padded
232 + # positions are invisible to attention (cu_seqlens masks them out).
233 + input_key = self.config.input_key
234 + src = mm_kwargs[input_key]
235 + n = src.shape[0]
236 + graph_meta.input_buffer.zero_()
237 + graph_meta.input_buffer[:n].copy_(src)
238 +
239 + # Copy metadata buffers using keys from config.buffer_keys.
240 + for key in self.config.buffer_keys:
241 + src = replay_buffers.get(key)
242 + if src is None:
243 + continue
244 + buf = graph_meta.metadata_buffers[key]
245 + if src.ndim == 0:
246 + buf.copy_(src)
247 + else:
248 + n = src.shape[0]
249 + buf.zero_()
250 + buf[:n].copy_(src)
251 +
252 + graph_meta.graph.replay()
253 +
254 + self.graph_hits += num_items
255 + return graph_meta.output_buffer
256 +
257 + def _execute_local(
258 + self,
259 + mm_kwargs: dict[str, Any],
260 + ) -> list[torch.Tensor]:
261 + """Execute encoder on local inputs using greedy-packed CUDA graphs.
262 +
263 + Sort images by output token count (smallest first), then greedily pack
264 + as many images as possible into each batch while staying within
265 + max_budget tokens and max_batch_size. Once a batch is finalised (next
266 + image would overflow either constraint), find the smallest fitting
267 + budget once for that batch.
268 +
269 + By exchange argument, greedy smallest-first packing minimises eager
270 + fallbacks -- any other ordering yields a higher token sum in some batch,
271 + making that batch more likely to exceed the budget.
272 +
273 + Stats note:
274 + graph_hits -- counted inside _run_budget_graph after successful replay.
275 + graph_misses -- counted here for single-image batches where the image
276 + exceeds max_budget. Batches split due to max_batch_size
277 + always satisfy total_tokens <= max_budget and therefore
278 + always find a valid budget (no miss).
279 + """
280 + num_items = self.model.get_encoder_cudagraph_num_items(mm_kwargs)
281 + max_budget = self.token_budgets[-1]
282 +
283 + per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs)
284 +
285 + # Sort ascending by output token count (smallest first)
286 + sorted_indices = sorted(range(num_items), key=lambda i: per_item_out_tokens[i])
287 +
288 + # Greedy pack against max_budget and max_batch_size.
289 + # _find_smallest_fitting_budget_given_tokens is called once per
290 + # finalised batch, not per image.
291 + batches: list[tuple[list[int], int | None]] = []
292 + current_batch: list[int] = []
293 + current_batch_tokens = 0
294 +
295 + for orig_idx in sorted_indices:
296 + item_tokens = per_item_out_tokens[orig_idx]
297 + if (
298 + current_batch_tokens + item_tokens <= max_budget
299 + and len(current_batch) < self.max_batch_size
300 + ):
301 + current_batch.append(orig_idx)
302 + current_batch_tokens += item_tokens
303 + else:
304 + if current_batch:
305 + batches.append(
306 + (
307 + current_batch,
308 + self._find_smallest_fitting_budget_given_tokens(
309 + current_batch_tokens
310 + ),
311 + )
312 + )
313 + current_batch = [orig_idx]
314 + current_batch_tokens = item_tokens
315 +
316 + if current_batch:
317 + batches.append(
318 + (
319 + current_batch,
320 + self._find_smallest_fitting_budget_given_tokens(
321 + current_batch_tokens
322 + ),
323 + )
324 + )
325 +
326 + # outputs_by_orig_idx maps each original image index to its output
327 + # tensor. Needed because greedy packing reorders images; we restore
328 + # the original order before returning.
329 + outputs_by_orig_idx: dict[int, torch.Tensor] = {}
330 +
331 + for batch_orig_indices, token_budget in batches:
332 + batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
333 + mm_kwargs, batch_orig_indices
334 + )
335 + batch_out_tokens = sum(per_item_out_tokens[i] for i in batch_orig_indices)
336 +
337 + if token_budget is None:
338 + # Single oversized image: item_tokens > max_budget.
339 + # graph_misses counted here for this eager fallback.
340 + logger.debug(
341 + "Encoder CUDA graph fallback to eager: no budget for "
342 + "%d tokens from %d images",
343 + batch_out_tokens,
344 + len(batch_orig_indices),
345 + )
346 + self.graph_misses += len(batch_orig_indices)
347 + with torch.inference_mode():
348 + raw = self.model.encoder_eager_forward(batch_mm_kwargs)
349 + self._scatter_output_slices(
350 + raw,
351 + batch_orig_indices,
352 + per_item_out_tokens,
353 + outputs_by_orig_idx,
354 + )
355 + else:
356 + logger.debug(
357 + "Encoder CUDA graph: batch_size=%d, tokens=%d, "
358 + "budget=%d, waste=%.1f%%",
359 + len(batch_orig_indices),
360 + batch_out_tokens,
361 + token_budget,
362 + (token_budget - batch_out_tokens) / token_budget * 100,
363 + )
364 + replay = self.model.prepare_encoder_cudagraph_replay_buffers(
365 + batch_mm_kwargs, self.max_batch_size
366 + )
367 +
368 + # graph_hits counted inside _run_budget_graph after replay.
369 + output = self._run_budget_graph(
370 + batch_mm_kwargs, token_budget, replay.buffers
371 + )
372 + assert output is not None
373 + self._scatter_output_slices(
374 + output,
375 + batch_orig_indices,
376 + per_item_out_tokens,
377 + outputs_by_orig_idx,
378 + clone=True,
379 + )
380 +
381 + # Return in original batch order (caller maps outputs to token positions)
382 + return [outputs_by_orig_idx[i] for i in range(num_items)]
383 +
384 + def _dp_shard(
385 + self,
386 + mm_kwargs: dict[str, Any],
387 + per_item_out_tokens: list[int],
388 + ) -> tuple[dict[str, Any], list[int], list[int], int]:
389 + """Distribute items across TP ranks for data-parallel execution.
390 +
391 + Uses get_load_balance_assignment() to balance load by input size,
392 + then select_encoder_cudagraph_items() to extract each rank's inputs.
393 +
394 + Returns:
395 + local_mm_kwargs: Inputs for this rank.
396 + image_rank_assignment: Flattened assignment order across all ranks.
397 + images_per_rank: Number of items per rank.
398 + max_output_tokens_per_rank: Max output tokens across all ranks
399 + (for padding during all_gather).
400 + """
401 + tp_size = get_tensor_model_parallel_world_size()
402 + current_rank = get_tensor_model_parallel_rank()
403 +
404 + per_item_input_sizes = self.model.get_encoder_cudagraph_per_item_input_sizes(
405 + mm_kwargs
406 + )
407 +
408 + (image_rank_assignment, images_per_rank, input_patches_per_rank) = (
409 + get_load_balance_assignment(per_item_input_sizes, tp_size)
410 + )
411 +
412 + # Extract local indices for this rank
413 + cum_images_per_rank = [0]
414 + for count in images_per_rank:
415 + cum_images_per_rank.append(cum_images_per_rank[-1] + count)
416 +
417 + local_indices = image_rank_assignment[
418 + cum_images_per_rank[current_rank] : cum_images_per_rank[current_rank + 1]
419 + ]
420 +
421 + if len(local_indices) > 0:
422 + local_mm_kwargs = self.model.select_encoder_cudagraph_items(
423 + mm_kwargs, local_indices
424 + )
425 + else:
426 + local_mm_kwargs = self.model.select_encoder_cudagraph_items(mm_kwargs, [])
427 +
428 + max_output_tokens_per_rank = (
429 + max(
430 + sum(
431 + per_item_out_tokens[i]
432 + for i in image_rank_assignment[
433 + cum_images_per_rank[r] : cum_images_per_rank[r + 1]
434 + ]
435 + )
436 + for r in range(tp_size)
437 + )
438 + if len(per_item_out_tokens) > 0
439 + else 0
440 + )
441 +
442 + return (
443 + local_mm_kwargs,
444 + image_rank_assignment,
445 + images_per_rank,
446 + max_output_tokens_per_rank,
447 + )
448 +
449 + def _dp_gather(
450 + self,
451 + local_outputs: list[torch.Tensor],
452 + per_item_out_tokens: list[int],
453 + image_rank_assignment: list[int],
454 + images_per_rank: list[int],
455 + max_output_tokens_per_rank: int,
456 + ) -> list[torch.Tensor]:
457 + """Gather outputs from all TP ranks and reorder to original sequence.
458 +
459 + Assumes 2D output tensors [tokens, hidden]. Follows the same
460 + pad -> all_gather -> unpad -> reorder algorithm as
461 + run_dp_sharded_mrope_vision_model() in the eager path.
462 + """
463 + hidden_size = self.config.out_hidden_size
464 + tp_size = len(images_per_rank)
465 +
466 + if len(local_outputs) > 0:
467 + local_concat = torch.cat(local_outputs, dim=0)
468 + else:
469 + local_concat = torch.empty(
470 + (0, hidden_size), device=self.device, dtype=self.dtype
471 + )
472 +
473 + # Pad to max_output_tokens_per_rank for all_gather
474 + current_len = local_concat.shape[0]
475 + if current_len < max_output_tokens_per_rank:
476 + padding = torch.empty(
477 + (max_output_tokens_per_rank - current_len, hidden_size),
478 + dtype=self.dtype,
479 + device=self.device,
480 + )
481 + local_padded = torch.cat([local_concat, padding], dim=0)
482 + else:
483 + local_padded = local_concat
484 +
485 + gathered = tensor_model_parallel_all_gather(local_padded, dim=0)
486 +
487 + # Unpad each rank's contribution
488 + rank_outputs: list[torch.Tensor] = []
489 + current_idx = 0
490 + for rank in range(tp_size):
491 + start = rank * max_output_tokens_per_rank
492 + rank_count = images_per_rank[rank]
493 + rank_indices = image_rank_assignment[current_idx : current_idx + rank_count]
494 + rank_tokens = sum(per_item_out_tokens[i] for i in rank_indices)
495 + current_idx += rank_count
496 + rank_outputs.append(gathered[start : start + rank_tokens])
497 +
498 + # Reorder to original sequence
499 + total_items = len(per_item_out_tokens)
500 + result: list[torch.Tensor | None] = [None] * total_items
501 + current_idx = 0
502 + for rank in range(tp_size):
503 + count = images_per_rank[rank]
504 + if count > 0:
505 + rank_items = image_rank_assignment[current_idx : current_idx + count]
506 + self._scatter_output_slices(
507 + rank_outputs[rank],
508 + rank_items,
509 + per_item_out_tokens,
510 + result,
511 + )
512 + current_idx += count
513 +
514 + return [t for t in result if t is not None]
515 +
516 + def execute(
517 + self,
518 + mm_kwargs: dict[str, Any],
519 + ) -> list[torch.Tensor]:
520 + """Execute encoder using CUDA graph with optional DP.
521 +
522 + Args:
523 + mm_kwargs: Multimodal keyword arguments containing the
524 + input tensor and grid dimensions.
525 +
526 + Returns:
527 + List of encoder outputs (one per item).
528 + """
529 + if self.use_dp:
530 + per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs)
531 +
532 + (
533 + local_mm_kwargs,
534 + image_rank_assignment,
535 + images_per_rank,
536 + max_output_tokens_per_rank,
537 + ) = self._dp_shard(mm_kwargs, per_item_out_tokens)
538 +
539 + local_outputs = self._execute_local(local_mm_kwargs)
540 +
541 + result = self._dp_gather(
542 + local_outputs,
543 + per_item_out_tokens,
544 + image_rank_assignment,
545 + images_per_rank,
546 + max_output_tokens_per_rank,
547 + )
548 + else:
549 + result = self._execute_local(mm_kwargs)
550 +
551 + # Log cumulative stats periodically
552 + stats = self.get_cumulative_stats()
553 + total_requests = self.graph_hits + self.graph_misses
554 + if total_requests > 0 and total_requests % self.log_stats_interval == 0:
555 + logger.debug(
556 + "Encoder CUDA graph cumulative stats: "
557 + "hits=%d, misses=%d, hit_rate=%.1f%%",
558 + stats["graph_hits"],
559 + stats["graph_misses"],
560 + stats["hit_rate"] * 100,
561 + )
562 +
563 + return result
564 +
565 + def get_cumulative_stats(self) -> dict[str, Any]:
566 + """Get cumulative CUDA graph statistics."""
567 + total_requests = self.graph_hits + self.graph_misses
568 + hit_rate = self.graph_hits / total_requests if total_requests > 0 else 0.0
569 +
570 + return {
571 + "graph_hits": self.graph_hits,
572 + "graph_misses": self.graph_misses,
573 + "hit_rate": hit_rate,
574 + "num_budgets": len(self.budget_graphs),
575 + "token_budgets": self.token_budgets,
576 + }
vllm/v1/worker/gpu/mm/encoder_cudagraph_defs.py new
+66
@@ -0,0 +1,66 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Data transfer objects for encoder CUDA graph management."""
4 +
5 +from dataclasses import dataclass
6 +from typing import Any
7 +
8 +import torch
9 +
10 +
11 +@dataclass
12 +class EncoderCudaGraphConfig:
13 + """Configuration for encoder CUDA graph management.
14 +
15 + Provided by the model at init time via
16 + ``get_encoder_cudagraph_config()``. Values are fixed for the
17 + lifetime of the manager.
18 + """
19 +
20 + modalities: list[str]
21 + """Supported modalities (e.g. ["image"])."""
22 +
23 + input_key: str
24 + """Key in mm_kwargs for the input tensor (e.g. "pixel_values")."""
25 +
26 + buffer_keys: list[str]
27 + """Keys for the tensor buffers recorded into the CUDA graph.
28 + Before replay the manager zeros then slice-copies new data
29 + into these buffers."""
30 +
31 + out_hidden_size: int
32 + """Output hidden dim of the vision encoder.
33 + Used for DP gather buffer allocation."""
34 +
35 +
36 +@dataclass
37 +class EncoderCudaGraphCaptureInputs:
38 + """Everything needed for one CUDA graph capture.
39 +
40 + Returned by ``prepare_encoder_cudagraph_capture_inputs()``.
41 + """
42 +
43 + mm_kwargs: dict[str, Any]
44 + """Dummy forward inputs (model-specific keys).
45 + For Qwen3-VL this contains pixel_values and grid_thw."""
46 +
47 + buffers: dict[str, torch.Tensor]
48 + """Precomputed tensor buffers that will be recorded into the
49 + CUDA graph. The manager stores references to these exact
50 + tensor objects and copies new data into them before each
51 + ``graph.replay()`` call (buffer identity invariant)."""
52 +
53 +
54 +@dataclass
55 +class EncoderCudaGraphReplayBuffers:
56 + """New buffer values for graph replay, computed by the model from
57 + actual batch inputs.
58 +
59 + Returned by ``prepare_encoder_cudagraph_replay_buffers()``.
60 + Keys match ``EncoderCudaGraphConfig.buffer_keys``.
61 + """
62 +
63 + buffers: dict[str, torch.Tensor | None]
64 + """Data to copy into the captured buffers before replay.
65 + ``None`` values leave the corresponding captured buffer
66 + unchanged."""
vllm/v1/worker/gpu_model_runner.py
+48 -1
@@ -207,6 +207,7 @@ from .utils import (
207 if TYPE_CHECKING:
208 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
209 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
210 + from vllm.v1.worker.gpu.mm.encoder_cudagraph import EncoderCudaGraphManager
211
212 logger = init_logger(__name__)
213
@@ -499,6 +500,9 @@ class GPUModelRunner(
500 self.encoder_cache: dict[str, torch.Tensor] = {}
501 self.late_interaction_runner = LateInteractionRunner()
502
503 + # Encoder CUDA graph manager (initialized after model load if enabled)
504 + self.encoder_cudagraph_manager: EncoderCudaGraphManager | None = None
505 +
506 self.use_aux_hidden_state_outputs = False
507 # Set up speculative decoding.
508 # NOTE(Jiayi): currently we put the entire draft model on
@@ -2664,7 +2668,19 @@ class GPUModelRunner(
2668 with self.timed_encoder_operation(
2669 should_time, mm_lora_refs, current_item_idx, num_items
2670 ):
2667 - batch_outputs = model.embed_multimodal(**mm_kwargs_batch)
2671 + cudagraph_output = None
2672 + if (
2673 + self.encoder_cudagraph_manager is not None
2674 + and self.encoder_cudagraph_manager.supports_modality(modality)
2675 + ):
2676 + cudagraph_output = self.encoder_cudagraph_manager.execute(
2677 + mm_kwargs_batch,
2678 + )
2679 +
2680 + if cudagraph_output is not None:
2681 + batch_outputs = cudagraph_output
2682 + else:
2683 + batch_outputs = model.embed_multimodal(**mm_kwargs_batch)
2684
2685 sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items)
2686 encoder_outputs.extend(batch_outputs)
@@ -5715,6 +5731,33 @@ class GPUModelRunner(
5731 )
5732 return 0
5733
5734 + # Initialize encoder CUDA graph manager if enabled.
5735 + # Use get_model() to unwrap CUDAGraphWrapper/UBatchWrapper,
5736 + # because @runtime_checkable Protocol isinstance() checks do not
5737 + # work through __getattr__ forwarding.
5738 + if (
5739 + self.compilation_config.cudagraph_mm_encoder
5740 + and self.supports_mm_inputs
5741 + and self.encoder_cudagraph_manager is None
5742 + ):
5743 + from vllm.model_executor.models.interfaces import (
5744 + SupportsEncoderCudaGraph,
5745 + supports_encoder_cudagraph,
5746 + )
5747 + from vllm.v1.worker.gpu.mm.encoder_cudagraph import (
5748 + EncoderCudaGraphManager,
5749 + )
5750 +
5751 + raw_model = self.get_model()
5752 + if supports_encoder_cudagraph(raw_model):
5753 + self.encoder_cudagraph_manager = EncoderCudaGraphManager(
5754 + vllm_config=self.vllm_config,
5755 + device=self.device,
5756 + dtype=self.dtype,
5757 + model=cast(SupportsEncoderCudaGraph, raw_model),
5758 + )
5759 + logger.info("Initialized EncoderCudaGraphManager for vision encoder")
5760 +
5761 compilation_counter.num_gpu_runner_capture_triggers += 1
5762
5763 start_time = time.perf_counter()
@@ -5738,6 +5781,10 @@ class GPUModelRunner(
5781 )
5782 torch.accelerator.synchronize()
5783
5784 + # Capture encoder CUDA graphs if enabled
5785 + if self.encoder_cudagraph_manager is not None:
5786 + self.encoder_cudagraph_manager.capture()
5787 +
5788 torch.accelerator.synchronize()
5789 end_free_gpu_memory = torch.cuda.mem_get_info()[0]
5790