| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import tempfile |
| 4 | from collections.abc import Iterable |
| 5 | from contextlib import contextmanager |
| 6 | from functools import partial |
| 7 | from typing import Any, TypeAlias |
| 8 | |
| 9 | import numpy as np |
| 10 | import pytest |
| 11 | import torch |
| 12 | import torch.nn as nn |
| 13 | from PIL import Image |
| 14 | |
| 15 | from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config |
| 16 | from vllm.config.cache import CacheConfig |
| 17 | from vllm.config.multimodal import ( |
| 18 | AudioDummyOptions, |
| 19 | BaseDummyOptions, |
| 20 | ImageDummyOptions, |
| 21 | MultiModalDummyOptions, |
| 22 | VideoDummyOptions, |
| 23 | ) |
| 24 | from vllm.distributed import ( |
| 25 | cleanup_dist_env_and_memory, |
| 26 | init_distributed_environment, |
| 27 | initialize_model_parallel, |
| 28 | ) |
| 29 | from vllm.model_executor.models.interfaces import supports_multimodal |
| 30 | from vllm.multimodal import MULTIMODAL_REGISTRY, BatchedTensorInputs |
| 31 | from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext |
| 32 | from vllm.multimodal.utils import group_and_batch_mm_kwargs |
| 33 | from vllm.platforms import current_platform |
| 34 | from vllm.tokenizers import cached_tokenizer_from_config |
| 35 | from vllm.utils.collection_utils import is_list_of |
| 36 | from vllm.utils.torch_utils import set_default_torch_dtype |
| 37 | |
| 38 | from ....utils import create_new_process_for_each_test |
| 39 | from ...registry import HF_EXAMPLE_MODELS |
| 40 | from ...utils import dummy_hf_overrides |
| 41 | from .test_common import get_model_ids_to_test, get_token_prompt |
| 42 | |
| 43 | ImageInput = list[Image.Image] |
| 44 | VideoInput: TypeAlias = ( |
| 45 | list[Image.Image] | list[np.ndarray] | list[tuple[np.ndarray, dict[str, Any]]] |
| 46 | ) |
| 47 | AudioInput = list[tuple[np.ndarray, int]] |
| 48 | |
| 49 | |
| 50 | def _resize_data( |
| 51 | _data: Image.Image | np.ndarray, size_factor: float |
| 52 | ) -> Image.Image | np.ndarray: |
| 53 | assert size_factor <= 1, "Size factor must be less than 1" |
| 54 | # Image input |
| 55 | if isinstance(_data, Image.Image): |
| 56 | W, H = _data.width, _data.height |
| 57 | W, H = map(lambda x: int(x * size_factor), (W, H)) |
| 58 | return _data.resize((W, H)) |
| 59 | # Video input with PIL Images |
| 60 | elif is_list_of(_data, Image.Image): |
| 61 | W, H = next(iter(_data)).width, next(iter(_data)).height |
| 62 | T = len(_data) |
| 63 | T, W, H = map(lambda x: max(int(x * size_factor), 2), (T, W, H)) |
| 64 | return [d.resize((W, H)) for d in _data[:T]] |
| 65 | # Video input with numpy arrays |
| 66 | elif isinstance(_data, np.ndarray) and _data.ndim >= 4: |
| 67 | T, H, W, C = _data.shape[-4:] |
| 68 | T, H, W = map(lambda x: max(int(x * size_factor), 2), (T, H, W)) |
| 69 | return _data[..., :T, :H, :W, :C] |
| 70 | # Audio input |
| 71 | elif isinstance(_data, np.ndarray) and _data.ndim == 1: |
| 72 | return _data[: int(len(_data) * size_factor)] |
| 73 | raise AssertionError("This line should be unreachable.") |
| 74 | |
| 75 | |
| 76 | def resize_mm_data( |
| 77 | data: ImageInput | VideoInput | AudioInput, size_factors: tuple[float, ...] |
| 78 | ) -> ImageInput | VideoInput | AudioInput: |
| 79 | size_factors = size_factors[: len(data)] |
| 80 | if is_list_of(data, (Image.Image, np.ndarray, list)): |
| 81 | return [_resize_data(d, s) for d, s in zip(data, size_factors)] |
| 82 | elif is_list_of(data, tuple): |
| 83 | return [_resize_data(d, s) for (d, _), s in zip(data, size_factors)] |
| 84 | raise ValueError("Unsupported multimodal data type.") |
| 85 | |
| 86 | |
| 87 | def create_batched_mm_kwargs( |
| 88 | model_config: ModelConfig, |
| 89 | processor: BaseMultiModalProcessor, |
| 90 | size_factors: tuple[float, ...] = (1.0, 0.5, 0.25), |
| 91 | ) -> Iterable[tuple[str, int, BatchedTensorInputs]]: |
| 92 | processing_info = processor.info |
| 93 | supported_mm_limits = processing_info.get_supported_mm_limits() |
| 94 | mm_counts = { |
| 95 | modality: 3 if limit is None else limit |
| 96 | for modality, limit in supported_mm_limits.items() |
| 97 | } |
| 98 | processor_inputs = processor.get_dummy_inputs( |
| 99 | seq_len=model_config.max_model_len, |
| 100 | mm_counts=mm_counts, |
| 101 | mm_options=MultiModalDummyOptions(), |
| 102 | ) |
| 103 | mm_items = processor_inputs.mm_data_items |
| 104 | resized_mm_data = { |
| 105 | modality: resize_mm_data(items.data, size_factors) |
| 106 | for modality, items in mm_items.items() |
| 107 | } |
| 108 | |
| 109 | # video metadata will be added back to the resized video data here. |
| 110 | token_prompt = get_token_prompt(processor, resized_mm_data) |
| 111 | |
| 112 | mm_kwargs = processor( |
| 113 | prompt=token_prompt, |
| 114 | mm_items=processor.info.parse_mm_data(resized_mm_data), |
| 115 | hf_processor_mm_kwargs=processor_inputs.hf_processor_mm_kwargs, |
| 116 | )["mm_kwargs"].require_data() |
| 117 | |
| 118 | return group_and_batch_mm_kwargs( |
| 119 | [ |
| 120 | (modality, item) |
| 121 | for modality in supported_mm_limits |
| 122 | for item in mm_kwargs[modality] |
| 123 | ] |
| 124 | ) |
| 125 | |
| 126 | |
| 127 | # TODO(Isotr0py): Don't initialize model during test |
| 128 | @contextmanager |
| 129 | def initialize_dummy_model( |
| 130 | model_cls: type[nn.Module], |
| 131 | model_config: ModelConfig, |
| 132 | ): |
| 133 | temp_file = tempfile.mkstemp()[1] |
| 134 | current_device = torch.get_default_device() |
| 135 | vllm_config = VllmConfig( |
| 136 | model_config=model_config, cache_config=CacheConfig(block_size=16) |
| 137 | ) |
| 138 | with set_current_vllm_config(vllm_config=vllm_config): |
| 139 | init_distributed_environment( |
| 140 | world_size=1, |
| 141 | rank=0, |
| 142 | distributed_init_method=f"file://{temp_file}", |
| 143 | local_rank=0, |
| 144 | backend="nccl", |
| 145 | ) |
| 146 | initialize_model_parallel(tensor_model_parallel_size=1) |
| 147 | |
| 148 | with set_default_torch_dtype(model_config.dtype): |
| 149 | torch.set_default_device(current_platform.device_type) |
| 150 | model = model_cls(vllm_config=vllm_config) |
| 151 | torch.set_default_device(current_device) |
| 152 | yield model |
| 153 | |
| 154 | del model |
| 155 | cleanup_dist_env_and_memory() |
| 156 | |
| 157 | |
| 158 | @create_new_process_for_each_test() |
| 159 | @pytest.mark.parametrize("model_id", get_model_ids_to_test()) |
| 160 | def test_model_tensor_schema(model_id: str): |
| 161 | if model_id == "moonshotai/Kimi-K2.5": |
| 162 | # FIXME(Isotr0py): Fix Kimi-K2.5's offline inference about vision chunks. |
| 163 | pytest.skip( |
| 164 | "Kimi-K2.5's offline inference has issues about vision chunks. Fix later." |
| 165 | ) |
| 166 | |
| 167 | if ( |
| 168 | model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" |
| 169 | and not current_platform.is_cuda_alike() |
| 170 | ): |
| 171 | pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") |
| 172 | |
| 173 | if model_id == "zai-org/GLM-5.3-Flash" and (current_platform.is_xpu()): |
| 174 | pytest.skip("GLM-5.3-Flash is not supported on XPU") |
| 175 | |
| 176 | model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) |
| 177 | model_info.check_available_online(on_fail="skip") |
| 178 | model_info.check_transformers_version( |
| 179 | on_fail="skip", |
| 180 | check_max_version=False, |
| 181 | check_version_reason="vllm", |
| 182 | ) |
| 183 | |
| 184 | model_arch = next( |
| 185 | arch for arch, info in HF_EXAMPLE_MODELS.hf_models.items() if info == model_info |
| 186 | ) |
| 187 | |
| 188 | hf_overrides_fn = partial( |
| 189 | dummy_hf_overrides, |
| 190 | model_arch=model_arch, |
| 191 | exist_overrides=model_info.hf_overrides, |
| 192 | use_original_num_layers=getattr(model_info, "use_original_num_layers", False), |
| 193 | ) |
| 194 | |
| 195 | # ROCm: Detect if model uses AWQ quantization and set appropriate dtype |
| 196 | if "awq" in model_id.lower() and current_platform.is_rocm(): |
| 197 | dtype = "float16" |
| 198 | else: |
| 199 | dtype = model_info.dtype |
| 200 | |
| 201 | model_config = ModelConfig( |
| 202 | model_id, |
| 203 | tokenizer=model_info.tokenizer or model_id, |
| 204 | tokenizer_mode=model_info.tokenizer_mode, |
| 205 | revision=model_info.revision, |
| 206 | trust_remote_code=model_info.trust_remote_code, |
| 207 | hf_overrides=hf_overrides_fn, |
| 208 | skip_tokenizer_init=model_info.require_embed_inputs, |
| 209 | enable_prompt_embeds=model_info.require_embed_inputs, |
| 210 | enable_mm_embeds=model_info.require_embed_inputs, |
| 211 | enforce_eager=model_info.enforce_eager, |
| 212 | dtype=dtype, |
| 213 | ) |
| 214 | |
| 215 | model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config) |
| 216 | assert supports_multimodal(model_cls) |
| 217 | |
| 218 | factories = model_cls._processor_factory |
| 219 | |
| 220 | # Capture helpers return capture-buffer containers (e.g. |
| 221 | # EncoderCudaGraphCaptureInputs), not TensorSchema mm inputs. |
| 222 | capture_helpers = {"prepare_encoder_cudagraph_capture_inputs"} |
| 223 | |
| 224 | inputs_parse_methods = [] |
| 225 | for attr_name in dir(model_cls): |
| 226 | if attr_name in capture_helpers: |
| 227 | continue |
| 228 | attr = getattr(model_cls, attr_name) |
| 229 | if hasattr(attr, "__annotations__"): |
| 230 | return_type = attr.__annotations__.get("return", None) |
| 231 | if return_type is not None and "Input" in str(return_type): |
| 232 | inputs_parse_methods.append(attr_name) |
| 233 | |
| 234 | if not any(inputs_parse_methods): |
| 235 | pytest.skip(f"{model_arch} does not support tensor schema validation.") |
| 236 | |
| 237 | ctx = InputProcessingContext( |
| 238 | model_config, |
| 239 | tokenizer=cached_tokenizer_from_config(model_config), |
| 240 | ) |
| 241 | processing_info = factories.info(ctx) |
| 242 | supported_mm_limits = processing_info.get_supported_mm_limits() |
| 243 | limit_mm_per_prompt = { |
| 244 | modality: 3 if limit is None else limit |
| 245 | for modality, limit in supported_mm_limits.items() |
| 246 | } |
| 247 | |
| 248 | def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions: |
| 249 | if modality == "video": |
| 250 | return VideoDummyOptions(count=count) |
| 251 | if modality == "image": |
| 252 | return ImageDummyOptions(count=count) |
| 253 | if modality == "audio": |
| 254 | return AudioDummyOptions(count=count) |
| 255 | return BaseDummyOptions(count=count) |
| 256 | |
| 257 | model_config.get_multimodal_config().limit_per_prompt = MultiModalDummyOptions( |
| 258 | { |
| 259 | modality: _to_dummy_options(modality, count) |
| 260 | for modality, count in limit_mm_per_prompt.items() |
| 261 | } |
| 262 | ) |
| 263 | processor = factories.build_processor(ctx) |
| 264 | |
| 265 | with initialize_dummy_model(model_cls, model_config) as model: |
| 266 | for modality, _, mm_kwargs in create_batched_mm_kwargs(model_config, processor): |
| 267 | for method_name in inputs_parse_methods: |
| 268 | print( |
| 269 | f"Testing `{method_name}` with modality={modality} " |
| 270 | f"and mm_kwargs{list(mm_kwargs.keys())}" |
| 271 | ) |
| 272 | getattr(model, method_name)(modality=modality, **mm_kwargs) |