@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 145 lines 3.5 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 import KeyeForConditionalGeneration
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) -> KeyeForConditionalGeneration:
38 model = object.__new__(KeyeForConditionalGeneration)
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 ) -> MultiModalFeatureSpec:
50 field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
51 return MultiModalFeatureSpec(
52 data=MultiModalKwargsItem(
53 {
54 field_name: MultiModalFieldElem(
55 data=torch.tensor(grid_thw),
56 field=None, # HACK.
57 ),
58 }
59 ),
60 modality=modality,
61 identifier="DUMMY",
62 mm_position=PlaceholderRange(offset=offset, length=length),
63 )
64
65
66 def test_get_mrope_input_positions_text_only():
67 model = make_model(DummyConfig())
68
69 positions, delta = model.get_mrope_input_positions(
70 input_tokens=[11, 12, 13, 14, 15],
71 mm_features=[],
72 )
73
74 expected = torch.tensor(
75 [
76 [0, 1, 2, 3, 4],
77 [0, 1, 2, 3, 4],
78 [0, 1, 2, 3, 4],
79 ]
80 )
81
82 assert torch.equal(positions, expected)
83 assert delta == 0
84
85
86 def test_get_mrope_input_positions_single_image():
87 model = make_model(DummyConfig())
88 mm_features = [
89 make_mm_feature(
90 modality="image",
91 offset=1,
92 length=4,
93 grid_thw=(1, 4, 4),
94 )
95 ]
96
97 positions, delta = model.get_mrope_input_positions(
98 input_tokens=[10, 20, 21, 22, 23, 30, 31],
99 mm_features=mm_features,
100 )
101
102 expected = torch.tensor(
103 [
104 [0, 1, 1, 1, 1, 3, 4],
105 [0, 1, 1, 2, 2, 3, 4],
106 [0, 1, 2, 1, 2, 3, 4],
107 ]
108 )
109
110 assert torch.equal(positions, expected)
111 assert delta == -2
112
113
114 def test_get_mrope_input_positions_interleaved_image_and_video():
115 model = make_model(DummyConfig())
116 mm_features = [
117 make_mm_feature(
118 modality="image",
119 offset=1,
120 length=4,
121 grid_thw=(1, 4, 4),
122 ),
123 make_mm_feature(
124 modality="video",
125 offset=7,
126 length=4,
127 grid_thw=[(2, 4, 2)],
128 ),
129 ]
130
131 positions, delta = model.get_mrope_input_positions(
132 input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 42, 43, 50, 51],
133 mm_features=mm_features,
134 )
135
136 expected = torch.tensor(
137 [
138 [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 7, 9, 10],
139 [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10],
140 [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 7, 9, 10],
141 ]
142 )
143
144 assert torch.equal(positions, expected)
145 assert delta == -2