| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for mllama's multimodal preprocessing and profiling.""" |
| 4 | |
| 5 | import pytest |
| 6 | from torch import prod |
| 7 | from transformers import Llama4Config |
| 8 | |
| 9 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 10 | |
| 11 | from ...utils import build_model_context |
| 12 | |
| 13 | |
| 14 | @pytest.mark.parametrize("model_id", ["meta-llama/Llama-Guard-4-12B"]) |
| 15 | @pytest.mark.parametrize("max_model_len", [4096, 8192, 25600, 131072]) |
| 16 | def test_profiling(model_id: str, max_model_len: int): |
| 17 | model_config_kwargs = { |
| 18 | "max_model_len": max_model_len, |
| 19 | } |
| 20 | mm_counts = {"image": 1} |
| 21 | ctx = build_model_context( |
| 22 | model_id, |
| 23 | model_config_kwargs=model_config_kwargs, |
| 24 | limit_mm_per_prompt=mm_counts, |
| 25 | ) |
| 26 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 27 | |
| 28 | mm_inputs = processor.get_dummy_mm_inputs(mm_counts) |
| 29 | |
| 30 | hf_config = ctx.get_hf_config(Llama4Config) |
| 31 | image_size = hf_config.vision_config.image_size |
| 32 | patch_size = hf_config.vision_config.patch_size |
| 33 | downsample_ratio = int( |
| 34 | round(1.0 / (hf_config.vision_config.pixel_shuffle_ratio**2)) |
| 35 | ) |
| 36 | tokens_per_patch = ((image_size // patch_size) ** 2) // downsample_ratio |
| 37 | |
| 38 | mm_data = mm_inputs["mm_kwargs"].get_data() |
| 39 | chunks_per_image = prod(mm_data["patches_per_image"]) |
| 40 | total_num_patches = chunks_per_image * tokens_per_patch |
| 41 | num_tiles = ( |
| 42 | mm_data["aspect_ratios"][0][0] * mm_data["aspect_ratios"][0][1] |
| 43 | ) # x-y separator tokens |
| 44 | total_tokens = ( |
| 45 | total_num_patches.item() + num_tiles.item() + 3 |
| 46 | ) # image start, image, image end |
| 47 | |
| 48 | assert total_num_patches == sum( |
| 49 | item.get_num_embeds() for item in mm_inputs["mm_placeholders"]["image"] |
| 50 | ) |
| 51 | assert total_tokens == sum( |
| 52 | placeholder.length for placeholder in mm_inputs["mm_placeholders"]["image"] |
| 53 | ) |