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 1, 23:45 UTC
[MM][Perf][CG] Support ViT full CUDA graph for InternVL (#41759)
Signed-off-by: oguz <oguzhankir17@gmail.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Oğuzhan KIR committed
Jun 4, 2026 at 05:24 UTC
f25952e59b4a9e21a62d169212f3b072790a409d
4 files changed
+183
-2
docs/design/cuda_graphs_multimodal.md
+1
@@ -82,6 +82,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
82
83
| Architecture | Models | CG for Image | CG for Video |
84
| ------------ | ------ | ------------ | ------------ |
85
+| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
86
| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
87
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
88
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
examples/generate/multimodal/vision_language_offline.py
+1
@@ -2554,6 +2554,7 @@ MODELS_NEED_VIDEO_METADATA = [
2554
2555
2556
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2557
+ "internvl_chat",
2558
"qwen2_5_vl",
2559
"qwen3_vl",
2560
"qwen3_vl_moe",
tests/models/multimodal/generation/test_vit_cudagraph.py
+15
@@ -43,6 +43,10 @@ def qwen_vl_chat_template(content: str) -> str:
43
return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
44
45
46
+def internvl_chat_template(content: str) -> str:
47
+ return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
48
+
49
+
50
def step3_vl_chat_template(content: str) -> str:
51
return (
52
"<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "
@@ -51,6 +55,17 @@ def step3_vl_chat_template(content: str) -> str:
55
56
57
MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
58
+ "internvl": VitCudagraphTestConfig(
59
+ model="OpenGVLab/InternVL3-1B",
60
+ num_video_frames=8,
61
+ image_prompt=internvl_chat_template("<image>\nWhat is in this image?"),
62
+ video_prompt=internvl_chat_template(
63
+ "<video>\nDescribe this video in one sentence."
64
+ ),
65
+ needs_video_metadata=False,
66
+ vllm_runner_kwargs={"trust_remote_code": True},
67
+ marks=[pytest.mark.core_model],
68
+ ),
69
"qwen2_5_vl": VitCudagraphTestConfig(
70
model="Qwen/Qwen2.5-VL-3B-Instruct",
71
image_prompt=qwen_vl_chat_template(
vllm/model_executor/models/internvl.py
+166
-2
@@ -10,7 +10,7 @@
10
from abc import abstractmethod
11
from collections.abc import Iterable, Mapping, Sequence
12
from functools import cached_property
13
-from typing import Annotated, Literal, TypeAlias, TypeVar
13
+from typing import Annotated, Any, Literal, TypeAlias, TypeVar
14
15
import torch
16
import torch.nn as nn
@@ -55,6 +55,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
56
from .interfaces import (
57
MultiModalEmbeddings,
58
+ SupportsEncoderCudaGraph,
59
SupportsLoRA,
60
SupportsMultiModal,
61
SupportsPP,
@@ -543,7 +544,13 @@ class InternVLMultiModalProcessor(
544
info=InternVLProcessingInfo,
545
dummy_inputs=InternVLDummyInputsBuilder,
546
)
546
-class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
547
+class InternVLChatModel(
548
+ nn.Module,
549
+ SupportsMultiModal,
550
+ SupportsPP,
551
+ SupportsLoRA,
552
+ SupportsEncoderCudaGraph,
553
+):
554
supports_encoder_tp_data = True
555
556
@classmethod
@@ -924,3 +931,160 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)
931
932
num_patches = num_vision_tokens // (self.patch_tokens + 1)
933
return num_patches * self.num_image_token
934
+
935
+ # -- SupportsEncoderCudaGraph protocol methods --
936
+
937
+ def get_encoder_cudagraph_config(self):
938
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
939
+
940
+ return EncoderCudaGraphConfig(
941
+ modalities=["image", "video"],
942
+ # InternVision uses standard ViT attention (no rotary embeddings,
943
+ # no variable-length sequence metadata), so the only graph-recorded
944
+ # buffer is pixel_values_flat itself.
945
+ buffer_keys=["pixel_values_flat"],
946
+ out_hidden_size=self.config.text_config.hidden_size,
947
+ )
948
+
949
+ def get_input_modality(
950
+ self,
951
+ mm_kwargs: dict[str, Any],
952
+ ) -> str:
953
+ if "pixel_values_flat" in mm_kwargs:
954
+ return "image"
955
+ return "video"
956
+
957
+ def get_encoder_cudagraph_budget_range(
958
+ self,
959
+ vllm_config: "VllmConfig",
960
+ ) -> tuple[int, int]:
961
+ # Min: 1 tile → num_image_token output tokens.
962
+ min_budget = self.num_image_token
963
+ max_budget = min(
964
+ vllm_config.scheduler_config.max_num_batched_tokens,
965
+ vllm_config.model_config.max_model_len,
966
+ )
967
+ return (min_budget, max_budget)
968
+
969
+ def _get_internvl_patches_list(
970
+ self,
971
+ mm_kwargs: dict[str, Any],
972
+ ) -> list[int]:
973
+ """Return per-item tile counts as a plain list of ints."""
974
+ if self.get_input_modality(mm_kwargs) == "image":
975
+ patches = mm_kwargs.get("image_num_patches", [])
976
+ else:
977
+ patches = mm_kwargs.get("video_num_patches", [])
978
+ if isinstance(patches, torch.Tensor):
979
+ return patches.tolist()
980
+ return [int(n) for n in patches]
981
+
982
+ def get_encoder_cudagraph_item_specs(
983
+ self,
984
+ mm_kwargs: dict[str, Any],
985
+ ):
986
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
987
+
988
+ return [
989
+ EncoderItemSpec(
990
+ input_size=n,
991
+ output_tokens=n * self.num_image_token,
992
+ )
993
+ for n in self._get_internvl_patches_list(mm_kwargs)
994
+ ]
995
+
996
+ def select_encoder_cudagraph_items(
997
+ self,
998
+ mm_kwargs: dict[str, Any],
999
+ indices: list[int],
1000
+ ) -> dict[str, Any]:
1001
+ modality = self.get_input_modality(mm_kwargs)
1002
+ pv_key = (
1003
+ "pixel_values_flat" if modality == "image" else "pixel_values_flat_video"
1004
+ )
1005
+ patches_key = (
1006
+ "image_num_patches" if modality == "image" else "video_num_patches"
1007
+ )
1008
+
1009
+ pixel_values = mm_kwargs[pv_key]
1010
+ patches_list = self._get_internvl_patches_list(mm_kwargs)
1011
+
1012
+ if len(indices) == 0:
1013
+ return {pv_key: pixel_values[:0], patches_key: []}
1014
+
1015
+ # Compute cumulative tile offsets for slicing pixel_values.
1016
+ cum_patches = [0]
1017
+ for n in patches_list:
1018
+ cum_patches.append(cum_patches[-1] + n)
1019
+
1020
+ selected_pv = torch.cat(
1021
+ [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
1022
+ )
1023
+ selected_patches = [patches_list[i] for i in indices]
1024
+
1025
+ return {pv_key: selected_pv, patches_key: selected_patches}
1026
+
1027
+ def prepare_encoder_cudagraph_capture_inputs(
1028
+ self,
1029
+ token_budget: int,
1030
+ max_batch_size: int,
1031
+ max_frames_per_batch: int,
1032
+ device: torch.device,
1033
+ dtype: torch.dtype,
1034
+ ):
1035
+ from vllm.v1.worker.encoder_cudagraph_defs import (
1036
+ EncoderCudaGraphCaptureInputs,
1037
+ )
1038
+
1039
+ # Size the buffer to hold the maximum possible tiles for this budget.
1040
+ total_tiles = max(token_budget // self.num_image_token, 1)
1041
+ image_size = self.config.vision_config.image_size
1042
+
1043
+ dummy_pixel_values = torch.randn(
1044
+ total_tiles, 3, image_size, image_size, device=device, dtype=dtype
1045
+ )
1046
+
1047
+ return EncoderCudaGraphCaptureInputs(
1048
+ values={"pixel_values_flat": dummy_pixel_values},
1049
+ )
1050
+
1051
+ def prepare_encoder_cudagraph_replay_buffers(
1052
+ self,
1053
+ mm_kwargs: dict[str, Any],
1054
+ max_batch_size: int,
1055
+ max_frames_per_batch: int,
1056
+ ):
1057
+ from vllm.v1.worker.encoder_cudagraph_defs import (
1058
+ EncoderCudaGraphReplayBuffers,
1059
+ )
1060
+
1061
+ modality = self.get_input_modality(mm_kwargs)
1062
+ pv_key = (
1063
+ "pixel_values_flat" if modality == "image" else "pixel_values_flat_video"
1064
+ )
1065
+ return EncoderCudaGraphReplayBuffers(
1066
+ values={"pixel_values_flat": mm_kwargs[pv_key]},
1067
+ )
1068
+
1069
+ def encoder_cudagraph_forward(
1070
+ self,
1071
+ values: dict[str, torch.Tensor],
1072
+ ) -> torch.Tensor:
1073
+ # The graph is always captured with pixel_values_flat as the input
1074
+ # buffer. During video replay the manager copies video tiles into
1075
+ # this same buffer before calling graph.replay(), so we always read
1076
+ # from pixel_values_flat here.
1077
+ pixel_values = values["pixel_values_flat"]
1078
+ out = self.extract_feature(pixel_values) # [N, num_image_token, H]
1079
+ return out.view(-1, self.config.text_config.hidden_size)
1080
+
1081
+ def encoder_eager_forward(
1082
+ self,
1083
+ mm_kwargs: dict[str, Any],
1084
+ ) -> torch.Tensor:
1085
+ if self.get_input_modality(mm_kwargs) == "image":
1086
+ pixel_values = mm_kwargs["pixel_values_flat"]
1087
+ else:
1088
+ pixel_values = mm_kwargs["pixel_values_flat_video"]
1089
+ out = self.extract_feature(pixel_values) # [N, num_image_token, H]
1090
+ return out.view(-1, self.config.text_config.hidden_size)