| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Batch-invariance regression tests for vision-language models (VLMs).""" |
| 4 | |
| 5 | import contextlib |
| 6 | import os |
| 7 | |
| 8 | import numpy as np |
| 9 | import pytest |
| 10 | import torch |
| 11 | from PIL import Image |
| 12 | from utils import _extract_step_logprobs, skip_if_not_cuda |
| 13 | |
| 14 | from vllm import LLM, SamplingParams |
| 15 | from vllm.inputs import TextPrompt |
| 16 | |
| 17 | VLM_TEST_MODEL = os.getenv("VLLM_VLM_TEST_MODEL", "Qwen/Qwen3-VL-2B-Instruct") |
| 18 | |
| 19 | IMAGE_PROMPT = ( |
| 20 | "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" |
| 21 | "Describe this image.<|im_end|><|im_start|>assistant\n" |
| 22 | ) |
| 23 | VIDEO_PROMPT = ( |
| 24 | "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" |
| 25 | "Describe this video.<|im_end|><|im_start|>assistant\n" |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | def _make_image(seed: int, size: int) -> Image.Image: |
| 30 | rng = np.random.default_rng(seed) |
| 31 | return Image.fromarray(rng.integers(0, 255, (size, size, 3), dtype=np.uint8)) |
| 32 | |
| 33 | |
| 34 | def _make_inputs(input_type: str, num_reqs: int) -> list[TextPrompt]: |
| 35 | inputs: list[TextPrompt] = [] |
| 36 | for i in range(num_reqs): |
| 37 | if input_type == "image": |
| 38 | prompt = IMAGE_PROMPT |
| 39 | mm_data = {"image": _make_image(1000 + i, size=256)} |
| 40 | else: |
| 41 | frames = [_make_image(1000 + i + f, size=128) for f in range(4)] |
| 42 | prompt = VIDEO_PROMPT |
| 43 | mm_data = { |
| 44 | "video": ( |
| 45 | frames, |
| 46 | { |
| 47 | "total_num_frames": len(frames), |
| 48 | "fps": 1.0, |
| 49 | "frames_indices": list(range(len(frames))), |
| 50 | "do_sample_frames": False, |
| 51 | }, |
| 52 | ) |
| 53 | } |
| 54 | inputs.append({"prompt": prompt, "multi_modal_data": mm_data}) |
| 55 | return inputs |
| 56 | |
| 57 | |
| 58 | def _assert_batch_invariant(llm: LLM, inputs: list[TextPrompt], sampling) -> None: |
| 59 | bs1 = [ |
| 60 | _extract_step_logprobs(llm.generate([inp], sampling, use_tqdm=False)[0]) |
| 61 | for inp in inputs |
| 62 | ] |
| 63 | bsN_outputs = llm.generate(inputs, sampling, use_tqdm=False) |
| 64 | assert len(bs1) == len(inputs) |
| 65 | assert len(bsN_outputs) == len(inputs) |
| 66 | bsN = [_extract_step_logprobs(out) for out in bsN_outputs] |
| 67 | |
| 68 | failures = [] |
| 69 | for i, ((lp1, t1), (lpN, tN)) in enumerate(zip(bs1, bsN)): |
| 70 | if lp1 is None or lpN is None: |
| 71 | failures.append(f"req {i}: selected-token logprobs are unavailable") |
| 72 | continue |
| 73 | if len(t1) != sampling.max_tokens or len(tN) != sampling.max_tokens: |
| 74 | failures.append( |
| 75 | f"req {i}: expected {sampling.max_tokens} tokens, " |
| 76 | f"got bs1={len(t1)} bsN={len(tN)}" |
| 77 | ) |
| 78 | continue |
| 79 | if t1 != tN: |
| 80 | failures.append(f"req {i}: token mismatch bs1={t1} bsN={tN}") |
| 81 | elif not torch.equal(lp1, lpN): |
| 82 | d = (lp1 - lpN).abs().max().item() |
| 83 | failures.append(f"req {i}: logprob mismatch max_diff={d:.3e}") |
| 84 | |
| 85 | if failures: |
| 86 | pytest.fail( |
| 87 | f"Batch invariance violated for {len(failures)}/{len(inputs)} " |
| 88 | f"requests:\n" + "\n".join(failures) |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | @skip_if_not_cuda |
| 93 | @pytest.mark.parametrize("input_type", ["image", "video"]) |
| 94 | @pytest.mark.parametrize("mm_encoder_attn_backend", ["FLASH_ATTN", "TORCH_SDPA"]) |
| 95 | def test_vlm_batch_invariance_bs1_vs_bsN( |
| 96 | input_type: str, |
| 97 | mm_encoder_attn_backend: str, |
| 98 | monkeypatch: pytest.MonkeyPatch, |
| 99 | ): |
| 100 | monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") |
| 101 | _run_vlm_batch_invariance(input_type, mm_encoder_attn_backend, enforce_eager=True) |
| 102 | |
| 103 | |
| 104 | @skip_if_not_cuda |
| 105 | def test_vlm_batch_invariance_default_execution(monkeypatch: pytest.MonkeyPatch): |
| 106 | monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") |
| 107 | _run_vlm_batch_invariance("image", "FLASH_ATTN", enforce_eager=False) |
| 108 | |
| 109 | |
| 110 | def _run_vlm_batch_invariance( |
| 111 | input_type: str, mm_encoder_attn_backend: str, enforce_eager: bool |
| 112 | ) -> None: |
| 113 | inputs = _make_inputs(input_type, num_reqs=4) |
| 114 | sampling = SamplingParams( |
| 115 | temperature=0.0, |
| 116 | max_tokens=8, |
| 117 | seed=1234, |
| 118 | ignore_eos=True, |
| 119 | logprobs=5, |
| 120 | ) |
| 121 | |
| 122 | llm = LLM( |
| 123 | model=VLM_TEST_MODEL, |
| 124 | dtype="bfloat16", |
| 125 | enforce_eager=enforce_eager, |
| 126 | max_num_seqs=16, |
| 127 | max_model_len=4096, |
| 128 | max_num_batched_tokens=2048, |
| 129 | gpu_memory_utilization=0.85, |
| 130 | enable_prefix_caching=False, |
| 131 | mm_processor_cache_gb=0, |
| 132 | mm_encoder_attn_backend=mm_encoder_attn_backend, |
| 133 | ) |
| 134 | try: |
| 135 | _assert_batch_invariant(llm, inputs, sampling) |
| 136 | finally: |
| 137 | with contextlib.suppress(Exception): |
| 138 | llm.llm_engine.engine_core.shutdown() |