| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass, field |
| 5 | from functools import partial |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.multimodal.video import sample_frames_from_video |
| 10 | from vllm.platforms import current_platform |
| 11 | |
| 12 | from ....conftest import IMAGE_ASSETS, VIDEO_ASSETS |
| 13 | from ...utils import dummy_hf_overrides |
| 14 | from .vlm_utils.builders import sample_frames_with_video_metadata |
| 15 | |
| 16 | |
| 17 | @dataclass |
| 18 | class VitCudagraphTestConfig: |
| 19 | model: str |
| 20 | modalities: list[str] = field(default_factory=lambda: ["image", "video"]) |
| 21 | image_prompt: str | None = None |
| 22 | video_prompt: str | None = None |
| 23 | dtype: str = "bfloat16" |
| 24 | max_model_len: int = 4096 |
| 25 | max_tokens: int = 64 |
| 26 | max_num_seqs: int = 2 |
| 27 | num_video_frames: int = 16 |
| 28 | needs_video_metadata: bool = False |
| 29 | vllm_runner_kwargs: dict = field(default_factory=dict) |
| 30 | compilation_config_overrides: dict = field(default_factory=dict) |
| 31 | marks: list = field(default_factory=list) |
| 32 | skip: bool = False |
| 33 | |
| 34 | |
| 35 | def params_with_marks( |
| 36 | configs: dict[str, VitCudagraphTestConfig], |
| 37 | ) -> list[pytest.param]: |
| 38 | return [ |
| 39 | pytest.param(model_id, marks=cfg.marks) for model_id, cfg in configs.items() |
| 40 | ] |
| 41 | |
| 42 | |
| 43 | def qwen_vl_chat_template(content: str) -> str: |
| 44 | return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n" |
| 45 | |
| 46 | |
| 47 | def internvl_chat_template(content: str) -> str: |
| 48 | return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n" |
| 49 | |
| 50 | |
| 51 | def kimi_vl_chat_template(content: str) -> str: |
| 52 | return ( |
| 53 | f"<|im_user|>user<|im_middle|>{content}<|im_end|>" |
| 54 | "<|im_assistant|>assistant<|im_middle|>" |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | def deepseek_v41_chat_template(content: str) -> str: |
| 59 | return f"<｜begin▁of▁sentence｜><｜User｜>{content}<｜Assistant｜></think>" |
| 60 | |
| 61 | |
| 62 | def step3_vl_chat_template(content: str) -> str: |
| 63 | return ( |
| 64 | "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n " |
| 65 | f"<im_patch>{content} <|EOT|><|BOT|>assistant\n" |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | def gemma3_chat_template(content: str) -> str: |
| 70 | return f"<bos><start_of_turn>user\n{content}<end_of_turn>\n<start_of_turn>model\n" |
| 71 | |
| 72 | |
| 73 | def ernie45_vl_chat_template(content: str) -> str: |
| 74 | return ( |
| 75 | f"<|begin_of_sentence|>User: {content}" |
| 76 | "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>\n" |
| 77 | "Assistant: <think></think>" |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | def minicpmv_25_chat_template(content: str) -> str: |
| 82 | """Llama3-style chat template used by MiniCPM-V 2.5.""" |
| 83 | return ( |
| 84 | f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n" |
| 85 | f"{content}" |
| 86 | f"<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | def minicpmv_chat_template(content: str) -> str: |
| 91 | """ChatML template used by MiniCPM-V 2.6 / 4.0 / 4.5.""" |
| 92 | return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n" |
| 93 | |
| 94 | |
| 95 | MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { |
| 96 | "gemma3": VitCudagraphTestConfig( |
| 97 | model="google/gemma-3-4b-it", |
| 98 | modalities=["image"], |
| 99 | image_prompt=gemma3_chat_template("<start_of_image>What is in this image?"), |
| 100 | compilation_config_overrides={ |
| 101 | "encoder_cudagraph_token_budgets": [512], |
| 102 | }, |
| 103 | dtype="bfloat16", |
| 104 | max_model_len=4096, |
| 105 | ), |
| 106 | "llama4": VitCudagraphTestConfig( |
| 107 | model="meta-llama/Llama-4-Scout-17B-16E-Instruct", |
| 108 | modalities=["image"], |
| 109 | image_prompt=( |
| 110 | "<|begin_of_text|><|header_start|>user<|header_end|>\n\n" |
| 111 | "<|image|>What is in this image?<|eot|>" |
| 112 | "<|header_start|>assistant<|header_end|>\n\n" |
| 113 | ), |
| 114 | max_model_len=4096, |
| 115 | max_tokens=32, |
| 116 | max_num_seqs=2, |
| 117 | vllm_runner_kwargs={ |
| 118 | "load_format": "dummy", |
| 119 | "hf_overrides": partial( |
| 120 | dummy_hf_overrides, |
| 121 | model_arch="Llama4ForConditionalGeneration", |
| 122 | ), |
| 123 | }, |
| 124 | marks=[pytest.mark.core_model], |
| 125 | ), |
| 126 | "qwen2_vl": VitCudagraphTestConfig( |
| 127 | model="Qwen/Qwen2-VL-2B-Instruct", |
| 128 | image_prompt=qwen_vl_chat_template( |
| 129 | "<|vision_start|><|image_pad|><|vision_end|>What is in this image?" |
| 130 | ), |
| 131 | video_prompt=qwen_vl_chat_template( |
| 132 | "<|vision_start|><|video_pad|><|vision_end|>" |
| 133 | "Describe this video in one sentence." |
| 134 | ), |
| 135 | needs_video_metadata=False, |
| 136 | marks=[pytest.mark.core_model], |
| 137 | ), |
| 138 | "qwen2_5_vl": VitCudagraphTestConfig( |
| 139 | model="Qwen/Qwen2.5-VL-3B-Instruct", |
| 140 | image_prompt=qwen_vl_chat_template( |
| 141 | "<|vision_start|><|image_pad|><|vision_end|>What is in this image?" |
| 142 | ), |
| 143 | video_prompt=qwen_vl_chat_template( |
| 144 | "<|vision_start|><|video_pad|><|vision_end|>" |
| 145 | "Describe this video in one sentence." |
| 146 | ), |
| 147 | needs_video_metadata=True, |
| 148 | marks=[pytest.mark.core_model], |
| 149 | ), |
| 150 | "kimi_vl": VitCudagraphTestConfig( |
| 151 | model="moonshotai/Kimi-VL-A3B-Instruct", |
| 152 | modalities=["image"], |
| 153 | image_prompt=kimi_vl_chat_template( |
| 154 | "<|media_start|>image<|media_content|><|media_pad|><|media_end|>" |
| 155 | "What is in this image?" |
| 156 | ), |
| 157 | needs_video_metadata=False, |
| 158 | # Single bucket sized to cover the test images' output tokens. |
| 159 | # The default auto-inferred range fans out into multiple power-of-2 |
| 160 | # buckets, each holding a full ViT capture pool. |
| 161 | compilation_config_overrides={ |
| 162 | "encoder_cudagraph_token_budgets": [1024], |
| 163 | }, |
| 164 | # Shrink to 1 text + 1 vision layer with random weights so the |
| 165 | # test runs on any CI GPU (incl. L4) and skips the multi-GiB |
| 166 | # weight download. The test only validates that encoder CG |
| 167 | # capture/replay functions correctly, not output quality. |
| 168 | vllm_runner_kwargs={ |
| 169 | "trust_remote_code": True, |
| 170 | "load_format": "dummy", |
| 171 | "hf_overrides": partial( |
| 172 | dummy_hf_overrides, |
| 173 | model_arch="KimiVLForConditionalGeneration", |
| 174 | ), |
| 175 | }, |
| 176 | marks=[pytest.mark.core_model], |
| 177 | ), |
| 178 | "minimax_m3": VitCudagraphTestConfig( |
| 179 | model="MiniMaxAI/MiniMax-M3", |
| 180 | modalities=["image"], |
| 181 | image_prompt=("]~b]user\n]<]image[>[What is in this image?[e~[\n]~b]ai\n"), |
| 182 | compilation_config_overrides={ |
| 183 | "encoder_cudagraph_token_budgets": [512], |
| 184 | }, |
| 185 | vllm_runner_kwargs={ |
| 186 | "trust_remote_code": True, |
| 187 | "load_format": "dummy", |
| 188 | "hf_overrides": partial( |
| 189 | dummy_hf_overrides, |
| 190 | model_arch="MiniMaxM3SparseForConditionalGeneration", |
| 191 | ), |
| 192 | "mm_processor_kwargs": {"max_long_side_pixel": 336}, |
| 193 | }, |
| 194 | marks=[pytest.mark.core_model], |
| 195 | ), |
| 196 | "qwen3_vl": VitCudagraphTestConfig( |
| 197 | model="Qwen/Qwen3-VL-2B-Instruct", |
| 198 | image_prompt=qwen_vl_chat_template( |
| 199 | "<|vision_start|><|image_pad|><|vision_end|>What is in this image?" |
| 200 | ), |
| 201 | video_prompt=qwen_vl_chat_template( |
| 202 | "<|vision_start|><|video_pad|><|vision_end|>" |
| 203 | "Describe this video in one sentence." |
| 204 | ), |
| 205 | needs_video_metadata=True, |
| 206 | marks=[pytest.mark.core_model], |
| 207 | ), |
| 208 | "qwen3_5": VitCudagraphTestConfig( |
| 209 | model="Qwen/Qwen3.5-0.8B", |
| 210 | image_prompt=qwen_vl_chat_template( |
| 211 | "<|vision_start|><|image_pad|><|vision_end|>What is in this image?" |
| 212 | ), |
| 213 | video_prompt=qwen_vl_chat_template( |
| 214 | "<|vision_start|><|video_pad|><|vision_end|>" |
| 215 | "Describe this video in one sentence." |
| 216 | ), |
| 217 | needs_video_metadata=True, |
| 218 | vllm_runner_kwargs={"enable_chunked_prefill": True}, |
| 219 | marks=[pytest.mark.core_model], |
| 220 | ), |
| 221 | "internvl": VitCudagraphTestConfig( |
| 222 | model="OpenGVLab/InternVL3-1B", |
| 223 | num_video_frames=8, |
| 224 | image_prompt=internvl_chat_template("<image>\nWhat is in this image?"), |
| 225 | video_prompt=internvl_chat_template( |
| 226 | "<video>\nDescribe this video in one sentence." |
| 227 | ), |
| 228 | needs_video_metadata=False, |
| 229 | vllm_runner_kwargs={"trust_remote_code": True}, |
| 230 | marks=[pytest.mark.core_model], |
| 231 | ), |
| 232 | "idefics3": VitCudagraphTestConfig( |
| 233 | model="HuggingFaceTB/SmolVLM-256M-Instruct", |
| 234 | modalities=["image"], |
| 235 | image_prompt=( |
| 236 | "<|begin_of_text|>User:<image>What is in this image?" |
| 237 | "<end_of_utterance>\nAssistant:" |
| 238 | ), |
| 239 | max_model_len=4096, |
| 240 | compilation_config_overrides={ |
| 241 | "encoder_cudagraph_token_budgets": [4096], |
| 242 | }, |
| 243 | vllm_runner_kwargs={"gpu_memory_utilization": 0.80}, |
| 244 | marks=[pytest.mark.core_model], |
| 245 | ), |
| 246 | "deepseek_v41": VitCudagraphTestConfig( |
| 247 | model="deepseek-ai/DeepSeek-V4.1-Flash", |
| 248 | modalities=["image"], |
| 249 | image_prompt=deepseek_v41_chat_template( |
| 250 | "<｜deepseek_image｜>\n\nWhat is in this image?" |
| 251 | ), |
| 252 | compilation_config_overrides={ |
| 253 | "encoder_cudagraph_token_budgets": [1024], |
| 254 | }, |
| 255 | vllm_runner_kwargs={ |
| 256 | "load_format": "dummy", |
| 257 | "attention_backend": ( |
| 258 | "ROCM_FLASHMLA_SPARSE_DSV4" |
| 259 | if current_platform.is_rocm() |
| 260 | else "FLASHMLA_SPARSE_DSV41" |
| 261 | ), |
| 262 | "hf_overrides": partial( |
| 263 | dummy_hf_overrides, |
| 264 | model_arch="DeepseekV41ForCausalLM", |
| 265 | exist_overrides={"vision_n_layers": 1}, |
| 266 | ), |
| 267 | }, |
| 268 | ), |
| 269 | "step3_vl": VitCudagraphTestConfig( |
| 270 | model="stepfun-ai/Step3-VL-10B", |
| 271 | modalities=["image"], |
| 272 | image_prompt=step3_vl_chat_template("What is in this image?"), |
| 273 | # Single bucket sized to cover the largest test image's output |
| 274 | # tokens (1152 > 1141 for cherry_blossom). The default auto- |
| 275 | # inferred range fans out into multiple power-of-2 buckets, each |
| 276 | # holding a full ViT capture pool. |
| 277 | compilation_config_overrides={ |
| 278 | "encoder_cudagraph_token_budgets": [1152], |
| 279 | }, |
| 280 | # Shrink to 1 text + 1 vision layer with random weights so the |
| 281 | # test runs on any CI GPU (incl. L4) and skips the 20 GiB weight |
| 282 | # download. The test only validates that encoder CG capture/ |
| 283 | # replay functions correctly, not output quality. |
| 284 | vllm_runner_kwargs={ |
| 285 | "load_format": "dummy", |
| 286 | "hf_overrides": partial( |
| 287 | dummy_hf_overrides, |
| 288 | model_arch="StepVLForConditionalGeneration", |
| 289 | ), |
| 290 | }, |
| 291 | ), |
| 292 | "ernie45_vl": VitCudagraphTestConfig( |
| 293 | model="baidu/ERNIE-4.5-VL-28B-A3B-PT", |
| 294 | # Image only: Ernie's resampler applies a temporal conv for video |
| 295 | # that changes the output token count, so video uses the eager path. |
| 296 | modalities=["image"], |
| 297 | image_prompt=ernie45_vl_chat_template("What is in this image?"), |
| 298 | # Ernie4_5_VLMoeModel is deliberately not torch-compiled, since its |
| 299 | # split of text and vision experts breaks compilation, so piecewise |
| 300 | # cudagraphs have nothing to partition. Only the encoder graphs this |
| 301 | # test covers are captured. |
| 302 | compilation_config_overrides={ |
| 303 | "cudagraph_mode": 2, |
| 304 | }, |
| 305 | # Shrink to 1 text + 1 vision layer with random weights so the test |
| 306 | # runs on any CI GPU and skips the ~56 GiB weight download. The test |
| 307 | # only validates encoder CG capture/replay, not output quality. |
| 308 | vllm_runner_kwargs={ |
| 309 | "load_format": "dummy", |
| 310 | "trust_remote_code": True, |
| 311 | "revision": "refs/pr/17", |
| 312 | "hf_overrides": partial( |
| 313 | dummy_hf_overrides, |
| 314 | model_arch="Ernie4_5_VLMoeForConditionalGeneration", |
| 315 | ), |
| 316 | }, |
| 317 | ), |
| 318 | "glm4_1v": VitCudagraphTestConfig( |
| 319 | model="zai-org/GLM-4.1V-9B-Thinking", |
| 320 | image_prompt=( |
| 321 | "[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n" |
| 322 | "<|begin_of_image|><|image|><|end_of_image|>" |
| 323 | "What is in this image?<|assistant|>assistant\n" |
| 324 | ), |
| 325 | video_prompt=( |
| 326 | "[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n" |
| 327 | "<|begin_of_video|><|video|><|end_of_video|>" |
| 328 | "Describe this video in one sentence<|assistant|>assistant\n" |
| 329 | ), |
| 330 | needs_video_metadata=True, |
| 331 | marks=[pytest.mark.core_model], |
| 332 | vllm_runner_kwargs={ |
| 333 | "load_format": "dummy", |
| 334 | "hf_overrides": partial( |
| 335 | dummy_hf_overrides, |
| 336 | model_arch="Glm4vForConditionalGeneration", |
| 337 | ), |
| 338 | }, |
| 339 | ), |
| 340 | "deepseek_ocr": VitCudagraphTestConfig( |
| 341 | model="deepseek-ai/DeepSeek-OCR", |
| 342 | modalities=["image"], |
| 343 | image_prompt="<image>\nWhat is in this image?", |
| 344 | marks=[pytest.mark.core_model], |
| 345 | compilation_config_overrides={ |
| 346 | "encoder_cudagraph_token_budgets": [272], |
| 347 | "mode": 0, |
| 348 | "cudagraph_mode": 2, |
| 349 | }, |
| 350 | vllm_runner_kwargs={ |
| 351 | "load_format": "dummy", |
| 352 | "hf_overrides": partial( |
| 353 | dummy_hf_overrides, |
| 354 | model_arch="DeepseekOCRForCausalLM", |
| 355 | ), |
| 356 | }, |
| 357 | skip=True, # TODO: Re-enable this once OOM issues are resolved on CI. |
| 358 | ), |
| 359 | "gemma4": VitCudagraphTestConfig( |
| 360 | model="google/gemma-4-E2B-it", |
| 361 | image_prompt=( |
| 362 | "<bos><start_of_turn>user\n<|image|>\nWhat is in this image?<end_of_turn>\n" |
| 363 | "<start_of_turn>model\n" |
| 364 | ), |
| 365 | video_prompt=( |
| 366 | "<bos><start_of_turn>user\n<|video|>\nDescribe this video in one sentence." |
| 367 | "<end_of_turn>\n<start_of_turn>model\n" |
| 368 | ), |
| 369 | # The 16-frame test video produces 1056 vision tokens. Capture only |
| 370 | # the smallest supported bucket that covers it instead of all default |
| 371 | # buckets through max_model_len, which adds unrelated memory pressure. |
| 372 | compilation_config_overrides={ |
| 373 | "encoder_cudagraph_token_budgets": [1120], |
| 374 | }, |
| 375 | needs_video_metadata=True, |
| 376 | marks=[pytest.mark.core_model], |
| 377 | ), |
| 378 | "minicpmv_25": VitCudagraphTestConfig( |
| 379 | model="openbmb/MiniCPM-Llama3-V-2_5", |
| 380 | modalities=["image"], |
| 381 | image_prompt=minicpmv_25_chat_template( |
| 382 | "(<image>./</image>)\nWhat is in this image?" |
| 383 | ), |
| 384 | # CI runs on 35GB MIG slices: every budget captures one graph per |
| 385 | # patch-grid bucket, so the default budgets through max_model_len |
| 386 | # OOM there. A small budget set covers the test images (each item |
| 387 | # is at most (max_slice_num + 1) * query_num = 640 tokens). |
| 388 | compilation_config_overrides={ |
| 389 | "encoder_cudagraph_token_budgets": [64, 1024], |
| 390 | }, |
| 391 | vllm_runner_kwargs={"trust_remote_code": True}, |
| 392 | marks=[pytest.mark.core_model], |
| 393 | ), |
| 394 | "minicpmv_26": VitCudagraphTestConfig( |
| 395 | model="openbmb/MiniCPM-V-2_6", |
| 396 | image_prompt=minicpmv_chat_template( |
| 397 | "(<image>./</image>)\nWhat is in this image?" |
| 398 | ), |
| 399 | video_prompt=minicpmv_chat_template( |
| 400 | "(<video>./</video>)\nDescribe this video in one sentence." |
| 401 | ), |
| 402 | max_model_len=2048, |
| 403 | # Fewer frames keep the video item within the 1024 budget. |
| 404 | num_video_frames=4, |
| 405 | compilation_config_overrides={ |
| 406 | "encoder_cudagraph_token_budgets": [64, 1024], |
| 407 | "encoder_cudagraph_max_frames_per_batch": 4, |
| 408 | }, |
| 409 | vllm_runner_kwargs={"trust_remote_code": True}, |
| 410 | marks=[pytest.mark.core_model], |
| 411 | ), |
| 412 | "minicpmv_40": VitCudagraphTestConfig( |
| 413 | model="openbmb/MiniCPM-V-4", |
| 414 | image_prompt=minicpmv_chat_template( |
| 415 | "(<image>./</image>)\nWhat is in this image?" |
| 416 | ), |
| 417 | video_prompt=minicpmv_chat_template( |
| 418 | "(<video>./</video>)\nDescribe this video in one sentence." |
| 419 | ), |
| 420 | max_model_len=2048, |
| 421 | num_video_frames=4, |
| 422 | compilation_config_overrides={ |
| 423 | "encoder_cudagraph_token_budgets": [64, 1024], |
| 424 | "encoder_cudagraph_max_frames_per_batch": 4, |
| 425 | }, |
| 426 | vllm_runner_kwargs={"trust_remote_code": True}, |
| 427 | marks=[pytest.mark.core_model], |
| 428 | ), |
| 429 | } |
| 430 | |
| 431 | |
| 432 | def get_compilation_config(config: VitCudagraphTestConfig): |
| 433 | return { |
| 434 | "cudagraph_mm_encoder": True, |
| 435 | "encoder_cudagraph_max_vision_items_per_batch": 1, |
| 436 | "encoder_cudagraph_max_frames_per_batch": 16, |
| 437 | **config.compilation_config_overrides, |
| 438 | } |
| 439 | |
| 440 | |
| 441 | # --------------------------------------------------------------------------- |
| 442 | # Tests |
| 443 | # --------------------------------------------------------------------------- |
| 444 | |
| 445 | |
| 446 | @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS)) |
| 447 | @pytest.mark.skipif( |
| 448 | not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" |
| 449 | ) |
| 450 | def test_vit_cudagraph_image(model_id, vllm_runner, image_assets): |
| 451 | config = MODEL_CONFIGS[model_id] |
| 452 | |
| 453 | if config.skip: |
| 454 | pytest.skip(f"{model_id} is marked to be skipped.") |
| 455 | |
| 456 | if "image" not in config.modalities: |
| 457 | pytest.skip(f"{model_id} does not support the image modality.") |
| 458 | |
| 459 | image_prompts = IMAGE_ASSETS.prompts( |
| 460 | { |
| 461 | "stop_sign": config.image_prompt, # type: ignore[typeddict-item] |
| 462 | "cherry_blossom": config.image_prompt, # type: ignore[typeddict-item] |
| 463 | } |
| 464 | ) |
| 465 | images = [[asset.pil_image] for asset in image_assets] |
| 466 | |
| 467 | with vllm_runner( |
| 468 | config.model, |
| 469 | dtype=config.dtype, |
| 470 | max_model_len=config.max_model_len, |
| 471 | max_num_seqs=config.max_num_seqs, |
| 472 | limit_mm_per_prompt={"image": 1}, |
| 473 | compilation_config=get_compilation_config(config), |
| 474 | **config.vllm_runner_kwargs, |
| 475 | ) as vllm_model: |
| 476 | outputs = vllm_model.generate_greedy( |
| 477 | image_prompts, config.max_tokens, images=images |
| 478 | ) |
| 479 | |
| 480 | # Basic validation that we got a response |
| 481 | assert len(outputs) == 2 |
| 482 | output_ids, output_text = outputs[0] |
| 483 | |
| 484 | # Ensure we got some output |
| 485 | assert len(output_ids) > 0 |
| 486 | assert len(output_text) > 0 |
| 487 | |
| 488 | # Ensure the output is a string |
| 489 | assert isinstance(output_text, str) |
| 490 | |
| 491 | |
| 492 | @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS)) |
| 493 | @pytest.mark.skipif( |
| 494 | not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" |
| 495 | ) |
| 496 | def test_vit_cudagraph_video(model_id, vllm_runner, video_assets): |
| 497 | config = MODEL_CONFIGS[model_id] |
| 498 | |
| 499 | if config.skip: |
| 500 | pytest.skip(f"{model_id} is marked to be skipped.") |
| 501 | |
| 502 | if "video" not in config.modalities: |
| 503 | pytest.skip(f"{model_id} does not support the video modality") |
| 504 | |
| 505 | video_prompts = VIDEO_ASSETS.prompts( |
| 506 | { |
| 507 | "baby_reading": config.video_prompt, # type: ignore[typeddict-item] |
| 508 | } |
| 509 | ) |
| 510 | if config.needs_video_metadata: |
| 511 | sampled_vids = [ |
| 512 | sample_frames_with_video_metadata( |
| 513 | (asset.np_ndarrays, asset.metadata), config.num_video_frames |
| 514 | ) |
| 515 | for asset in video_assets |
| 516 | ] |
| 517 | else: |
| 518 | sampled_vids = [ |
| 519 | sample_frames_from_video(asset.np_ndarrays, config.num_video_frames) |
| 520 | for asset in video_assets |
| 521 | ] |
| 522 | videos = [sampled_vids[0]] |
| 523 | |
| 524 | with vllm_runner( |
| 525 | config.model, |
| 526 | dtype=config.dtype, |
| 527 | max_model_len=config.max_model_len, |
| 528 | max_num_seqs=config.max_num_seqs, |
| 529 | limit_mm_per_prompt={"video": 1}, |
| 530 | compilation_config=get_compilation_config(config), |
| 531 | **config.vllm_runner_kwargs, |
| 532 | ) as vllm_model: |
| 533 | outputs = vllm_model.generate_greedy( |
| 534 | video_prompts, config.max_tokens, videos=videos |
| 535 | ) |
| 536 | |
| 537 | # Basic validation that we got a response |
| 538 | assert len(outputs) == 1 |
| 539 | output_ids, output_text = outputs[0] |
| 540 | |
| 541 | # Ensure we got some output |
| 542 | assert len(output_ids) > 0 |
| 543 | assert len(output_text) > 0 |
| 544 | |
| 545 | # Ensure the output is a string |
| 546 | assert isinstance(output_text, str) |
| 547 | |
| 548 | |
| 549 | def _check_eonly_encoder_outputs(worker, batches): |
| 550 | """Compare real E-only replay with eager, retaining outputs across replay.""" |
| 551 | import torch |
| 552 | |
| 553 | from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner |
| 554 | |
| 555 | runner = worker.model_runner |
| 556 | assert isinstance(runner, MMEncoderModelRunner) |
| 557 | manager = runner.model_state.encoder_runner.cudagraph_manager |
| 558 | assert manager is not None and manager.is_captured() |
| 559 | retained: list[tuple[torch.Tensor, torch.Tensor]] = [] |
| 560 | max_error = 0.0 |
| 561 | stats = [] |
| 562 | with torch.inference_mode(): |
| 563 | for batch in batches: |
| 564 | kwargs = {key: value.to(runner.device) for key, value in batch.items()} |
| 565 | actual = manager.execute(kwargs) |
| 566 | expected = runner.model.embed_multimodal(**kwargs) |
| 567 | assert len(actual) == len(expected) |
| 568 | for output, reference in zip(actual, expected): |
| 569 | assert torch.isfinite(output).all() |
| 570 | # Two BF16 ulps at unit scale, fixed before model validation. |
| 571 | torch.testing.assert_close(output, reference, rtol=0.016, atol=0.016) |
| 572 | max_error = max(max_error, (output - reference).abs().max().item()) |
| 573 | for output, snapshot in retained: |
| 574 | torch.testing.assert_close(output, snapshot, rtol=0, atol=0) |
| 575 | retained.extend((output, output.clone()) for output in actual) |
| 576 | stats.append(manager.get_cumulative_stats()) |
| 577 | assert stats[0]["graph_hits"] > 0 |
| 578 | assert stats[1]["graph_hits"] > stats[0]["graph_hits"] |
| 579 | assert stats[2]["graph_misses"] > stats[1]["graph_misses"] |
| 580 | assert stats[3]["graph_hits"] > stats[2]["graph_hits"] |
| 581 | assert stats[4]["graph_hits"] > stats[3]["graph_hits"] |
| 582 | assert stats[5]["graph_hits"] > stats[4]["graph_hits"] |
| 583 | assert stats[6]["graph_hits"] > stats[5]["graph_hits"] |
| 584 | return {"max_abs_error": max_error, "stats": stats} |
| 585 | |
| 586 | |
| 587 | @pytest.mark.parametrize( |
| 588 | "model_id", |
| 589 | params_with_marks({key: MODEL_CONFIGS[key] for key in ("qwen3_vl", "qwen3_5")}), |
| 590 | ) |
| 591 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA graphs") |
| 592 | def test_eonly_vit_cudagraph_outputs( |
| 593 | model_id, vllm_runner, image_assets, video_assets, monkeypatch |
| 594 | ): |
| 595 | """Exercise startup capture, image/video replay, fallback and output ownership.""" |
| 596 | import numpy as np |
| 597 | from PIL import Image |
| 598 | from transformers import AutoProcessor |
| 599 | |
| 600 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 601 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 602 | config = MODEL_CONFIGS[model_id] |
| 603 | token_budgets = [64, 256, 2048] |
| 604 | processor = AutoProcessor.from_pretrained(config.model) |
| 605 | first, second = [asset.pil_image for asset in image_assets] |
| 606 | batches = [ |
| 607 | [first.resize((224, 224))], |
| 608 | [second.resize((448, 224)), first.resize((224, 224))], |
| 609 | [second.resize((1792, 1792))], |
| 610 | [second.resize((224, 224))], |
| 611 | [first.resize((1280, 720))], |
| 612 | [first.resize((224, 224))], |
| 613 | ] |
| 614 | frames = sample_frames_from_video(video_assets[0].np_ndarrays, 2) |
| 615 | video = np.stack( |
| 616 | [np.asarray(Image.fromarray(frame).resize((224, 224))) for frame in frames] |
| 617 | ) |
| 618 | with vllm_runner( |
| 619 | config.model, |
| 620 | dtype=config.dtype, |
| 621 | mm_encoder_only=True, |
| 622 | # FA padding NaNs are tracked separately in #57136. |
| 623 | mm_encoder_attn_backend="FLASHINFER", |
| 624 | enable_prefix_caching=False, |
| 625 | max_model_len=4096, |
| 626 | max_num_seqs=2, |
| 627 | limit_mm_per_prompt={"image": 2, "video": 1}, |
| 628 | compilation_config={ |
| 629 | "cudagraph_mode": "NONE", |
| 630 | "cudagraph_mm_encoder": True, |
| 631 | "encoder_cudagraph_token_budgets": token_budgets, |
| 632 | "encoder_cudagraph_max_vision_items_per_batch": 2, |
| 633 | "encoder_cudagraph_max_frames_per_batch": 2, |
| 634 | }, |
| 635 | ) as model: |
| 636 | mm_config = model.llm.llm_engine.vllm_config.model_config.multimodal_config |
| 637 | processor_kwargs = mm_config.merge_mm_processor_kwargs({}) |
| 638 | inputs = [ |
| 639 | dict( |
| 640 | processor.image_processor( |
| 641 | images=images, return_tensors="pt", **processor_kwargs |
| 642 | ) |
| 643 | ) |
| 644 | for images in batches |
| 645 | ] |
| 646 | inputs.append( |
| 647 | dict( |
| 648 | processor.video_processor( |
| 649 | videos=[video], |
| 650 | do_sample_frames=False, |
| 651 | return_tensors="pt", |
| 652 | **processor_kwargs, |
| 653 | ) |
| 654 | ) |
| 655 | ) |
| 656 | results = model.llm.collective_rpc( |
| 657 | partial(_check_eonly_encoder_outputs, batches=inputs) |
| 658 | ) |
| 659 | for result in results: |
| 660 | assert result["stats"][0]["token_budgets"] == token_budgets |
| 661 | print(f"E-only {model_id}: {results}") |