@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 80118853f4

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 5, 04:45 UTC

[MM][Perf][CG] Support ViT full CUDA graph for Qwen3-VL video inference (#38061)

Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.io>

Shanshan Shen committed Apr 14, 2026 at 16:49 UTC 80118853f42a728805e6a861ffb50f3b876a9939

7 files changed +583 -68

docs/design/cuda_graphs_multimodal.md

+66 -7

@@ -28,6 +28,7 @@ Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g.

28 class BudgetGraphMetadata:

29 token_budget: int

30 max_batch_size: int

31 + max_frames_per_batch: int

32 graph: torch.cuda.CUDAGraph

33 input_buffer: torch.Tensor # e.g. pixel_values

34 metadata_buffers: dict[str, torch.Tensor] # e.g. embeddings, seq metadata

@@ -51,6 +52,15 @@ For each graph replay:

53 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`.

55 +### Video inference support (experimental)

56 +

57 +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.

58 +

59 +!!! note

60 + Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.

61 +

62 + Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.

63 +

64 ## Model integration via `SupportsEncoderCudaGraph`

66 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:

@@ -65,12 +75,17 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra

75 * `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay.

76 * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).

77 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.

68 -

69 -Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`).

78 +* `get_input_modality(...)` - return the modality of the inputs.

80 !!! note

81 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.

83 +**Supported models:**

84 +

86 +| ------------ | ------ | ------------ | ------------ |

87 +| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |

88 +

89 !!! note

90 Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.

@@ -80,10 +95,13 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:

96 * `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.

97 * `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.

83 -* `encoder_cudagraph_max_images_per_batch` (`int`, default `0`) — maximum number of images per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.

98 +* `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`.

99 +* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).

100

101 ## Usage guide

102

103 +### Image inference

104 +

105 Enable encoder CUDA Graphs via `compilation_config`:

106

107 ```bash

@@ -95,7 +113,7 @@ With explicit budgets:

113

114 ```bash

115 vllm serve Qwen/Qwen3-VL-32B \

98 - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}'

116 + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}'

117 ```

118

119 Python example:

@@ -107,7 +125,7 @@ compilation_config = {

125 "cudagraph_mm_encoder": True,

126 # Optional: override auto-inferred budgets

127 # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],

110 - # "encoder_cudagraph_max_images_per_batch": 8,

128 + # "encoder_cudagraph_max_vision_items_per_batch": 8,

129 }

130

131 model = vllm.LLM(

@@ -118,6 +136,44 @@ model = vllm.LLM(

136

137 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).

138

139 +### Video inference

140 +

141 +Enable encoder CUDA Graphs via `compilation_config`:

142 +

143 +```bash

144 +vllm serve Qwen/Qwen3-VL-32B \

145 + --limit-mm-per-prompt '{"image": 0}' \

146 + --compilation-config '{"cudagraph_mm_encoder": true}'

147 +```

148 +

149 +With explicit budgets:

150 +

151 +```bash

152 +vllm serve Qwen/Qwen3-VL-32B \

153 + --limit-mm-per-prompt '{"image": 0}' \

154 + --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}'

155 +```

156 +

157 +Python example:

158 +

159 +```python

160 +import vllm

161 +

162 +compilation_config = {

163 + "cudagraph_mm_encoder": True,

164 + # Optional: override auto-inferred budgets

165 + # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],

166 + # "encoder_cudagraph_max_vision_items_per_batch": 8,

167 + # "encoder_cudagraph_max_frames_per_batch": 64,

168 +}

169 +

170 +model = vllm.LLM(

171 + model="Qwen/Qwen3-VL-32B",

172 + limit_mm_per_prompt='{"image": 0}',

173 + compilation_config=compilation_config,

174 +)

175 +```

176 +

177 ## About the Performance

178

179 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.

@@ -140,7 +196,7 @@ vllm bench mm-processor \

196 --num-prompts 3000 --num-warmups 300 \

197 --max-model-len 32768 --seed 42 \

198 --mm-encoder-attn-backend FLASH_ATTN \

143 - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'

199 + --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}'

200 ```

201

202 ### Multi-GPU (4x GB200, TP=4, DP=4)

@@ -165,5 +221,8 @@ vllm bench mm-processor \

221 --max-model-len 8192 --seed 42 \

222 --mm-encoder-attn-backend FLASHINFER \

223 --tensor-parallel-size 4 --mm-encoder-tp-mode data \

168 - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'

224 + --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}'

225 ```

226 +

227 +!!! note

228 + Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061).

tests/v1/cudagraph/test_encoder_cudagraph.py

+312 -1

@@ -6,8 +6,10 @@ Test organization:

6 No GPU required:

7 - TestFindBudgetGraph — greedy budget selection logic

8 - TestGetCumulativeStats — hit/miss rate statistics

9 + - TestGetInputModality — modality routing from mm_kwargs keys

10 GPU required:

11 - TestEncoderCudaGraphCaptureReplay — capture, replay, fallback, counters, chunking

12 + - TestEncoderCudaGraphVideoReplay — video modality capture, replay

13 """

15 from typing import Any

@@ -205,11 +207,19 @@ class SimpleMockViTModel(torch.nn.Module):

207 def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:

208 return EncoderCudaGraphConfig(

209 modalities=["image"],

208 - input_key="pixel_values",

210 + input_key_by_modality={

211 + "image": "pixel_values",

212 + },

213 buffer_keys=["dummy_buf"],

214 out_hidden_size=_HIDDEN,

215 )

216

217 + def get_input_modality(

218 + self,

219 + mm_kwargs: dict[str, Any],

220 + ) -> str:

221 + return "image"

222 +

223 def get_encoder_cudagraph_budget_range(

224 self,

225 vllm_config,

@@ -268,6 +278,7 @@ class SimpleMockViTModel(torch.nn.Module):

278 self,

279 token_budget: int,

280 max_batch_size: int,

281 + max_frames_per_batch: int,

282 device: torch.device,

283 dtype: torch.dtype,

284 ) -> EncoderCudaGraphCaptureInputs:

@@ -294,6 +305,7 @@ class SimpleMockViTModel(torch.nn.Module):

305 self,

306 mm_kwargs: dict[str, Any],

307 max_batch_size: int,

308 + max_frames_per_batch: int,

309 ) -> EncoderCudaGraphReplayBuffers:

310 grid_thw = mm_kwargs["image_grid_thw"]

311 n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)

@@ -327,11 +339,16 @@ def _make_manager_for_gpu(

339 max_batch_size: int,

340 device: torch.device,

341 dtype: torch.dtype,

342 + *,

343 + max_frames_per_batch: int | None = None,

344 ) -> EncoderCudaGraphManager:

345 """Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests."""

346 mgr = object.__new__(EncoderCudaGraphManager)

347 mgr.token_budgets = sorted(token_budgets)

348 mgr.max_batch_size = max_batch_size

349 + mgr.max_frames_per_batch = (

350 + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size * 2

351 + )

352 mgr.use_dp = False

353 mgr.budget_graphs = {}

354 mgr.graph_hits = 0

@@ -366,6 +383,18 @@ def _make_mm_kwargs(

383 }

384

385

386 +def _make_video_mm_kwargs(

387 + grid_thw_list: list[list[int]],

388 + device: torch.device,

389 + dtype: torch.dtype,

390 +) -> dict[str, Any]:

391 + """Create video mm_kwargs (pixel_values_videos / video_grid_thw) for testing."""

392 + return {

393 + "pixel_values_videos": _make_pixel_values(grid_thw_list, device, dtype),

394 + "video_grid_thw": grid_thw_list,

395 + }

396 +

397 +

398 # ---------------------------------------------------------------------------

399 # GPU tests — capture, replay, fallback, counters, chunking

400 # ---------------------------------------------------------------------------

@@ -449,3 +478,285 @@ class TestEncoderCudaGraphCaptureReplay:

478 assert len(result) == n_images

479 for out in result:

480 assert out.shape == (4, _HIDDEN)

481 +

482 +

483 +# ---------------------------------------------------------------------------

484 +# SimpleMockViTVideoModel — extends SimpleMockViTModel with video support

485 +# ---------------------------------------------------------------------------

486 +

487 +

488 +class SimpleMockViTVideoModel(SimpleMockViTModel):

489 + """ViT mock that supports both image and video modalities.

490 +

491 + Reuses SimpleMockViTModel's NN weights and _forward() logic.

492 + Only the protocol methods that are key-dependent are overridden.

493 + """

494 +

495 + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:

496 + return EncoderCudaGraphConfig(

497 + modalities=["image", "video"],

498 + input_key_by_modality={

499 + "image": "pixel_values",

500 + "video": "pixel_values_videos",

501 + },

502 + buffer_keys=["dummy_buf"],

503 + out_hidden_size=_HIDDEN,

504 + )

505 +

506 + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:

507 + return "video" if "video_grid_thw" in mm_kwargs else "image"

508 +

509 + # ------------------------------------------------------------------

510 + # Private helpers — route to the correct mm_kwargs keys

511 + # ------------------------------------------------------------------

512 +

513 + def _get_grid_thw(self, mm_kwargs: dict[str, Any]) -> list[list[int]]:

514 + key = (

515 + "video_grid_thw"

516 + if self.get_input_modality(mm_kwargs) == "video"

517 + else "image_grid_thw"

518 + )

519 + return mm_kwargs[key]

520 +

521 + def _get_pixel_values(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:

522 + key = (

523 + "pixel_values_videos"

524 + if self.get_input_modality(mm_kwargs) == "video"

525 + else "pixel_values"

526 + )

527 + return mm_kwargs[key]

528 +

529 + # ------------------------------------------------------------------

530 + # Protocol overrides that depend on modality keys

531 + # ------------------------------------------------------------------

532 +

533 + def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:

534 + return len(self._get_grid_thw(mm_kwargs))

535 +

536 + def get_encoder_cudagraph_per_item_output_tokens(

537 + self, mm_kwargs: dict[str, Any]

538 + ) -> list[int]:

539 + m = _SPATIAL_MERGE

540 + return [t * (h // m) * (w // m) for t, h, w in self._get_grid_thw(mm_kwargs)]

541 +

542 + def get_encoder_cudagraph_per_item_input_sizes(

543 + self, mm_kwargs: dict[str, Any]

544 + ) -> list[int]:

545 + return [t * h * w for t, h, w in self._get_grid_thw(mm_kwargs)]

546 +

547 + def select_encoder_cudagraph_items(

548 + self, mm_kwargs: dict[str, Any], indices: list[int]

549 + ) -> dict[str, Any]:

550 + modality = self.get_input_modality(mm_kwargs)

551 + pv_key = "pixel_values_videos" if modality == "video" else "pixel_values"

552 + grid_key = "video_grid_thw" if modality == "video" else "image_grid_thw"

553 +

554 + grid_thw = self._get_grid_thw(mm_kwargs)

555 + pixel_values = self._get_pixel_values(mm_kwargs)

556 +

557 + if len(indices) == 0:

558 + return {pv_key: pixel_values[:0], grid_key: []}

559 +

560 + patches_per_item = [t * h * w for t, h, w in grid_thw]

561 + cum_patches = [0]

562 + for p in patches_per_item:

563 + cum_patches.append(cum_patches[-1] + p)

564 +

565 + selected_pv = torch.cat(

566 + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]

567 + )

568 + return {pv_key: selected_pv, grid_key: [grid_thw[i] for i in indices]}

569 +

570 + def prepare_encoder_cudagraph_capture_inputs(

571 + self,

572 + token_budget: int,

573 + max_batch_size: int,

574 + max_frames_per_batch: int,

575 + device: torch.device,

576 + dtype: torch.dtype,

577 + ) -> EncoderCudaGraphCaptureInputs:

578 + per_item_output = token_budget // max_batch_size

579 + frames_per_item = max_frames_per_batch // max_batch_size

580 + if frames_per_item > 1:

581 + # Video-format capture: size cu_seqlens for T frames per item.

582 + tokens_per_frame = (

583 + per_item_output + frames_per_item - 1

584 + ) // frames_per_item

585 + grid_config = [

586 + [frames_per_item, _SPATIAL_MERGE, tokens_per_frame * _SPATIAL_MERGE]

587 + for _ in range(max_batch_size)

588 + ]

589 + else:

590 + grid_config = [

591 + [1, _SPATIAL_MERGE, per_item_output * _SPATIAL_MERGE]

592 + for _ in range(max_batch_size)

593 + ]

594 + total_patches = _count_input_patches(grid_config)

595 + # Use pixel_values (image key) for capture — same patch shape as video.

596 + dummy_pixel_values = torch.randn(

597 + total_patches, _FLAT, device=device, dtype=dtype

598 + )

599 + n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)

600 + dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)

601 + return EncoderCudaGraphCaptureInputs(

602 + mm_kwargs={

603 + "pixel_values": dummy_pixel_values,

604 + "image_grid_thw": grid_config,

605 + },

606 + buffers={"dummy_buf": dummy_buf},

607 + )

608 +

609 + def prepare_encoder_cudagraph_replay_buffers(

610 + self,

611 + mm_kwargs: dict[str, Any],

612 + max_batch_size: int,

613 + max_frames_per_batch: int,

614 + ) -> EncoderCudaGraphReplayBuffers:

615 + n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE)

616 + p = next(self.parameters())

617 + dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)

618 + return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})

619 +

620 + def encoder_cudagraph_forward(

621 + self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor]

622 + ) -> torch.Tensor:

623 + return self._forward(self._get_pixel_values(mm_kwargs))

624 +

625 + def encoder_eager_forward(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:

626 + return self._forward(self._get_pixel_values(mm_kwargs))

627 +

628 +

629 +# ---------------------------------------------------------------------------

630 +# No-GPU tests — get_input_modality routing

631 +# ---------------------------------------------------------------------------

632 +

633 +

634 +class TestGetInputModality:

635 + """get_input_modality returns correct modality based on mm_kwargs keys."""

636 +

637 + def test_image_only_model_always_returns_image(self):

638 + model = SimpleMockViTModel()

639 + mm_kwargs = {

640 + "pixel_values": torch.zeros(1, _FLAT),

641 + "image_grid_thw": [[1, 4, 4]],

642 + }

643 + assert model.get_input_modality(mm_kwargs) == "image"

644 +

645 + def test_video_model_returns_image_for_image_kwargs(self):

646 + model = SimpleMockViTVideoModel()

647 + mm_kwargs = {

648 + "pixel_values": torch.zeros(1, _FLAT),

649 + "image_grid_thw": [[1, 4, 4]],

650 + }

651 + assert model.get_input_modality(mm_kwargs) == "image"

652 +

653 + def test_video_model_returns_video_for_video_kwargs(self):

654 + model = SimpleMockViTVideoModel()

655 + mm_kwargs = {

656 + "pixel_values_videos": torch.zeros(8, _FLAT),

657 + "video_grid_thw": [[2, 4, 4]],

658 + }

659 + assert model.get_input_modality(mm_kwargs) == "video"

660 +

661 + def test_video_model_config_has_both_modalities(self):

662 + model = SimpleMockViTVideoModel()

663 + cfg = model.get_encoder_cudagraph_config()

664 + assert "image" in cfg.modalities

665 + assert "video" in cfg.modalities

666 + assert cfg.input_key_by_modality["image"] == "pixel_values"

667 + assert cfg.input_key_by_modality["video"] == "pixel_values_videos"

668 +

669 +

670 +# ---------------------------------------------------------------------------

671 +# GPU tests — video capture, replay, fallback, and mixed image+video

672 +# ---------------------------------------------------------------------------

673 +

674 +_VIDEO_MAX_BATCH = 4

675 +_VIDEO_MAX_FRAMES = 8 # 2 frames per item at max_batch_size=4

676 +

677 +

678 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")

679 +class TestEncoderCudaGraphVideoReplay:

680 + def setup_method(self):

681 + self.device = torch.device("cuda:0")

682 + self.dtype = torch.float16

683 + self.model = SimpleMockViTVideoModel().to(self.device).half()

684 + self.mgr = _make_manager_for_gpu(

685 + self.model,

686 + _BUDGETS,

687 + _VIDEO_MAX_BATCH,

688 + self.device,

689 + self.dtype,

690 + max_frames_per_batch=_VIDEO_MAX_FRAMES,

691 + )

692 + self.mgr.capture()

693 +

694 + # --- capture ---

695 +

696 + def test_capture_creates_one_graph_per_budget(self):

697 + assert len(self.mgr.budget_graphs) == len(_BUDGETS)

698 + assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)

699 +

700 + # --- output shape ---

701 +

702 + def test_video_execute_returns_one_tensor_per_video(self):

703 + # T=2, 4x4 → 2*(4//2)*(4//2) = 8 tokens per video

704 + grid_thw = [[2, 4, 4], [2, 4, 4]]

705 + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)

706 + result = self.mgr.execute(mm_kwargs)

707 + assert result is not None

708 + assert len(result) == 2

709 +

710 + def test_video_output_tokens_per_item(self):

711 + # T=2,4x4 → 8 tokens; T=1,4x4 → 4 tokens

712 + grid_thw = [[2, 4, 4], [1, 4, 4]]

713 + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)

714 + result = self.mgr.execute(mm_kwargs)

715 + assert result is not None

716 + assert result[0].shape == (8, _HIDDEN)

717 + assert result[1].shape == (4, _HIDDEN)

718 +

719 + # --- budget fallback ---

720 +

721 + def test_video_eager_fallback_when_tokens_exceed_all_budgets(self):

722 + # T=2, 18x18 → 2*(18//2)*(18//2) = 162 tokens > max budget 64

723 + grid_thw = [[2, 18, 18]]

724 + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)

725 + result = self.mgr.execute(mm_kwargs)

726 + assert result is not None

727 + assert len(result) == 1

728 + assert result[0].shape == (162, _HIDDEN)

729 + assert self.mgr.graph_misses == 1

730 +

731 + # --- counters ---

732 +

733 + def test_video_hit_counter_increments_by_num_videos(self):

734 + grid_thw = [[2, 4, 4], [1, 4, 4]]

735 + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)

736 + self.mgr.execute(mm_kwargs)

737 + assert self.mgr.graph_hits == 2

738 +

739 + def test_video_miss_counter_increments_for_oversized_video(self):

740 + grid_thw = [[2, 18, 18]] # 162 tokens > 64

741 + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)

742 + self.mgr.execute(mm_kwargs)

743 + assert self.mgr.graph_misses == 1

744 +

745 + # --- image and video sharing the same manager ---

746 +

747 + def test_image_and_video_share_manager(self):

748 + """Image and video inputs can both be executed through the same manager."""

749 + img_grid = [[1, 4, 4], [1, 4, 4]]

750 + img_result = self.mgr.execute(

751 + _make_mm_kwargs(img_grid, self.device, self.dtype)

752 + )

753 +

754 + vid_grid = [[2, 4, 4]]

755 + vid_result = self.mgr.execute(

756 + _make_video_mm_kwargs(vid_grid, self.device, self.dtype)

757 + )

758 +

759 + assert len(img_result) == 2

760 + assert len(vid_result) == 1

761 + assert img_result[0].shape == (4, _HIDDEN)

762 + assert vid_result[0].shape == (8, _HIDDEN)

vllm/config/compilation.py

+20 -4

@@ -519,13 +519,21 @@ class CompilationConfig:

519 User-provided values override auto-inference.

520 Example: [2048, 4096, 8192, 13824]"""

521

522 - encoder_cudagraph_max_images_per_batch: int = 0

523 - """Maximum number of images per batch for encoder CUDA graph capture.

522 + encoder_cudagraph_max_vision_items_per_batch: int = 0

523 + """Maximum number of images/videos per batch for encoder CUDA graph capture.

524 Determines the fixed batch size used during graph capture.

525 If 0 (default), auto-inferred as max_budget // min_budget from the

526 model's budget range. User-provided positive value overrides

527 auto-inference."""

528

529 + encoder_cudagraph_max_frames_per_batch: int = 0

530 + """Maximum total video frames per batch for encoder CUDA graph capture.

531 + Controls the cu_seqlens buffer size (one entry per attention sequence,

532 + i.e. one per video frame). If 0 (default), auto-inferred per budget

533 + level as token_budget (tight bound: packing guarantees

534 + sum(T_i) <= token_budget). Positive value overrides auto-inference

535 + and applies to all budget levels."""

536 +

537 # Inductor capture

538 compile_sizes: list[int | str] | None = None

539 """Sizes to compile for inductor. In addition

@@ -964,10 +972,18 @@ class CompilationConfig:

972 # Validate encoder CUDA graph configuration

973 if (

974 self.cudagraph_mm_encoder

967 - and self.encoder_cudagraph_max_images_per_batch < 0

975 + and self.encoder_cudagraph_max_vision_items_per_batch < 0

976 + ):

977 + raise ValueError(

978 + "encoder_cudagraph_max_vision_items_per_batch must be "

979 + "non-negative (0 = auto-infer)"

980 + )

981 + if (

982 + self.cudagraph_mm_encoder

983 + and self.encoder_cudagraph_max_frames_per_batch < 0

984 ):

985 raise ValueError(

970 - "encoder_cudagraph_max_images_per_batch must be "

986 + "encoder_cudagraph_max_frames_per_batch must be "

987 "non-negative (0 = auto-infer)"

988 )

989

vllm/model_executor/models/interfaces.py

+10 -1

@@ -1524,6 +1524,13 @@ class SupportsEncoderCudaGraph(Protocol):

1524

1525 def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": ...

1526

1527 + def get_input_modality(

1528 + self,

1529 + mm_kwargs: dict[str, Any],

1530 + ) -> str:

1531 + """Return the modality of the inputs."""

1532 + ...

1533 +

1534 def get_encoder_cudagraph_budget_range(

1535 self,

1536 vllm_config: "VllmConfig",

@@ -1536,7 +1543,7 @@ class SupportsEncoderCudaGraph(Protocol):

1543 (e.g. max_num_batched_tokens)

1544

1545 Used when ``encoder_cudagraph_token_budgets`` and/or

1539 - ``encoder_cudagraph_max_images_per_batch`` are not explicitly

1546 + ``encoder_cudagraph_max_vision_items_per_batch`` are not explicitly

1547 specified by the user.

1548 """

1549 ...

@@ -1590,6 +1597,7 @@ class SupportsEncoderCudaGraph(Protocol):

1597 self,

1598 token_budget: int,

1599 max_batch_size: int,

1600 + max_frames_per_batch: int,

1601 device: torch.device,

1602 dtype: torch.dtype,

1603 ) -> "EncoderCudaGraphCaptureInputs":

@@ -1600,6 +1608,7 @@ class SupportsEncoderCudaGraph(Protocol):

1608 self,

1609 mm_kwargs: dict[str, Any],

1610 max_batch_size: int,

1611 + max_frames_per_batch: int,

1612 ) -> "EncoderCudaGraphReplayBuffers":

1613 """Compute buffer values from actual batch inputs for replay."""

1614 ...

vllm/model_executor/models/qwen3_vl.py

+138 -42

@@ -99,6 +99,7 @@ from vllm.tokenizers.registry import cached_tokenizer_from_config

99 from vllm.triton_utils import HAS_TRITON, tl, triton

100 from vllm.utils.collection_utils import is_list_of

101 from vllm.utils.math_utils import round_up

102 +from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

103

104 from .interfaces import (

105 MultiModalEmbeddings,

@@ -689,6 +690,7 @@ class Qwen3_VisionTransformer(nn.Module):

690 grid_thw_list: list[list[int]],

691 *,

692 max_batch_size: int | None = None,

693 + max_frames_per_batch: int | None = None,

694 max_seqlen_override: int | None = None,

695 device: torch.device | None = None,

696 ) -> dict[str, torch.Tensor | None]:

@@ -701,6 +703,10 @@ class Qwen3_VisionTransformer(nn.Module):

703 grid_thw_list: Grid configurations as list of [t, h, w].

704 max_batch_size: If set, pad cu_seqlens to this size

705 (needed for CUDA graph capture/replay).

706 + max_frames_per_batch: If set, overrides max_batch_size for

707 + cu_seqlens padding. For video inputs each item contributes

708 + T attention sequences (frames); this sizes the buffer to

709 + the total frame budget so video replays never overflow.

710 max_seqlen_override: If set, use this value for max_seqlen

711 instead of computing from cu_seqlens (needed for CUDA

712 graph capture to cover worst-case replay scenarios).

@@ -725,15 +731,21 @@ class Qwen3_VisionTransformer(nn.Module):

731 )

732 cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])

733

728 - # Pad cu_seqlens if max_batch_size specified

729 - if max_batch_size is not None:

734 + # Pad cu_seqlens to the required number of sequences.

735 + # For videos each item contributes T frames = T attention sequences,

736 + # so the total can exceed max_batch_size. max_frames_per_batch

737 + # overrides the pad target when set.

738 + pad_to = (

739 + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size

740 + )

741 + if pad_to is not None:

742 num_seqs = len(cu_seqlens) - 1

731 - if num_seqs < max_batch_size:

743 + if num_seqs < pad_to:

744 cu_seqlens = np.concatenate(

745 [

746 cu_seqlens,

747 np.full(

736 - max_batch_size - num_seqs,

748 + pad_to - num_seqs,

749 cu_seqlens[-1],

750 dtype=np.int32,

751 ),

@@ -1737,9 +1749,21 @@ class Qwen3VLForConditionalGeneration(

1749 EncoderCudaGraphConfig,

1750 )

1751

1752 + modalities = ["image"]

1753 + # NOTE: When EVS (Efficient Video Sampling) pruning is enabled, the number

1754 + # of tokens becomes data-dependent (i.e., the retained tokens are

1755 + # dynamically selected based on inter-frame differences) and therefore

1756 + # cannot be captured by CUDA Graphs. As a result, video CUDA Graphs are

1757 + # only enabled when EVS is disabled.

1758 + if not self.is_multimodal_pruning_enabled:

1759 + modalities.append("video")

1760 +

1761 return EncoderCudaGraphConfig(

1741 - modalities=["image"],

1742 - input_key="pixel_values",

1762 + modalities=modalities,

1763 + input_key_by_modality={

1764 + "image": "pixel_values",

1765 + "video": "pixel_values_videos",

1766 + },

1767 buffer_keys=[

1768 "pos_embeds",

1769 "rotary_pos_emb_cos",

@@ -1751,49 +1775,86 @@ class Qwen3VLForConditionalGeneration(

1775 out_hidden_size=self.visual.out_hidden_size,

1776 )

1777

1778 + def get_input_modality(

1779 + self,

1780 + mm_kwargs: dict[str, Any],

1781 + ) -> str:

1782 + if "image_grid_thw" in mm_kwargs:

1783 + return "image"

1784 + return "video"

1785 +

1786 def get_encoder_cudagraph_budget_range(

1787 self,

1788 vllm_config,

1789 ) -> tuple[int, int]:

1790 # Min: estimated smallest possible encoder input.

1759 - # 224x224 image → 16x16 patches, spatial_merge_size=2 → 8x8 = 64 tokens

1791 + # 224x224 image → 16x16 patches (patch_size=14)

1792 + # spatial_merge_size=2 → 8x8 = 64 tokens

1793 min_budget = 64

1794 # Max: capped by max_num_batched_tokens

1795 max_budget = vllm_config.scheduler_config.max_num_batched_tokens

1796 return (min_budget, max_budget)

1797

1798 + def _get_pixel_values_by_modality(

1799 + self,

1800 + mm_kwargs: dict[str, Any],

1801 + ) -> torch.Tensor:

1802 + if self.get_input_modality(mm_kwargs) == "image":

1803 + pixel_values = mm_kwargs["pixel_values"]

1804 + else:

1805 + pixel_values = mm_kwargs["pixel_values_videos"]

1806 + return pixel_values

1807 +

1808 + def _get_grid_thw_by_modality(

1809 + self,

1810 + mm_kwargs: dict[str, Any],

1811 + ) -> list[tuple[int, int, int]]:

1812 + grid_thw_key = f"{self.get_input_modality(mm_kwargs)}_grid_thw"

1813 + grid_thw = mm_kwargs[grid_thw_key]

1814 + if not isinstance(grid_thw, list):

1815 + grid_thw = grid_thw.tolist()

1816 + return grid_thw

1817 +

1818 def get_encoder_cudagraph_num_items(

1819 self,

1820 mm_kwargs: dict[str, Any],

1821 ) -> int:

1769 - return len(mm_kwargs["image_grid_thw"])

1822 + return len(self._get_grid_thw_by_modality(mm_kwargs))

1823

1824 def get_encoder_cudagraph_per_item_output_tokens(

1825 self,

1826 mm_kwargs: dict[str, Any],

1827 ) -> list[int]:

1828 m = self.visual.spatial_merge_size

1776 - return [t * (h // m) * (w // m) for t, h, w in mm_kwargs["image_grid_thw"]]

1829 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1830 + return [t * (h // m) * (w // m) for t, h, w in grid_thw]

1831

1832 def get_encoder_cudagraph_per_item_input_sizes(

1833 self,

1834 mm_kwargs: dict[str, Any],

1835 ) -> list[int]:

1782 - return [t * h * w for t, h, w in mm_kwargs["image_grid_thw"]]

1836 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1837 + return [t * h * w for t, h, w in grid_thw]

1838

1839 def select_encoder_cudagraph_items(

1840 self,

1841 mm_kwargs: dict[str, Any],

1842 indices: list[int],

1843 ) -> dict[str, Any]:

1789 - grid_thw = mm_kwargs["image_grid_thw"]

1790 - pixel_values = mm_kwargs["pixel_values"]

1844 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1845 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1846

1847 if len(indices) == 0:

1793 - return {

1794 - "pixel_values": pixel_values[:0],

1795 - "image_grid_thw": [],

1796 - }

1848 + if self.get_input_modality(mm_kwargs) == "image":

1849 + return {

1850 + "pixel_values": pixel_values[:0],

1851 + "image_grid_thw": [],

1852 + }

1853 + else:

1854 + return {

1855 + "pixel_values_videos": pixel_values[:0],

1856 + "video_grid_thw": [],

1857 + }

1858

1859 # Compute cumulative patch offsets for slicing pixel_values

1860 patches_per_item = [t * h * w for t, h, w in grid_thw]

@@ -1806,15 +1867,22 @@ class Qwen3VLForConditionalGeneration(

1867 )

1868 selected_grid = [grid_thw[i] for i in indices]

1869

1809 - return {

1810 - "pixel_values": selected_pv,

1811 - "image_grid_thw": selected_grid,

1812 - }

1870 + if self.get_input_modality(mm_kwargs) == "image":

1871 + return {

1872 + "pixel_values": selected_pv,

1873 + "image_grid_thw": selected_grid,

1874 + }

1875 + else:

1876 + return {

1877 + "pixel_values_videos": selected_pv,

1878 + "video_grid_thw": selected_grid,

1879 + }

1880

1881 def prepare_encoder_cudagraph_capture_inputs(

1882 self,

1883 token_budget: int,

1884 max_batch_size: int,

1885 + max_frames_per_batch: int,

1886 device: torch.device,

1887 dtype: torch.dtype,

1888 ):

@@ -1823,14 +1891,35 @@ class Qwen3VLForConditionalGeneration(

1891 )

1892

1893 spatial_merge_size = self.visual.spatial_merge_size

1826 - per_image_output = token_budget // max_batch_size

1827 -

1828 - # Synthetic rectangular grid: [1, merge, per_image_output * merge]

1829 - # produces exactly per_image_output tokens per image.

1830 - grid_config = [

1831 - [1, spatial_merge_size, per_image_output * spatial_merge_size]

1832 - for _ in range(max_batch_size)

1833 - ]

1894 + per_mm_item_output = token_budget // max_batch_size

1895 +

1896 + frames_per_item = max_frames_per_batch // max_batch_size

1897 + if frames_per_item > 1:

1898 + # Build the capture grid using a video-format layout so that

1899 + # cu_seqlens is sized for video replays from the start.

1900 + # cu_seqlens has one entry per attention sequence (one per frame),

1901 + # so using T > 1 per item makes the buffer large enough without

1902 + # relying solely on padding.

1903 + # Ceiling ensures frames_per_item * tokens_per_frame >= per_mm_item_output

1904 + # so the pixel_values buffer covers any valid single-item replay.

1905 + tokens_per_frame = (

1906 + per_mm_item_output + frames_per_item - 1

1907 + ) // frames_per_item

1908 + # Video-format grid_config (T=frames_per_item).

1909 + grid_config = [

1910 + [

1911 + frames_per_item,

1912 + spatial_merge_size,

1913 + tokens_per_frame * spatial_merge_size,

1914 + ]

1915 + for _ in range(max_batch_size)

1916 + ]

1917 + else:

1918 + # Image-format grid_config (T=1).

1919 + grid_config = [

1920 + [1, spatial_merge_size, per_mm_item_output * spatial_merge_size]

1921 + for _ in range(max_batch_size)

1922 + ]

1923

1924 # Create dummy pixel_values

1925 patch_embed = self.visual.patch_embed

@@ -1848,15 +1937,18 @@ class Qwen3VLForConditionalGeneration(

1937 # Override max_seqlen with a safe upper bound for capture.

1938 # max_seqlen.item() gets baked into the CUDA graph (not replayed),

1939 # so the capture value must cover any replay scenario.

1851 - # Worst case: 1 image consuming the full budget ->

1940 + # Worst case: 1 item consuming the full budget ->

1941 # seq_len = token_budget * spatial_merge_size^2.

1942 buffers = self.visual.prepare_encoder_metadata(

1943 grid_config,

1944 max_batch_size=max_batch_size,

1945 + max_frames_per_batch=max_frames_per_batch,

1946 max_seqlen_override=token_budget * (spatial_merge_size**2),

1947 device=device,

1948 )

1949

1950 + # Just use image-modality dummy input_buffer for capturing, since it's also

1951 + # compatible for video inputs (has the same shape: [num_patches, C*T*P*P]).

1952 mm_kwargs = {

1953 "pixel_values": dummy_pixel_values,

1954 "image_grid_thw": grid_config,

@@ -1871,17 +1963,21 @@ class Qwen3VLForConditionalGeneration(

1963 self,

1964 mm_kwargs: dict[str, Any],

1965 max_batch_size: int,

1966 + max_frames_per_batch: int,

1967 ):

1875 - from vllm.v1.worker.encoder_cudagraph_defs import (

1876 - EncoderCudaGraphReplayBuffers,

1877 - )

1968 + modality = self.get_input_modality(mm_kwargs)

1969 + grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)

1970

1879 - grid_thw_list = mm_kwargs["image_grid_thw"]

1880 -

1881 - buffers = self.visual.prepare_encoder_metadata(

1882 - grid_thw_list,

1883 - max_batch_size=max_batch_size,

1884 - )

1971 + if modality == "image":

1972 + buffers = self.visual.prepare_encoder_metadata(

1973 + grid_thw_list,

1974 + max_batch_size=max_batch_size,

1975 + )

1976 + else:

1977 + buffers = self.visual.prepare_encoder_metadata(

1978 + grid_thw_list,

1979 + max_frames_per_batch=max_frames_per_batch,

1980 + )

1981

1982 return EncoderCudaGraphReplayBuffers(buffers=buffers)

1983

@@ -1890,16 +1986,16 @@ class Qwen3VLForConditionalGeneration(

1986 mm_kwargs: dict[str, Any],

1987 buffers: dict[str, torch.Tensor],

1988 ) -> torch.Tensor:

1893 - pixel_values = mm_kwargs["pixel_values"]

1894 - grid_thw = mm_kwargs["image_grid_thw"]

1989 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1990 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1991 return self.visual(pixel_values, grid_thw, encoder_metadata=buffers)

1992

1993 def encoder_eager_forward(

1994 self,

1995 mm_kwargs: dict[str, Any],

1996 ) -> torch.Tensor:

1901 - pixel_values = mm_kwargs["pixel_values"]

1902 - grid_thw = mm_kwargs["image_grid_thw"]

1997 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1998 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1999 return self.visual(pixel_values, grid_thw)

2000

2001 def _parse_and_validate_image_input(

vllm/v1/worker/encoder_cudagraph.py

+33 -11

@@ -36,6 +36,7 @@ class BudgetGraphMetadata:

37 token_budget: int

38 max_batch_size: int # Max number of images/videos per batch

39 + max_frames_per_batch: int # Max total frames per batch (for video)

40 graph: torch.cuda.CUDAGraph

41 # The input tensor updated before replay (e.g. pixel_values)

42 input_buffer: torch.Tensor

@@ -66,12 +67,13 @@ class EncoderCudaGraphManager:

68 comp_config = vllm_config.compilation_config

69 user_budgets = comp_config.encoder_cudagraph_token_budgets

69 - user_max_images = comp_config.encoder_cudagraph_max_images_per_batch

70 + user_max_mm_items = comp_config.encoder_cudagraph_max_vision_items_per_batch

71 + user_max_frames = comp_config.encoder_cudagraph_max_frames_per_batch

71 - if user_budgets and user_max_images > 0:

73 + if user_budgets and user_max_mm_items > 0:

74 # Fully user-specified

75 self.token_budgets = sorted(user_budgets)

74 - self.max_batch_size = user_max_images

76 + self.max_batch_size = user_max_mm_items

77 else:

78 # Auto-infer missing values from model

79 min_budget, max_budget = model.get_encoder_cudagraph_budget_range(

@@ -83,9 +85,15 @@ class EncoderCudaGraphManager:

85 else self._generate_budgets(min_budget, max_budget)

86 )

87 self.max_batch_size = (

86 - user_max_images if user_max_images > 0 else max_budget // min_budget

88 + user_max_mm_items if user_max_mm_items > 0 else max_budget // min_budget

89 )

91 + if user_max_frames > 0:

92 + self.max_frames_per_batch = user_max_frames

93 + else:

94 + # TODO(shen-shanshan): optimize this auto-infer for max_frames_per_batch.

95 + self.max_frames_per_batch = self.max_batch_size * 2

96 +

97 mm_config = vllm_config.model_config.multimodal_config

98 self.use_dp = (

99 mm_config is not None

@@ -100,9 +108,10 @@ class EncoderCudaGraphManager:

108

109 logger.info(

110 "EncoderCudaGraphManager initialized with "

103 - "budgets=%s, max_batch_size=%d, use_dp=%s",

111 + "budgets=%s, max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",

112 self.token_budgets,

113 self.max_batch_size,

114 + self.max_frames_per_batch if self.max_frames_per_batch > 0 else "auto",

115 self.use_dp,

116 )

117

@@ -136,13 +145,19 @@ class EncoderCudaGraphManager:

145 def _capture_budget_graph(self, token_budget: int):

146 """Capture CUDA graph for a single token budget."""

147 logger.debug(

139 - "Capturing encoder cudagraph for budget=%d, max_batch_size=%d",

148 + "Capturing encoder cudagraph for budget=%d, max_batch_size=%d, "

149 + "max_frames_per_batch=%d",

150 token_budget,

151 self.max_batch_size,

152 + self.max_frames_per_batch,

153 )

154

155 capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs(

145 - token_budget, self.max_batch_size, self.device, self.dtype

156 + token_budget,

157 + self.max_batch_size,

158 + self.max_frames_per_batch,

159 + self.device,

160 + self.dtype,

161 )

162

163 mm_kwargs = capture_inputs.mm_kwargs

@@ -157,10 +172,14 @@ class EncoderCudaGraphManager:

172 output = self.model.encoder_cudagraph_forward(mm_kwargs, buffers)

173 output_buffer.copy_(output)

174

160 - input_key = self.config.input_key

175 + # Since the image and video modalities share the same per-patch shape,

176 + # so we can use the image dummy inputs to capture CUDA graph for both

177 + # image and video.

178 + input_key = self.config.input_key_by_modality["image"]

179 self.budget_graphs[token_budget] = BudgetGraphMetadata(

180 token_budget=token_budget,

181 max_batch_size=self.max_batch_size,

182 + max_frames_per_batch=self.max_frames_per_batch,

183 graph=graph,

184 input_buffer=mm_kwargs[input_key],

185 metadata_buffers=buffers,

@@ -230,10 +249,11 @@ class EncoderCudaGraphManager:

249 # Copy the input tensor. Buffers are sized for the full budget;

250 # actual inputs may be smaller. Zero then slice-copy so padded

251 # positions are invisible to attention (cu_seqlens masks them out).

233 - input_key = self.config.input_key

252 + input_key = self.config.input_key_by_modality[

253 + self.model.get_input_modality(mm_kwargs)

254 + ]

255 src = mm_kwargs[input_key]

256 n = src.shape[0]

236 - graph_meta.input_buffer.zero_()

257 graph_meta.input_buffer[:n].copy_(src)

258

259 # Copy metadata buffers using keys from config.buffer_keys.

@@ -362,7 +382,9 @@ class EncoderCudaGraphManager:

382 (token_budget - batch_out_tokens) / token_budget * 100,

383 )

384 replay = self.model.prepare_encoder_cudagraph_replay_buffers(

365 - batch_mm_kwargs, self.max_batch_size

385 + batch_mm_kwargs,

386 + self.max_batch_size,

387 + self.max_frames_per_batch,

388 )

389

390 # graph_hits counted inside _run_budget_graph after replay.

vllm/v1/worker/encoder_cudagraph_defs.py

+4 -2

@@ -20,8 +20,10 @@ class EncoderCudaGraphConfig:

20 modalities: list[str]

21 """Supported modalities (e.g. ["image"])."""

23 - input_key: str

24 - """Key in mm_kwargs for the input tensor (e.g. "pixel_values")."""

23 + input_key_by_modality: dict[str, str]

24 + """Per-modality input tensor key mapping, e.g.

25 + {"image": "pixel_values", "video": "pixel_values_videos"}.

26 + """

28 buffer_keys: list[str]

29 """Keys for the tensor buffers recorded into the CUDA graph.