diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 342dcb8b6f..b23b8d5057 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass, field +from functools import partial import pytest @@ -10,6 +11,7 @@ from vllm.platforms import current_platform from ....conftest import IMAGE_ASSETS, VIDEO_ASSETS from ....utils import create_new_process_for_each_test +from ...utils import dummy_hf_overrides from .vlm_utils.builders import sample_frames_with_video_metadata @@ -26,6 +28,7 @@ class VitCudagraphTestConfig: num_video_frames: int = 16 needs_video_metadata: bool = False vllm_runner_kwargs: dict = field(default_factory=dict) + compilation_config_overrides: dict = field(default_factory=dict) marks: list = field(default_factory=list) @@ -99,17 +102,36 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { ), "step3_vl": VitCudagraphTestConfig( model="stepfun-ai/Step3-VL-10B", + modalities=["image"], image_prompt=step3_vl_chat_template("What is in this image?"), - video_prompt=None, + # Single bucket sized to cover the largest test image's output + # tokens (1152 > 1141 for cherry_blossom). The default auto- + # inferred range fans out into multiple power-of-2 buckets, each + # holding a full ViT capture pool. + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [1152], + }, + # Shrink to 1 text + 1 vision layer with random weights so the + # test runs on any CI GPU (incl. L4) and skips the 20 GiB weight + # download. The test only validates that encoder CG capture/ + # replay functions correctly, not output quality. + vllm_runner_kwargs={ + "load_format": "dummy", + "hf_overrides": partial( + dummy_hf_overrides, + model_arch="StepVLForConditionalGeneration", + ), + }, ), } -def get_compilation_config(): +def get_compilation_config(config: VitCudagraphTestConfig): return { "cudagraph_mm_encoder": True, "encoder_cudagraph_max_vision_items_per_batch": 1, "encoder_cudagraph_max_frames_per_batch": 16, + **config.compilation_config_overrides, } @@ -141,7 +163,7 @@ def test_vit_cudagraph_image(model_id, vllm_runner, image_assets): max_model_len=config.max_model_len, max_num_seqs=config.max_num_seqs, limit_mm_per_prompt={"image": 1}, - compilation_config=get_compilation_config(), + compilation_config=get_compilation_config(config), **config.vllm_runner_kwargs, ) as vllm_model: outputs = vllm_model.generate_greedy( @@ -194,7 +216,7 @@ def test_vit_cudagraph_video(model_id, vllm_runner, video_assets): max_model_len=config.max_model_len, max_num_seqs=config.max_num_seqs, limit_mm_per_prompt={"video": 1}, - compilation_config=get_compilation_config(), + compilation_config=get_compilation_config(config), **config.vllm_runner_kwargs, ) as vllm_model: outputs = vllm_model.generate_greedy(