[MM][CG] Support ViT full CUDA graph for Idefics3 and SmolVLM (#47625)
Signed-off-by: Ryouichi Kurihara <ryouichi.kurihara.628@gmail.com>
Chipmunk committed
Aug 29, 2026 at 02:16 UTC
ffe690eca95636e57de9e683445c6c0e3f6e3fa9
3 files changed
+240
-33
tests/models/multimodal/generation/test_vit_cudagraph.py
+14
@@ -193,6 +193,20 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
193
vllm_runner_kwargs={"trust_remote_code": True},
194
marks=[pytest.mark.core_model],
195
),
196
+ "idefics3": VitCudagraphTestConfig(
197
+ model="HuggingFaceTB/SmolVLM-256M-Instruct",
198
+ modalities=["image"],
199
+ image_prompt=(
200
+ "<|begin_of_text|>User:<image>What is in this image?"
201
+ "<end_of_utterance>\nAssistant:"
202
+ ),
203
+ max_model_len=4096,
204
+ compilation_config_overrides={
205
+ "encoder_cudagraph_token_budgets": [4096],
206
+ },
207
+ vllm_runner_kwargs={"gpu_memory_utilization": 0.80},
208
+ marks=[pytest.mark.core_model],
209
+ ),
210
"step3_vl": VitCudagraphTestConfig(
211
model="stepfun-ai/Step3-VL-10B",
212
modalities=["image"],
vllm/model_executor/models/idefics2_vision_model.py
+19
-13
@@ -75,20 +75,12 @@ class Idefics2VisionEmbeddings(nn.Module):
75
self.num_positions = self.num_patches
76
self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
77
78
- def forward(
78
+ def get_position_ids(
79
self,
80
- pixel_values: torch.FloatTensor,
80
patch_attention_mask: torch.BoolTensor,
81
tgt_sizes: torch.IntTensor | None = None,
82
) -> torch.Tensor:
84
- batch_size, _, max_im_h, max_im_w = pixel_values.shape
85
- target_dtype = self.patch_embedding.weight.dtype
86
- patch_embeds = self.patch_embedding(pixel_values.to(target_dtype))
87
- embeddings = patch_embeds.flatten(2).transpose(1, 2)
88
- max_nb_patches_h, max_nb_patches_w = (
89
- max_im_h // self.patch_size,
90
- max_im_w // self.patch_size,
91
- )
83
+ batch_size, max_nb_patches_h, max_nb_patches_w = patch_attention_mask.shape
84
boundaries = torch.arange(
85
1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side
86
)
@@ -119,9 +111,21 @@ class Idefics2VisionEmbeddings(nn.Module):
111
bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w
112
).flatten()
113
position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids
122
- position_ids = position_ids.to(
123
- self.position_embedding.weight.device, non_blocking=True
124
- )
114
+
115
+ return position_ids.to(self.position_embedding.weight.device, non_blocking=True)
116
+
117
+ def forward(
118
+ self,
119
+ pixel_values: torch.FloatTensor,
120
+ patch_attention_mask: torch.BoolTensor,
121
+ tgt_sizes: torch.IntTensor | None = None,
122
+ position_ids: torch.Tensor | None = None,
123
+ ) -> torch.Tensor:
124
+ target_dtype = self.patch_embedding.weight.dtype
125
+ patch_embeds = self.patch_embedding(pixel_values.to(target_dtype))
126
+ embeddings = patch_embeds.flatten(2).transpose(1, 2)
127
+ if position_ids is None:
128
+ position_ids = self.get_position_ids(patch_attention_mask, tgt_sizes)
129
embeddings += self.position_embedding(position_ids)
130
return embeddings
131
@@ -423,6 +427,7 @@ class Idefics2VisionTransformer(nn.Module):
427
pixel_values,
428
patch_attention_mask: torch.BoolTensor | None = None,
429
tgt_sizes: torch.IntTensor | None = None,
430
+ position_ids: torch.Tensor | None = None,
431
) -> torch.Tensor:
432
batch_size = pixel_values.size(0)
433
@@ -446,6 +451,7 @@ class Idefics2VisionTransformer(nn.Module):
451
pixel_values=pixel_values,
452
patch_attention_mask=patch_attention_mask,
453
tgt_sizes=tgt_sizes,
454
+ position_ids=position_ids,
455
)
456
457
# Align with HuggingFace NaViT SigLIP in MiniCPMV/O:
vllm/model_executor/models/idefics3.py
+207
-20
@@ -17,7 +17,7 @@
17
"""Inference-only Idefics3 model compatible with HuggingFace weights."""
18
19
from collections.abc import Iterable, Mapping, Sequence
20
-from typing import Annotated, Literal, TypeAlias
20
+from typing import Annotated, Any, Literal, TypeAlias
21
22
import torch
23
from torch import nn
@@ -60,6 +60,7 @@ from .idefics2_vision_model import (
60
)
61
from .interfaces import (
62
MultiModalEmbeddings,
63
+ SupportsEncoderCudaGraph,
64
SupportsLoRA,
65
SupportsMultiModal,
66
)
@@ -442,6 +443,9 @@ class Idefics3Model(nn.Module):
443
self,
444
pixel_values: torch.Tensor,
445
pixel_attention_mask: torch.Tensor,
446
+ *,
447
+ remove_padding: bool = True,
448
+ position_ids: torch.Tensor | None = None,
449
) -> torch.Tensor:
450
# NOTE: we skip the step to select the vision feature layer since
451
# this is already done inside the vision tower
@@ -449,35 +453,52 @@ class Idefics3Model(nn.Module):
453
dtype=self.vision_model.embeddings.patch_embedding.weight.dtype
454
) # fp16 compatibility
455
452
- # Remove padding images - padding images are full 0.
453
- nb_values_per_image = pixel_values.shape[1:].numel()
454
- real_images_inds = (pixel_values == 0.0).sum(
455
- dim=(-1, -2, -3)
456
- ) != nb_values_per_image
457
- with gpu_sync_allowed():
458
- pixel_values = pixel_values[real_images_inds].contiguous()
456
+ if remove_padding:
457
+ # Remove padding images - padding images are full 0.
458
+ nb_values_per_image = pixel_values.shape[1:].numel()
459
+ real_images_inds = (pixel_values == 0.0).sum(
460
+ dim=(-1, -2, -3)
461
+ ) != nb_values_per_image
462
+ with gpu_sync_allowed():
463
+ pixel_values = pixel_values[real_images_inds].contiguous()
464
460
- # Handle the vision attention mask
461
- # Remove padding images from the mask
462
- pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous()
465
+ # Handle the vision attention mask
466
+ # Remove padding images from the mask
467
+ pixel_attention_mask = pixel_attention_mask[
468
+ real_images_inds
469
+ ].contiguous()
470
464
- patch_size = self.config.vision_config.patch_size
465
- patches_subgrid = pixel_attention_mask.unfold(
466
- dimension=1, size=patch_size, step=patch_size
467
- )
468
- patches_subgrid = patches_subgrid.unfold(
469
- dimension=2, size=patch_size, step=patch_size
470
- )
471
- patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool()
471
+ patch_attention_mask = self.get_patch_attention_mask(pixel_attention_mask)
472
473
# Get sequence from the vision encoder
474
image_hidden_states = self.vision_model(
475
pixel_values=pixel_values,
476
patch_attention_mask=patch_attention_mask,
477
+ position_ids=position_ids,
478
)
479
480
return image_hidden_states
481
482
+ def get_patch_attention_mask(
483
+ self,
484
+ pixel_attention_mask: torch.Tensor,
485
+ ) -> torch.Tensor:
486
+ patch_size = self.config.vision_config.patch_size
487
+ patches_subgrid = pixel_attention_mask.unfold(
488
+ dimension=1, size=patch_size, step=patch_size
489
+ )
490
+ patches_subgrid = patches_subgrid.unfold(
491
+ dimension=2, size=patch_size, step=patch_size
492
+ )
493
+ return (patches_subgrid.sum(dim=(-1, -2)) > 0).bool()
494
+
495
+ def get_position_ids(
496
+ self,
497
+ pixel_attention_mask: torch.Tensor,
498
+ ) -> torch.Tensor:
499
+ patch_attention_mask = self.get_patch_attention_mask(pixel_attention_mask)
500
+ return self.vision_model.embeddings.get_position_ids(patch_attention_mask)
501
+
502
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
503
return self.text_model.embed_input_ids(input_ids)
504
@@ -502,7 +523,11 @@ class Idefics3Model(nn.Module):
523
info=Idefics3ProcessingInfo,
524
dummy_inputs=Idefics3DummyInputsBuilder,
525
)
505
-class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA):
526
+class Idefics3ForConditionalGeneration(
527
+ nn.Module, SupportsMultiModal, SupportsLoRA, SupportsEncoderCudaGraph
528
+):
529
+ supports_encoder_cudagraph = True
530
+
531
packed_modules_mapping = {
532
"qkv_proj": [
533
"q_proj",
@@ -606,6 +631,168 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo
631
num_patches = image_input["num_patches"]
632
return [e.flatten(0, 1) for e in image_features.split(num_patches.tolist())]
633
634
+ # -- SupportsEncoderCudaGraph protocol methods --
635
+
636
+ def get_encoder_cudagraph_config(self):
637
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
638
+
639
+ return EncoderCudaGraphConfig(
640
+ modalities=["image"],
641
+ buffer_keys=["pixel_values", "pixel_attention_mask", "position_ids"],
642
+ out_hidden_size=self.config.text_config.hidden_size,
643
+ )
644
+
645
+ def get_encoder_cudagraph_budget_range(
646
+ self,
647
+ vllm_config: VllmConfig,
648
+ ) -> tuple[int, int]:
649
+ min_budget = self.model.image_seq_len
650
+ max_budget = min(
651
+ vllm_config.scheduler_config.max_num_batched_tokens,
652
+ vllm_config.model_config.max_model_len,
653
+ )
654
+ return (min_budget, max_budget)
655
+
656
+ def _get_num_patches_list(self, mm_kwargs: dict[str, Any]) -> list[int]:
657
+ num_patches = mm_kwargs["num_patches"]
658
+ if isinstance(num_patches, torch.Tensor):
659
+ return num_patches.tolist()
660
+ return [int(n) for n in num_patches]
661
+
662
+ def get_encoder_cudagraph_item_specs(
663
+ self,
664
+ mm_kwargs: dict[str, Any],
665
+ ):
666
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
667
+
668
+ return [
669
+ EncoderItemSpec(
670
+ input_size=num_patches,
671
+ output_tokens=num_patches * self.model.image_seq_len,
672
+ )
673
+ for num_patches in self._get_num_patches_list(mm_kwargs)
674
+ ]
675
+
676
+ def select_encoder_cudagraph_items(
677
+ self,
678
+ mm_kwargs: dict[str, Any],
679
+ indices: list[int],
680
+ ) -> dict[str, Any]:
681
+ pixel_values = mm_kwargs["pixel_values"]
682
+ pixel_attention_mask = mm_kwargs["pixel_attention_mask"]
683
+ num_patches_list = self._get_num_patches_list(mm_kwargs)
684
+
685
+ if len(indices) == 0:
686
+ return {
687
+ "pixel_values": pixel_values[:0],
688
+ "pixel_attention_mask": pixel_attention_mask[:0],
689
+ "num_patches": torch.empty(
690
+ (0,), dtype=torch.long, device=pixel_values.device
691
+ ),
692
+ }
693
+
694
+ cum_patches = [0]
695
+ for num_patches in num_patches_list:
696
+ cum_patches.append(cum_patches[-1] + num_patches)
697
+
698
+ selected_pixel_values = torch.cat(
699
+ [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
700
+ )
701
+ selected_attention_mask = torch.cat(
702
+ [pixel_attention_mask[cum_patches[i] : cum_patches[i + 1]] for i in indices]
703
+ )
704
+ selected_num_patches = torch.tensor(
705
+ [num_patches_list[i] for i in indices],
706
+ dtype=torch.long,
707
+ device=pixel_values.device,
708
+ )
709
+
710
+ return {
711
+ "pixel_values": selected_pixel_values,
712
+ "pixel_attention_mask": selected_attention_mask,
713
+ "num_patches": selected_num_patches,
714
+ }
715
+
716
+ def prepare_encoder_cudagraph_capture_inputs(
717
+ self,
718
+ token_budget: int,
719
+ max_batch_size: int,
720
+ max_frames_per_batch: int,
721
+ device: torch.device,
722
+ dtype: torch.dtype,
723
+ path: str = "default",
724
+ ):
725
+ from vllm.v1.worker.encoder_cudagraph_defs import (
726
+ EncoderCudaGraphCaptureInputs,
727
+ )
728
+
729
+ image_size = self.config.vision_config.image_size
730
+ num_tiles = max(
731
+ (token_budget + self.model.image_seq_len - 1) // self.model.image_seq_len,
732
+ 1,
733
+ )
734
+ dummy_pixel_values = torch.randn(
735
+ num_tiles, 3, image_size, image_size, device=device, dtype=dtype
736
+ )
737
+ dummy_pixel_attention_mask = torch.ones(
738
+ num_tiles, image_size, image_size, device=device, dtype=torch.bool
739
+ )
740
+ dummy_position_ids = self.model.get_position_ids(dummy_pixel_attention_mask)
741
+
742
+ return EncoderCudaGraphCaptureInputs(
743
+ values={
744
+ "pixel_values": dummy_pixel_values,
745
+ "pixel_attention_mask": dummy_pixel_attention_mask,
746
+ "position_ids": dummy_position_ids,
747
+ }
748
+ )
749
+
750
+ def prepare_encoder_cudagraph_replay_buffers(
751
+ self,
752
+ mm_kwargs: dict[str, Any],
753
+ max_batch_size: int,
754
+ max_frames_per_batch: int,
755
+ path: str = "default",
756
+ ):
757
+ from vllm.v1.worker.encoder_cudagraph_defs import (
758
+ EncoderCudaGraphReplayBuffers,
759
+ )
760
+
761
+ return EncoderCudaGraphReplayBuffers(
762
+ values={
763
+ "pixel_values": mm_kwargs["pixel_values"],
764
+ "pixel_attention_mask": mm_kwargs["pixel_attention_mask"],
765
+ "position_ids": self.model.get_position_ids(
766
+ mm_kwargs["pixel_attention_mask"]
767
+ ),
768
+ }
769
+ )
770
+
771
+ def encoder_cudagraph_forward(
772
+ self,
773
+ inputs: dict[str, torch.Tensor],
774
+ path: str = "default",
775
+ ) -> torch.Tensor:
776
+ image_features = self.model.image_pixels_to_features(
777
+ inputs["pixel_values"],
778
+ pixel_attention_mask=inputs["pixel_attention_mask"],
779
+ remove_padding=False,
780
+ position_ids=inputs["position_ids"],
781
+ )
782
+ image_features = self.model.connector(image_features)
783
+ return image_features.flatten(0, 1)
784
+
785
+ def encoder_eager_forward(
786
+ self,
787
+ mm_kwargs: dict[str, Any],
788
+ path: str = "default",
789
+ ) -> torch.Tensor:
790
+ image_input = self._parse_and_validate_image_input(**mm_kwargs)
791
+ assert image_input is not None and image_input["type"] == "pixel_values"
792
+ image_features = self._process_image_pixels(image_input)
793
+ image_features = self.model.connector(image_features)
794
+ return image_features.flatten(0, 1)
795
+
796
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
797
image_input = self._parse_and_validate_image_input(**kwargs)
798
if image_input is None: