| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import NamedTuple |
| 4 | |
| 5 | import pytest |
| 6 | from PIL import Image |
| 7 | |
| 8 | from vllm import LLM, SamplingParams |
| 9 | from vllm.assets.image import ImageAsset |
| 10 | from vllm.config import AttentionConfig, KVTransferConfig |
| 11 | from vllm.multimodal.utils import encode_image_url |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | MODEL_NAME = "RedHatAI/Qwen2.5-VL-3B-Instruct-quantized.w8a8" |
| 15 | |
| 16 | SAMPLING_PARAMS = SamplingParams(temperature=0.0, top_k=1, max_tokens=128) |
| 17 | |
| 18 | TEXT_PROMPTS = [ |
| 19 | "What's in the image(s)? Around 30 words. What's special in 2nd image?", |
| 20 | "The future of AI is", |
| 21 | ] |
| 22 | |
| 23 | |
| 24 | class InputCase(NamedTuple): |
| 25 | text: str |
| 26 | img: list[Image] |
| 27 | expected_len: int |
| 28 | info: str |
| 29 | |
| 30 | |
| 31 | def _check_path_len(path): |
| 32 | """Return the latest length in path.""" |
| 33 | return len(list(path.iterdir())) |
| 34 | |
| 35 | |
| 36 | def _list_path(path): |
| 37 | """Return the list of foldername (hashes generated) under the path.""" |
| 38 | return list(path.iterdir()) |
| 39 | |
| 40 | |
| 41 | def run_test( |
| 42 | tmp_path, |
| 43 | processor, |
| 44 | llm: LLM, |
| 45 | question: str, |
| 46 | image_urls: list[Image], |
| 47 | expected_len: int, |
| 48 | info: str, |
| 49 | ): |
| 50 | """One individual test to process the prompt and output base on 1 set of input |
| 51 | Then check if the length in the storage path matches the expected length |
| 52 | `info` introduces details or purpose of the individual test |
| 53 | """ |
| 54 | print(f"***info: {info}***") |
| 55 | print(f"**Expected storage path length after llm generate: {expected_len}**") |
| 56 | process_prompt(processor, llm, question, image_urls) |
| 57 | |
| 58 | print(f"Path matched expected length: {_check_path_len(tmp_path)}") |
| 59 | print(f"Hashes under the storage path: {_list_path(tmp_path)}") |
| 60 | |
| 61 | assert _check_path_len(tmp_path) == expected_len, ( |
| 62 | f"Expect storage path length {expected_len} ;", |
| 63 | f"but end up {_check_path_len(tmp_path)} instead. ", |
| 64 | f"Info: {info}", |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | def process_prompt(processor, llm: LLM, question: str, image_urls: list[Image]): |
| 69 | """Form the prompt based on the text and image input, then llm generate output.""" |
| 70 | placeholders = [ |
| 71 | { |
| 72 | "type": "image_url", |
| 73 | "image_url": {"url": encode_image_url(image_pil)}, |
| 74 | } |
| 75 | for image_pil in image_urls |
| 76 | ] |
| 77 | |
| 78 | messages = [ |
| 79 | {"role": "system", "content": "You are a helpful assistant."}, |
| 80 | { |
| 81 | "role": "user", |
| 82 | "content": [ |
| 83 | *placeholders, |
| 84 | {"type": "text", "text": question}, |
| 85 | ], |
| 86 | }, |
| 87 | ] |
| 88 | |
| 89 | prompt = processor.apply_chat_template( |
| 90 | messages, tokenize=False, add_generation_prompt=True |
| 91 | ) |
| 92 | |
| 93 | outputs = llm.generate( |
| 94 | { |
| 95 | "prompt": prompt, |
| 96 | **({"multi_modal_data": {"image": [*image_urls]}} if image_urls else {}), |
| 97 | }, |
| 98 | sampling_params=SAMPLING_PARAMS, |
| 99 | ) |
| 100 | |
| 101 | print("-" * 50) |
| 102 | print("Output:") |
| 103 | for o in outputs: |
| 104 | generated_text = o.outputs[0].text |
| 105 | print(generated_text) |
| 106 | print("-" * 50) |
| 107 | |
| 108 | |
| 109 | @pytest.mark.parametrize( |
| 110 | "attn_backend", |
| 111 | ( |
| 112 | ["FLASH_ATTN", "TRITON_ATTN"] |
| 113 | if current_platform.is_cuda() |
| 114 | else ["TRITON_ATTN"] |
| 115 | if current_platform.is_rocm() or current_platform.is_xpu() |
| 116 | else [] |
| 117 | ), |
| 118 | ) |
| 119 | def test_shared_storage_connector_hashes(tmp_path, attn_backend): |
| 120 | """Tests that ExampleConnector saves KV to the storage locations |
| 121 | with proper hashes; that are unique for inputs with identical text but |
| 122 | different images (same size), or same multiple images but different orders. |
| 123 | """ |
| 124 | # Using tmp_path as the storage path to store KV |
| 125 | print(f"KV storage path at: {str(tmp_path)}") |
| 126 | |
| 127 | # don't put this import at the top level |
| 128 | # it will call torch.accelerator.device_count() |
| 129 | from transformers import AutoProcessor |
| 130 | |
| 131 | # Create processor to handle the chat prompt |
| 132 | processor = AutoProcessor.from_pretrained(MODEL_NAME) |
| 133 | |
| 134 | # Prepare images for the tests |
| 135 | # Resize to the same size to check hashes correctness |
| 136 | image_1 = ImageAsset("stop_sign").pil_image.resize((1280, 720)) |
| 137 | image_2 = ImageAsset("cherry_blossom").pil_image.resize((1280, 720)) |
| 138 | |
| 139 | # Make sure that they are not the same picture |
| 140 | assert image_1 != image_2, "The images should not be identical" |
| 141 | |
| 142 | # Create the LLM instance |
| 143 | llm = LLM( |
| 144 | model=MODEL_NAME, |
| 145 | max_model_len=8192, |
| 146 | max_num_seqs=1, |
| 147 | gpu_memory_utilization=0.4, |
| 148 | attention_config=AttentionConfig(backend=attn_backend), |
| 149 | enforce_eager=True, |
| 150 | kv_transfer_config=KVTransferConfig( |
| 151 | kv_connector="ExampleConnector", |
| 152 | kv_role="kv_both", |
| 153 | kv_connector_extra_config={"shared_storage_path": str(tmp_path)}, |
| 154 | ), |
| 155 | limit_mm_per_prompt={"image": 2}, |
| 156 | ) |
| 157 | |
| 158 | # Prepare the input cases |
| 159 | input_cases = [ |
| 160 | InputCase( |
| 161 | text=TEXT_PROMPTS[0], |
| 162 | img=[image_1], |
| 163 | expected_len=1, |
| 164 | info="image_1 single input the first time.", |
| 165 | ), |
| 166 | InputCase( |
| 167 | text=TEXT_PROMPTS[0], |
| 168 | img=[image_2], |
| 169 | expected_len=2, |
| 170 | info=( |
| 171 | "image_2 single input the first time. " |
| 172 | "It is in same pixel size with image_1, yet it " |
| 173 | "should be able to form a new unique hash." |
| 174 | ), |
| 175 | ), |
| 176 | InputCase( |
| 177 | text=TEXT_PROMPTS[0], |
| 178 | img=[image_1], |
| 179 | expected_len=2, |
| 180 | info=( |
| 181 | "image_1 single input the 2nd time. " |
| 182 | "It should not form another new hash." |
| 183 | ), |
| 184 | ), |
| 185 | InputCase( |
| 186 | text=TEXT_PROMPTS[0], |
| 187 | img=[image_2], |
| 188 | expected_len=2, |
| 189 | info=( |
| 190 | "image_2 single input the 2nd time. " |
| 191 | "It should not form another new hash." |
| 192 | ), |
| 193 | ), |
| 194 | InputCase( |
| 195 | text=TEXT_PROMPTS[0], |
| 196 | img=[image_1, image_2], |
| 197 | expected_len=3, |
| 198 | info="image_1 with image_2 input the first time.", |
| 199 | ), |
| 200 | InputCase( |
| 201 | text=TEXT_PROMPTS[0], |
| 202 | img=[image_2, image_1], |
| 203 | expected_len=4, |
| 204 | info="The image order is swapped. Should form new hash.", |
| 205 | ), |
| 206 | InputCase( |
| 207 | text=TEXT_PROMPTS[0], |
| 208 | img=[image_1, image_2], |
| 209 | expected_len=4, |
| 210 | info=( |
| 211 | "[image_1, image_2] input the 2nd time. " |
| 212 | "It should not form another new hash." |
| 213 | ), |
| 214 | ), |
| 215 | InputCase( |
| 216 | text=TEXT_PROMPTS[0], |
| 217 | img=[image_2, image_1], |
| 218 | expected_len=4, |
| 219 | info=( |
| 220 | "[image_2, image_1] input the 2nd time. " |
| 221 | "It should not form another new hash." |
| 222 | ), |
| 223 | ), |
| 224 | InputCase( |
| 225 | text=TEXT_PROMPTS[0], |
| 226 | img=[], |
| 227 | expected_len=5, |
| 228 | info="Pure text input test as a case-control", |
| 229 | ), |
| 230 | InputCase( |
| 231 | text=TEXT_PROMPTS[0], |
| 232 | img=[], |
| 233 | expected_len=5, |
| 234 | info="Identical pure text input as a case-control", |
| 235 | ), |
| 236 | InputCase( |
| 237 | text=TEXT_PROMPTS[1], |
| 238 | img=[], |
| 239 | expected_len=6, |
| 240 | info="Another pure text input as a case-control", |
| 241 | ), |
| 242 | ] |
| 243 | |
| 244 | # Run tests |
| 245 | for case_id, (text, img, expected_len, info) in enumerate(input_cases): |
| 246 | print("\n", "=" * 25, f"Below running input case: {case_id}", "=" * 25) |
| 247 | run_test(tmp_path, processor, llm, text, img, expected_len, info) |
| 248 | |
| 249 | print("All tests passed successfully!") |