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, 16:15 UTC
[CI/Build] Add e2e test for ViT CUDA graph (#40780)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Apr 24, 2026 at 18:12 UTC
b5587e1013d0e352bb33c30b456d5221aebecd8c
2 files changed
+167
.buildkite/test_areas/models_multimodal.yaml
+1
@@ -28,6 +28,7 @@ steps:
28
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
29
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
30
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
31
+ - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
32
mirror:
33
amd:
34
device: mi325_1
tests/models/multimodal/generation/test_vit_cudagraph.py
new
+166
@@ -0,0 +1,166 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass, field
5
+
6
+import pytest
7
+
8
+from vllm.multimodal.video import sample_frames_from_video
9
+from vllm.platforms import current_platform
10
+
11
+from ....conftest import IMAGE_ASSETS, VIDEO_ASSETS
12
+from ....utils import create_new_process_for_each_test
13
+from .vlm_utils.builders import sample_frames_with_video_metadata
14
+
15
+
16
+@dataclass
17
+class VitCudagraphTestConfig:
18
+ model: str
19
+ modalities: list[str] = field(default_factory=lambda: ["image", "video"])
20
+ image_prompt: str | None = None
21
+ video_prompt: str | None = None
22
+ dtype: str = "bfloat16"
23
+ max_model_len: int = 4096
24
+ max_tokens: int = 64
25
+ max_num_seqs: int = 2
26
+ num_video_frames: int = 16
27
+ needs_video_metadata: bool = False
28
+ vllm_runner_kwargs: dict = field(default_factory=dict)
29
+ marks: list = field(default_factory=list)
30
+
31
+
32
+def params_with_marks(
33
+ configs: dict[str, VitCudagraphTestConfig],
34
+) -> list[pytest.param]:
35
+ return [
36
+ pytest.param(model_id, marks=cfg.marks) for model_id, cfg in configs.items()
37
+ ]
38
+
39
+
40
+def qwen_vl_chat_template(content: str) -> str:
41
+ return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
42
+
43
+
44
+MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
45
+ "qwen3_vl": VitCudagraphTestConfig(
46
+ model="Qwen/Qwen3-VL-2B-Instruct",
47
+ image_prompt=qwen_vl_chat_template(
48
+ "<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
49
+ ),
50
+ video_prompt=qwen_vl_chat_template(
51
+ "<|vision_start|><|video_pad|><|vision_end|>"
52
+ "Describe this video in one sentence."
53
+ ),
54
+ needs_video_metadata=True,
55
+ marks=[pytest.mark.core_model],
56
+ ),
57
+ # TODO: Add more models below.
58
+}
59
+
60
+
61
+def get_compilation_config():
62
+ return {
63
+ "cudagraph_mm_encoder": True,
64
+ "encoder_cudagraph_max_vision_items_per_batch": 1,
65
+ "encoder_cudagraph_max_frames_per_batch": 16,
66
+ }
67
+
68
+
69
+# ---------------------------------------------------------------------------
70
+# Tests
71
+# ---------------------------------------------------------------------------
72
+
73
+
74
+@pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))
75
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
76
+@create_new_process_for_each_test()
77
+def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):
78
+ config = MODEL_CONFIGS[model_id]
79
+
80
+ if "image" not in config.modalities:
81
+ pytest.skip(f"{model_id} does not support the image modality.")
82
+
83
+ image_prompts = IMAGE_ASSETS.prompts(
84
+ {
85
+ "stop_sign": config.image_prompt, # type: ignore[typeddict-item]
86
+ "cherry_blossom": config.image_prompt, # type: ignore[typeddict-item]
87
+ }
88
+ )
89
+ images = [[asset.pil_image] for asset in image_assets]
90
+
91
+ with vllm_runner(
92
+ config.model,
93
+ dtype=config.dtype,
94
+ max_model_len=config.max_model_len,
95
+ max_num_seqs=config.max_num_seqs,
96
+ limit_mm_per_prompt={"image": 1},
97
+ compilation_config=get_compilation_config(),
98
+ **config.vllm_runner_kwargs,
99
+ ) as vllm_model:
100
+ outputs = vllm_model.generate_greedy(
101
+ image_prompts, config.max_tokens, images=images
102
+ )
103
+
104
+ # Basic validation that we got a response
105
+ assert len(outputs) == 2
106
+ output_ids, output_text = outputs[0]
107
+
108
+ # Ensure we got some output
109
+ assert len(output_ids) > 0
110
+ assert len(output_text) > 0
111
+
112
+ # Ensure the output is a string
113
+ assert isinstance(output_text, str)
114
+
115
+
116
+@pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))
117
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
118
+@create_new_process_for_each_test()
119
+def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):
120
+ config = MODEL_CONFIGS[model_id]
121
+
122
+ if "video" not in config.modalities:
123
+ pytest.skip(f"{model_id} does not support the video modality")
124
+
125
+ video_prompts = VIDEO_ASSETS.prompts(
126
+ {
127
+ "baby_reading": config.video_prompt, # type: ignore[typeddict-item]
128
+ }
129
+ )
130
+ if config.needs_video_metadata:
131
+ sampled_vids = [
132
+ sample_frames_with_video_metadata(
133
+ (asset.np_ndarrays, asset.metadata), config.num_video_frames
134
+ )
135
+ for asset in video_assets
136
+ ]
137
+ else:
138
+ sampled_vids = [
139
+ sample_frames_from_video(asset.np_ndarrays, config.num_video_frames)
140
+ for asset in video_assets
141
+ ]
142
+ videos = [sampled_vids[0]]
143
+
144
+ with vllm_runner(
145
+ config.model,
146
+ dtype=config.dtype,
147
+ max_model_len=config.max_model_len,
148
+ max_num_seqs=config.max_num_seqs,
149
+ limit_mm_per_prompt={"video": 1},
150
+ compilation_config=get_compilation_config(),
151
+ **config.vllm_runner_kwargs,
152
+ ) as vllm_model:
153
+ outputs = vllm_model.generate_greedy(
154
+ video_prompts, config.max_tokens, videos=videos
155
+ )
156
+
157
+ # Basic validation that we got a response
158
+ assert len(outputs) == 1
159
+ output_ids, output_text = outputs[0]
160
+
161
+ # Ensure we got some output
162
+ assert len(output_ids) > 0
163
+ assert len(output_text) > 0
164
+
165
+ # Ensure the output is a string
166
+ assert isinstance(output_text, str)