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 4, 08:45 UTC
[MM][CG] Support ViT CG for Qwen2.5-VL (#40830)
Signed-off-by: John Calderon <jcalderon@nvidia.com>
John Calderon committed
May 1, 2026 at 23:10 UTC
964a4bc2a57aca2a42d04538b27cab4d333d0f5d
5 files changed
+539
-22
docs/design/cuda_graphs_multimodal.md
+2
@@ -86,9 +86,11 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
86
| Architecture | Models | CG for Image | CG for Video |
87
| ------------ | ------ | ------------ | ------------ |
88
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
89
+| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
90
91
!!! note
92
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
93
+ For Qwen2.5-VL only FA2 and FA3 has been tested.
94
95
## Configuration
96
examples/generate/multimodal/vision_language_offline.py
+1
@@ -2466,6 +2466,7 @@ MODELS_NEED_VIDEO_METADATA = [
2466
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2467
"qwen3_vl",
2468
"qwen3_vl_moe",
2469
+ "qwen2_5_vl",
2470
]
2471
2472
tests/models/multimodal/generation/test_qwen2_5_vl.py
+95
@@ -3,6 +3,7 @@
3
4
import pytest
5
6
+from vllm.assets.image import ImageAsset
7
from vllm.multimodal.video import sample_frames_from_video
8
9
from ....conftest import VIDEO_ASSETS
@@ -11,6 +12,7 @@ models = ["Qwen/Qwen2.5-VL-3B-Instruct"]
12
target_dtype = "bfloat16"
13
14
VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
15
+IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
16
17
18
def qwen2_5_vl_chat_template(*query):
@@ -28,6 +30,25 @@ VIDEO_PROMPTS = VIDEO_ASSETS.prompts(
30
)
31
32
33
+WINDOW_ATTN_IMAGE_PROMPT = qwen2_5_vl_chat_template(
34
+ IMAGE_PLACEHOLDER,
35
+ "Describe the image.",
36
+)
37
+
38
+
39
+def _window_attention_regression_image():
40
+ # image from regression issue: https://github.com/vllm-project/vllm/issues/15122
41
+ image = ImageAsset("hato").pil_image
42
+ return image.resize((image.width // 2, image.height // 2))
43
+
44
+
45
+def _encoder_cudagraph_config(*, max_vision_items: int) -> dict:
46
+ return {
47
+ "cudagraph_mm_encoder": True,
48
+ "encoder_cudagraph_max_vision_items_per_batch": max_vision_items,
49
+ }
50
+
51
+
52
@pytest.mark.core_model
53
@pytest.mark.parametrize("model", models)
54
@pytest.mark.parametrize("video_pruning_rate", [0.0, 0.75])
@@ -146,3 +167,77 @@ def test_qwen2_5_vl_evs_batched_videos(
167
168
# Ensure the output is a string
169
assert isinstance(output_text, str)
170
+
171
+
172
+@pytest.mark.core_model
173
+@pytest.mark.parametrize("model", models)
174
+@pytest.mark.parametrize("dtype", [target_dtype])
175
+@pytest.mark.parametrize("max_tokens", [128])
176
+@pytest.mark.parametrize("use_bytecode_hook", [True, False])
177
+def test_qwen2_5_vl_window_attention_image(
178
+ vllm_runner,
179
+ model,
180
+ dtype: str,
181
+ max_tokens: int,
182
+ use_bytecode_hook: bool,
183
+ monkeypatch,
184
+) -> None:
185
+ """Regression test for Qwen2.5 window-attention image path."""
186
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
187
+
188
+ prompt = [WINDOW_ATTN_IMAGE_PROMPT]
189
+ images = [[_window_attention_regression_image()]]
190
+
191
+ with vllm_runner(
192
+ model,
193
+ runner="generate",
194
+ max_model_len=4096,
195
+ dtype=dtype,
196
+ limit_mm_per_prompt={"image": 1},
197
+ compilation_config=_encoder_cudagraph_config(max_vision_items=1),
198
+ ) as vllm_model:
199
+ outputs = vllm_model.generate_greedy(prompt, max_tokens, images=images)
200
+
201
+ assert len(outputs) == 1
202
+ output_ids, output_text = outputs[0]
203
+ assert len(output_ids) > 0
204
+ assert len(output_text) > 0
205
+ assert isinstance(output_text, str)
206
+
207
+
208
+@pytest.mark.core_model
209
+@pytest.mark.parametrize("model", models)
210
+@pytest.mark.parametrize("dtype", [target_dtype])
211
+@pytest.mark.parametrize("max_tokens", [128])
212
+@pytest.mark.parametrize("use_bytecode_hook", [True, False])
213
+def test_qwen2_5_vl_window_attention_image_batch(
214
+ vllm_runner,
215
+ model,
216
+ dtype: str,
217
+ max_tokens: int,
218
+ use_bytecode_hook: bool,
219
+ monkeypatch,
220
+) -> None:
221
+ """Regression test window-attention with a small image batch."""
222
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
223
+
224
+ image = _window_attention_regression_image()
225
+ prompts = [WINDOW_ATTN_IMAGE_PROMPT, WINDOW_ATTN_IMAGE_PROMPT]
226
+ images = [[image], [image]]
227
+
228
+ with vllm_runner(
229
+ model,
230
+ runner="generate",
231
+ max_model_len=4096,
232
+ max_num_seqs=2,
233
+ dtype=dtype,
234
+ limit_mm_per_prompt={"image": 1},
235
+ compilation_config=_encoder_cudagraph_config(max_vision_items=2),
236
+ ) as vllm_model:
237
+ outputs = vllm_model.generate_greedy(prompts, max_tokens, images=images)
238
+
239
+ assert len(outputs) == 2
240
+ for output_ids, output_text in outputs:
241
+ assert len(output_ids) > 0
242
+ assert len(output_text) > 0
243
+ assert isinstance(output_text, str)
tests/models/multimodal/generation/test_vit_cudagraph.py
+12
-1
@@ -54,7 +54,18 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
54
needs_video_metadata=True,
55
marks=[pytest.mark.core_model],
56
),
57
- # TODO: Add more models below.
57
+ "qwen2_5_vl": VitCudagraphTestConfig(
58
+ model="Qwen/Qwen2.5-VL-3B-Instruct",
59
+ image_prompt=qwen_vl_chat_template(
60
+ "<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
61
+ ),
62
+ video_prompt=qwen_vl_chat_template(
63
+ "<|vision_start|><|video_pad|><|vision_end|>"
64
+ "Describe this video in one sentence."
65
+ ),
66
+ needs_video_metadata=False,
67
+ marks=[pytest.mark.core_model],
68
+ ),
69
}
70
71
vllm/model_executor/models/qwen2_5_vl.py
+429
-21
@@ -85,11 +85,13 @@ from vllm.sequence import IntermediateTensors
85
from vllm.utils.platform_utils import is_pin_memory_available
86
from vllm.utils.tensor_schema import TensorSchema, TensorShape
87
from vllm.v1.attention.backends.registry import AttentionBackendEnum
88
+from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
89
90
from .interfaces import (
91
MultiModalEmbeddings,
92
SupportsEagle,
93
SupportsEagle3,
94
+ SupportsEncoderCudaGraph,
95
SupportsLoRA,
96
SupportsMRoPE,
97
SupportsMultiModal,
@@ -771,22 +773,54 @@ class Qwen2_5_VisionTransformer(nn.Module):
773
inv[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype)
774
return inv
775
774
- def forward(
776
+ def prepare_encoder_metadata(
777
self,
776
- x: torch.Tensor,
778
grid_thw: list[list[int]],
778
- ) -> torch.Tensor:
779
+ *,
780
+ max_batch_size: int | None = None,
781
+ max_frames_per_batch: int | None = None,
782
+ max_window_seqs_per_batch: int | None = None,
783
+ max_seqlen_override: int | None = None,
784
+ max_seqlen_window_override: int | None = None,
785
+ device: torch.device | None = None,
786
+ ) -> dict[str, torch.Tensor]:
787
+ """Compute encoder metadata from grid_thw.
788
+
789
+ Shared by the eager forward path, CUDA graph capture, and
790
+ CUDA graph replay to avoid duplicated implementation.
791
+
792
+ Args:
793
+ grid_thw: Grid configurations as list of [t, h, w].
794
+ max_batch_size: If set, pad cu_seqlens to this size
795
+ (needed for CUDA graph capture/replay).
796
+ max_frames_per_batch: If set, overrides max_batch_size for
797
+ cu_seqlens padding. For video inputs each item contributes
798
+ T attention sequences (frames); this sizes the buffer to
799
+ the total frame budget so video replays never overflow.
800
+ max_window_seqs_per_batch: If set, pad cu_window_seqlens to this
801
+ number of window sequences. This keeps cu_window_seqlens shape
802
+ stable across capture/replay for CUDA graph safety.
803
+ max_seqlen_override: If set, use this value for max_seqlen
804
+ instead of computing from cu_seqlens (needed for CUDA
805
+ graph capture to cover worst-case replay scenarios).
806
+ max_seqlen_window_override: If set, use this value for
807
+ window-attention max_seqlen instead of computing from
808
+ cu_window_seqlens (needed for CUDA graph capture to
809
+ cover worst-case replay scenarios).
810
+ device: Device to place tensors on. Defaults to self.device.
811
+ """
812
+
813
+ if device is None:
814
+ device = self.device
815
+ metadata: dict[str, torch.Tensor] = {}
816
+
817
# patchify
780
- seq_len, _ = x.size()
818
rotary_pos_emb_cos = []
819
rotary_pos_emb_sin = []
820
window_index: list = []
821
cu_window_seqlens: list = [torch.tensor([0], dtype=torch.int32)]
822
cu_seqlens: list = []
823
787
- hidden_states = x.to(device=self.device, dtype=self.dtype)
788
- hidden_states = self.patch_embed(hidden_states)
789
-
824
window_index_id = 0
825
cu_window_seqlens_last = 0
826
for t, h, w in grid_thw:
@@ -825,23 +859,99 @@ class Qwen2_5_VisionTransformer(nn.Module):
859
cu_seqlens = torch.cumsum(cu_seqlens, dim=0, dtype=torch.int32)
860
cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)
861
862
+ # Pad cu_seqlens to the required number of sequences.
863
+ # For videos each item contributes T frames = T attention sequences,
864
+ # so the total can exceed max_batch_size. max_frames_per_batch
865
+ # overrides the pad target when set.
866
+ pad_to = (
867
+ max_frames_per_batch if max_frames_per_batch is not None else max_batch_size
868
+ )
869
+ if pad_to is not None:
870
+ num_seqs = len(cu_seqlens) - 1
871
+ if num_seqs < pad_to:
872
+ cu_seqlens = torch.cat(
873
+ (
874
+ cu_seqlens,
875
+ torch.full(
876
+ (pad_to - num_seqs,),
877
+ cu_seqlens[-1],
878
+ dtype=cu_seqlens.dtype,
879
+ device=cu_seqlens.device,
880
+ ),
881
+ )
882
+ )
883
+
884
+ # Pad cu_window_seqlens to a stable number of window sequences.
885
+ # Like cu_seqlens, we repeat the last cumulative offset so padded
886
+ # entries represent empty sequences.
887
+ if max_window_seqs_per_batch is not None:
888
+ num_window_seqs = len(cu_window_seqlens) - 1
889
+ if num_window_seqs < max_window_seqs_per_batch:
890
+ cu_window_seqlens = torch.cat(
891
+ (
892
+ cu_window_seqlens,
893
+ torch.full(
894
+ (max_window_seqs_per_batch - num_window_seqs,),
895
+ cu_window_seqlens[-1],
896
+ dtype=cu_window_seqlens.dtype,
897
+ device=cu_window_seqlens.device,
898
+ ),
899
+ )
900
+ )
901
+
902
# transformers
903
# pre-compute seqlens for window/full attn to reduce cuMemcpy operations
830
- max_seqlen_full = self.compute_attn_mask_seqlen(cu_seqlens)
831
- max_seqlen_window = self.compute_attn_mask_seqlen(cu_window_seqlens)
904
+ if max_seqlen_override is None:
905
+ max_seqlen_full = self.compute_attn_mask_seqlen(cu_seqlens)
906
+ else:
907
+ max_seqlen_full = torch.tensor(max_seqlen_override, dtype=torch.int32)
908
+ if max_seqlen_window_override is None:
909
+ max_seqlen_window = self.compute_attn_mask_seqlen(cu_window_seqlens)
910
+ else:
911
+ max_seqlen_window = torch.tensor(
912
+ max_seqlen_window_override, dtype=torch.int32
913
+ )
914
833
- cu_seqlens = cu_seqlens.to(device=self.device, non_blocking=True)
834
- cu_window_seqlens = cu_window_seqlens.to(device=self.device, non_blocking=True)
835
- rotary_pos_emb_cos = rotary_pos_emb_cos.to(
836
- device=self.device, non_blocking=True
837
- )
838
- rotary_pos_emb_sin = rotary_pos_emb_sin.to(
839
- device=self.device, non_blocking=True
840
- )
841
- window_index = window_index.to(device=hidden_states.device, non_blocking=True)
842
- reverse_indices = reverse_indices.to(
843
- device=hidden_states.device, non_blocking=True
844
- )
915
+ cu_seqlens = cu_seqlens.to(device=device, non_blocking=True)
916
+ cu_window_seqlens = cu_window_seqlens.to(device=device, non_blocking=True)
917
+ rotary_pos_emb_cos = rotary_pos_emb_cos.to(device=device, non_blocking=True)
918
+ rotary_pos_emb_sin = rotary_pos_emb_sin.to(device=device, non_blocking=True)
919
+ window_index = window_index.to(device=device, non_blocking=True)
920
+ reverse_indices = reverse_indices.to(device=device, non_blocking=True)
921
+
922
+ metadata["rotary_pos_emb_cos"] = rotary_pos_emb_cos
923
+ metadata["rotary_pos_emb_sin"] = rotary_pos_emb_sin
924
+ metadata["window_index"] = window_index
925
+ metadata["reverse_indices"] = reverse_indices
926
+ metadata["cu_seqlens"] = cu_seqlens
927
+ metadata["cu_window_seqlens"] = cu_window_seqlens
928
+ metadata["max_seqlen_full"] = max_seqlen_full
929
+ metadata["max_seqlen_window"] = max_seqlen_window
930
+
931
+ return metadata
932
+
933
+ def forward(
934
+ self,
935
+ x: torch.Tensor,
936
+ grid_thw: list[list[int]],
937
+ *,
938
+ encoder_metadata: dict[str, torch.Tensor] | None = None,
939
+ ) -> torch.Tensor:
940
+ hidden_states = x.to(device=self.device, dtype=self.dtype)
941
+ hidden_states = self.patch_embed(hidden_states)
942
+
943
+ seq_len = hidden_states.shape[0]
944
+ if encoder_metadata is None:
945
+ encoder_metadata = self.prepare_encoder_metadata(grid_thw)
946
+
947
+ rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"]
948
+ rotary_pos_emb_sin = encoder_metadata["rotary_pos_emb_sin"]
949
+ window_index = encoder_metadata["window_index"]
950
+ reverse_indices = encoder_metadata["reverse_indices"]
951
+ cu_seqlens = encoder_metadata["cu_seqlens"]
952
+ cu_window_seqlens = encoder_metadata["cu_window_seqlens"]
953
+ max_seqlen_full = encoder_metadata["max_seqlen_full"]
954
+ max_seqlen_window = encoder_metadata["max_seqlen_window"]
955
956
hidden_states = hidden_states.reshape(
957
seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
@@ -1003,6 +1113,7 @@ class Qwen2_5_VLMultiModalProcessor(Qwen2VLMultiModalProcessor):
1113
class Qwen2_5_VLForConditionalGeneration(
1114
nn.Module,
1115
SupportsMultiModal,
1116
+ SupportsEncoderCudaGraph,
1117
SupportsLoRA,
1118
SupportsPP,
1119
SupportsQuant,
@@ -1124,6 +1235,7 @@ class Qwen2_5_VLForConditionalGeneration(
1235
1236
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
1237
self.config = config
1238
+ self.model_config = vllm_config.model_config
1239
self.vllm_config = vllm_config
1240
self.multimodal_config = multimodal_config
1241
self.video_pruning_rate = multimodal_config.video_pruning_rate
@@ -1447,6 +1559,302 @@ class Qwen2_5_VLForConditionalGeneration(
1559
multimodal_embeddings += tuple(video_embeddings)
1560
return multimodal_embeddings
1561
1562
+ # -- SupportsEncoderCudaGraph protocol methods --
1563
+
1564
+ def get_encoder_cudagraph_config(self):
1565
+ from vllm.v1.worker.encoder_cudagraph_defs import (
1566
+ EncoderCudaGraphConfig,
1567
+ )
1568
+
1569
+ # NOTE: With EVS pruning enabled, multimodal embeddings are post-processed
1570
+ # (append positions for image and prune+append positions for video) in
1571
+ # embed_multimodal(). The encoder CUDA graph path bypasses that postprocess
1572
+ # hook, so disable CUDA graph for all modalities to avoid inconsistent
1573
+ # embedding formats between eager and cudagraph paths.
1574
+ modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]
1575
+
1576
+ return EncoderCudaGraphConfig(
1577
+ modalities=modalities,
1578
+ input_key_by_modality={
1579
+ "image": "pixel_values",
1580
+ "video": "pixel_values_videos",
1581
+ },
1582
+ buffer_keys=[
1583
+ "rotary_pos_emb_cos",
1584
+ "rotary_pos_emb_sin",
1585
+ "window_index",
1586
+ "reverse_indices",
1587
+ "cu_seqlens",
1588
+ "cu_window_seqlens",
1589
+ "max_seqlen_full",
1590
+ "max_seqlen_window",
1591
+ ],
1592
+ out_hidden_size=self.visual.out_hidden_size,
1593
+ )
1594
+
1595
+ def get_input_modality(
1596
+ self,
1597
+ mm_kwargs: dict[str, Any],
1598
+ ) -> str:
1599
+ if "image_grid_thw" in mm_kwargs:
1600
+ return "image"
1601
+ return "video"
1602
+
1603
+ def get_max_frames_per_video(self) -> int:
1604
+ mm_registry = MULTIMODAL_REGISTRY
1605
+ info = mm_registry.get_processing_info(self.model_config)
1606
+ max_frames_per_video = info.get_num_frames_with_most_features(
1607
+ seq_len=self.model_config.max_model_len,
1608
+ mm_counts={"video": self.multimodal_config.get_limit_per_prompt("video")},
1609
+ )
1610
+ return max_frames_per_video
1611
+
1612
+ def get_encoder_cudagraph_budget_range(
1613
+ self,
1614
+ vllm_config: VllmConfig,
1615
+ ) -> tuple[int, int]:
1616
+ # Min: estimated smallest possible encoder input.
1617
+ # 224x224 image → 16x16 patches (patch_size=14)
1618
+ # spatial_merge_size=2 → 8x8 = 64 tokens
1619
+ min_budget = 64
1620
+ # Max: capped by max_num_batched_tokens
1621
+ max_budget = min(
1622
+ vllm_config.scheduler_config.max_num_batched_tokens,
1623
+ self.model_config.max_model_len,
1624
+ )
1625
+ return (min_budget, max_budget)
1626
+
1627
+ def _get_pixel_values_by_modality(
1628
+ self,
1629
+ mm_kwargs: dict[str, Any],
1630
+ ) -> torch.Tensor:
1631
+ if self.get_input_modality(mm_kwargs) == "image":
1632
+ pixel_values = mm_kwargs["pixel_values"]
1633
+ else:
1634
+ pixel_values = mm_kwargs["pixel_values_videos"]
1635
+ return pixel_values
1636
+
1637
+ def _get_grid_thw_by_modality(
1638
+ self,
1639
+ mm_kwargs: dict[str, Any],
1640
+ ) -> list[tuple[int, int, int]]:
1641
+ grid_thw_key = f"{self.get_input_modality(mm_kwargs)}_grid_thw"
1642
+ grid_thw = mm_kwargs[grid_thw_key]
1643
+ if not isinstance(grid_thw, list):
1644
+ grid_thw = grid_thw.tolist()
1645
+ return grid_thw
1646
+
1647
+ def get_encoder_cudagraph_num_items(
1648
+ self,
1649
+ mm_kwargs: dict[str, Any],
1650
+ ) -> int:
1651
+ return len(self._get_grid_thw_by_modality(mm_kwargs))
1652
+
1653
+ def get_encoder_cudagraph_per_item_output_tokens(
1654
+ self,
1655
+ mm_kwargs: dict[str, Any],
1656
+ ) -> list[int]:
1657
+ m = self.visual.spatial_merge_size
1658
+ grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
1659
+ return [t * (h // m) * (w // m) for t, h, w in grid_thw]
1660
+
1661
+ def get_encoder_cudagraph_per_item_input_sizes(
1662
+ self,
1663
+ mm_kwargs: dict[str, Any],
1664
+ ) -> list[int]:
1665
+ grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
1666
+ return [t * h * w for t, h, w in grid_thw]
1667
+
1668
+ def select_encoder_cudagraph_items(
1669
+ self,
1670
+ mm_kwargs: dict[str, Any],
1671
+ indices: list[int],
1672
+ ) -> dict[str, Any]:
1673
+ grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
1674
+ pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
1675
+
1676
+ if len(indices) == 0:
1677
+ if self.get_input_modality(mm_kwargs) == "image":
1678
+ return {
1679
+ "pixel_values": pixel_values[:0],
1680
+ "image_grid_thw": [],
1681
+ }
1682
+ else:
1683
+ return {
1684
+ "pixel_values_videos": pixel_values[:0],
1685
+ "video_grid_thw": [],
1686
+ }
1687
+
1688
+ # Compute cumulative patch offsets for slicing pixel_values
1689
+ patches_per_item = [t * h * w for t, h, w in grid_thw]
1690
+ cum_patches = [0]
1691
+ for p in patches_per_item:
1692
+ cum_patches.append(cum_patches[-1] + p)
1693
+
1694
+ selected_pv = torch.cat(
1695
+ [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
1696
+ )
1697
+ selected_grid = [grid_thw[i] for i in indices]
1698
+
1699
+ if self.get_input_modality(mm_kwargs) == "image":
1700
+ return {
1701
+ "pixel_values": selected_pv,
1702
+ "image_grid_thw": selected_grid,
1703
+ }
1704
+ else:
1705
+ return {
1706
+ "pixel_values_videos": selected_pv,
1707
+ "video_grid_thw": selected_grid,
1708
+ }
1709
+
1710
+ def prepare_encoder_cudagraph_capture_inputs(
1711
+ self,
1712
+ token_budget: int,
1713
+ max_batch_size: int,
1714
+ max_frames_per_batch: int,
1715
+ device: torch.device,
1716
+ dtype: torch.dtype,
1717
+ ):
1718
+ from vllm.v1.worker.encoder_cudagraph_defs import (
1719
+ EncoderCudaGraphCaptureInputs,
1720
+ )
1721
+
1722
+ spatial_merge_size = self.visual.spatial_merge_size
1723
+ max_window_seqs_per_batch = min(
1724
+ self.vllm_config.scheduler_config.max_num_batched_tokens,
1725
+ self.model_config.max_model_len,
1726
+ )
1727
+ # Use ceil here (not floor) so total captured capacity is never smaller
1728
+ # than token_budget when token_budget is not divisible by max_batch_size
1729
+ # (e.g., 324 budget with max_batch_size=8). Floor under-allocates
1730
+ # input_buffer and can fail replay copy for valid single-item batches.
1731
+ per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size
1732
+
1733
+ frames_per_item = max_frames_per_batch // max_batch_size
1734
+ if frames_per_item > 1:
1735
+ # Build the capture grid using a video-format layout so that
1736
+ # cu_seqlens is sized for video replays from the start.
1737
+ # cu_seqlens has one entry per attention sequence (one per frame),
1738
+ # so using T > 1 per item makes the buffer large enough without
1739
+ # relying solely on padding.
1740
+ # Ceiling ensures frames_per_item * tokens_per_frame >= per_mm_item_output
1741
+ # so the pixel_values buffer covers any valid single-item replay.
1742
+ tokens_per_frame = (
1743
+ per_mm_item_output + frames_per_item - 1
1744
+ ) // frames_per_item
1745
+ # Video-format grid_config (T=frames_per_item).
1746
+ grid_config = [
1747
+ [
1748
+ frames_per_item,
1749
+ spatial_merge_size,
1750
+ tokens_per_frame * spatial_merge_size,
1751
+ ]
1752
+ for _ in range(max_batch_size)
1753
+ ]
1754
+ else:
1755
+ # Image-format grid_config (T=1).
1756
+ grid_config = [
1757
+ [1, spatial_merge_size, per_mm_item_output * spatial_merge_size]
1758
+ for _ in range(max_batch_size)
1759
+ ]
1760
+
1761
+ # Create dummy pixel_values
1762
+ patch_embed = self.visual.patch_embed
1763
+ in_channels = patch_embed.proj.in_channels
1764
+ patch_size = patch_embed.patch_size
1765
+ temporal_patch_size = patch_embed.temporal_patch_size
1766
+ total_patches = sum(t * h * w for t, h, w in grid_config)
1767
+ flattened_patch_size = (
1768
+ in_channels * temporal_patch_size * patch_size * patch_size
1769
+ )
1770
+ dummy_pixel_values = torch.randn(
1771
+ total_patches, flattened_patch_size, device=device, dtype=dtype
1772
+ )
1773
+
1774
+ # Override max_seqlen with a safe upper bound for capture.
1775
+ # max_seqlen.item() gets baked into the CUDA graph (not replayed),
1776
+ # so the capture value must cover any replay scenario.
1777
+ # Worst case: 1 item consuming the full budget ->
1778
+ # seq_len = token_budget * spatial_merge_size^2.
1779
+ # For window-attention, each local window is bounded by fixed geometry:
1780
+ # (window_size / patch_size / spatial_merge_size)^2 windows in merged
1781
+ # token space, multiplied by spatial_merge_size^2 to map back to the
1782
+ # unmerged sequence length used by attention kernels.
1783
+ vit_merger_window_size = (
1784
+ self.visual.window_size
1785
+ // self.visual.spatial_merge_size
1786
+ // self.visual.patch_size
1787
+ )
1788
+ max_seqlen_window_override = vit_merger_window_size**2 * (spatial_merge_size**2)
1789
+ buffers = self.visual.prepare_encoder_metadata(
1790
+ grid_config,
1791
+ max_batch_size=max_batch_size,
1792
+ max_frames_per_batch=max_frames_per_batch,
1793
+ max_window_seqs_per_batch=max_window_seqs_per_batch,
1794
+ max_seqlen_override=token_budget * (spatial_merge_size**2),
1795
+ max_seqlen_window_override=max_seqlen_window_override,
1796
+ device=device,
1797
+ )
1798
+
1799
+ # Just use image-modality dummy input_buffer for capturing, since it's also
1800
+ # compatible for video inputs (has the same shape: [num_patches, C*T*P*P]).
1801
+ mm_kwargs = {
1802
+ "pixel_values": dummy_pixel_values,
1803
+ "image_grid_thw": grid_config,
1804
+ }
1805
+
1806
+ return EncoderCudaGraphCaptureInputs(
1807
+ mm_kwargs=mm_kwargs,
1808
+ buffers=buffers,
1809
+ )
1810
+
1811
+ def prepare_encoder_cudagraph_replay_buffers(
1812
+ self,
1813
+ mm_kwargs: dict[str, Any],
1814
+ max_batch_size: int,
1815
+ max_frames_per_batch: int,
1816
+ ):
1817
+ modality = self.get_input_modality(mm_kwargs)
1818
+ grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
1819
+
1820
+ if modality == "image":
1821
+ buffers = self.visual.prepare_encoder_metadata(
1822
+ grid_thw_list,
1823
+ max_batch_size=max_batch_size,
1824
+ max_window_seqs_per_batch=min(
1825
+ self.vllm_config.scheduler_config.max_num_batched_tokens,
1826
+ self.model_config.max_model_len,
1827
+ ),
1828
+ )
1829
+ else:
1830
+ buffers = self.visual.prepare_encoder_metadata(
1831
+ grid_thw_list,
1832
+ max_frames_per_batch=max_frames_per_batch,
1833
+ max_window_seqs_per_batch=min(
1834
+ self.vllm_config.scheduler_config.max_num_batched_tokens,
1835
+ self.model_config.max_model_len,
1836
+ ),
1837
+ )
1838
+
1839
+ return EncoderCudaGraphReplayBuffers(buffers=buffers)
1840
+
1841
+ def encoder_cudagraph_forward(
1842
+ self,
1843
+ mm_kwargs: dict[str, Any],
1844
+ buffers: dict[str, torch.Tensor],
1845
+ ) -> torch.Tensor:
1846
+ pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
1847
+ grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
1848
+ return self.visual(pixel_values, grid_thw, encoder_metadata=buffers)
1849
+
1850
+ def encoder_eager_forward(
1851
+ self,
1852
+ mm_kwargs: dict[str, Any],
1853
+ ) -> torch.Tensor:
1854
+ pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
1855
+ grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
1856
+ return self.visual(pixel_values, grid_thw)
1857
+
1858
def forward(
1859
self,
1860
input_ids: torch.Tensor | None,