@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 145 lines 3.57 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass, field
4
5 import pytest
6 import torch
7
8 from vllm.model_executor.models.keye_vl1_5 import KeyeVL1_5ForConditionalGeneration
9 from vllm.multimodal.inputs import (
10 MultiModalFeatureSpec,
11 MultiModalFieldElem,
12 MultiModalKwargsItem,
13 PlaceholderRange,
14 )
15
16 pytestmark = pytest.mark.skip_global_cleanup
17
18
19 @pytest.fixture(autouse=True, scope="module")
20 def _force_cpu_default_device():
21 original = torch.get_default_device()
22 torch.set_default_device("cpu")
23 yield
24 torch.set_default_device(original)
25
26
27 @dataclass
28 class DummyVisionConfig:
29 spatial_merge_size: int = 2
30
31
32 @dataclass
33 class DummyConfig:
34 vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
35
36
37 def make_model(config: DummyConfig) -> KeyeVL1_5ForConditionalGeneration:
38 model = object.__new__(KeyeVL1_5ForConditionalGeneration)
39 model.config = config
40 return model
41
42
43 def make_mm_feature(
44 *,
45 modality: str,
46 offset: int,
47 length: int,
48 grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
49 is_embed: list[bool] | None = None,
50 ) -> MultiModalFeatureSpec:
51 field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
52 return MultiModalFeatureSpec(
53 data=MultiModalKwargsItem(
54 {
55 field_name: MultiModalFieldElem(
56 data=torch.tensor(grid_thw),
57 field=None, # HACK.
58 ),
59 }
60 ),
61 modality=modality,
62 identifier="DUMMY",
63 mm_position=PlaceholderRange(
64 offset=offset,
65 length=length,
66 is_embed=None if is_embed is None else torch.tensor(is_embed),
67 ),
68 )
69
70
71 def test_get_mrope_input_positions_text_only():
72 model = make_model(DummyConfig())
73
74 positions, delta = model.get_mrope_input_positions(
75 input_tokens=[11, 12, 13, 14, 15],
76 mm_features=[],
77 )
78
79 expected = torch.tensor(
80 [
81 [0, 1, 2, 3, 4],
82 [0, 1, 2, 3, 4],
83 [0, 1, 2, 3, 4],
84 ]
85 )
86
87 assert torch.equal(positions, expected)
88 assert delta == 0
89
90
91 def test_get_mrope_input_positions_single_image():
92 model = make_model(DummyConfig())
93 mm_features = [
94 make_mm_feature(
95 modality="image",
96 offset=1,
97 length=4,
98 grid_thw=(1, 4, 4),
99 )
100 ]
101
102 positions, delta = model.get_mrope_input_positions(
103 input_tokens=[10, 20, 21, 22, 23, 30, 31],
104 mm_features=mm_features,
105 )
106
107 expected = torch.tensor(
108 [
109 [0, 1, 1, 1, 1, 3, 4],
110 [0, 1, 1, 2, 2, 3, 4],
111 [0, 1, 2, 1, 2, 3, 4],
112 ]
113 )
114
115 assert torch.equal(positions, expected)
116 assert delta == -2
117
118
119 def test_get_mrope_input_positions_video_uses_embed_ranges():
120 model = make_model(DummyConfig())
121 mm_features = [
122 make_mm_feature(
123 modality="video",
124 offset=1,
125 length=8,
126 grid_thw=[(2, 4, 2)],
127 is_embed=[False, False, True, True, False, False, True, True],
128 )
129 ]
130
131 positions, delta = model.get_mrope_input_positions(
132 input_tokens=[10, 101, 102, 20, 21, 103, 104, 30, 31, 40, 41],
133 mm_features=mm_features,
134 )
135
136 expected = torch.tensor(
137 [
138 [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
139 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
140 [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
141 ]
142 )
143
144 assert torch.equal(positions, expected)
145 assert delta == 0