| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.model_executor.models.llava_next_video import ( |
| 9 | LlavaNextVideoForConditionalGeneration, |
| 10 | ) |
| 11 | from vllm.model_executor.models.vision import get_vision_encoder_info |
| 12 | |
| 13 | from ...utils import build_model_context |
| 14 | |
| 15 | |
| 16 | class _StubModel: |
| 17 | """Carries only the attributes the token-count method reads from |
| 18 | `self`, so the real methods can be exercised without constructing |
| 19 | the full `nn.Module` (vision tower, language model, etc.).""" |
| 20 | |
| 21 | patch_grid_length: int |
| 22 | pooled_grid_length: int |
| 23 | |
| 24 | |
| 25 | get_mm_lora_token_counts = ( |
| 26 | LlavaNextVideoForConditionalGeneration.get_mm_lora_token_counts |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | @pytest.mark.parametrize("model_id", ["llava-hf/LLaVA-NeXT-Video-7B-hf"]) |
| 31 | def test_num_mm_tokens_match_real_config(model_id): |
| 32 | """The stored grid lengths must match what `__init__` derives from |
| 33 | the real HF config, and the LoRA token counts must preserve the |
| 34 | frame-level scaling exactly.""" |
| 35 | ctx = build_model_context(model_id, limit_mm_per_prompt={"video": 1}) |
| 36 | config = ctx.model_config.hf_config |
| 37 | |
| 38 | vision_encoder_info = get_vision_encoder_info(config) |
| 39 | patch_grid_length = vision_encoder_info.get_patch_grid_length() |
| 40 | pooled_grid_length = math.ceil(patch_grid_length / config.spatial_pool_stride) |
| 41 | |
| 42 | stub = _StubModel() |
| 43 | stub.patch_grid_length = patch_grid_length |
| 44 | stub.pooled_grid_length = pooled_grid_length |
| 45 | |
| 46 | for num_frames in (1, 2, 8, 16, 32): |
| 47 | num_video_tokens = num_frames * pooled_grid_length**2 |
| 48 | |
| 49 | encoder_tokens, connector_tokens = get_mm_lora_token_counts( |
| 50 | stub, |
| 51 | modality="video", |
| 52 | mm_kwargs=None, |
| 53 | num_mm_embeds=num_video_tokens, |
| 54 | ) |
| 55 | assert encoder_tokens == num_frames * patch_grid_length**2 |
| 56 | assert connector_tokens == num_video_tokens |
| 57 | |
| 58 | |
| 59 | @pytest.mark.parametrize( |
| 60 | ("patch_grid_length", "pooled_grid_length", "num_frames"), |
| 61 | [ |
| 62 | (24, 12, 1), # llava-hf/LLaVA-NeXT-Video-7B-hf: 336 / 14, stride 2 |
| 63 | (24, 12, 16), |
| 64 | (27, 14, 5), # non-power-of-2 pooled grid (ceil rounding) |
| 65 | (16, 8, 6), |
| 66 | ], |
| 67 | ) |
| 68 | def test_num_mm_tokens_roundtrip(patch_grid_length, pooled_grid_length, num_frames): |
| 69 | stub = _StubModel() |
| 70 | stub.patch_grid_length = patch_grid_length |
| 71 | stub.pooled_grid_length = pooled_grid_length |
| 72 | |
| 73 | num_video_tokens = num_frames * pooled_grid_length**2 |
| 74 | |
| 75 | encoder_tokens, connector_tokens = get_mm_lora_token_counts( |
| 76 | stub, |
| 77 | modality="video", |
| 78 | mm_kwargs=None, |
| 79 | num_mm_embeds=num_video_tokens, |
| 80 | ) |
| 81 | assert encoder_tokens == num_frames * patch_grid_length**2 |
| 82 | assert connector_tokens == num_video_tokens |
| 83 | |
| 84 | |
| 85 | def test_num_mm_tokens_zero(): |
| 86 | stub = _StubModel() |
| 87 | stub.patch_grid_length = 24 |
| 88 | stub.pooled_grid_length = 12 |
| 89 | |
| 90 | assert get_mm_lora_token_counts( |
| 91 | stub, |
| 92 | modality="video", |
| 93 | mm_kwargs=None, |
| 94 | num_mm_embeds=0, |
| 95 | ) == (0, 0) |