@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 237 lines 7.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import dataclasses
4 import random
5 from dataclasses import dataclass
6
7 import pytest
8 import torch
9
10 from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
11 from vllm.multimodal.inputs import (
12 MultiModalFeatureSpec,
13 MultiModalFieldElem,
14 MultiModalKwargsItem,
15 PlaceholderRange,
16 )
17
18
19 @pytest.fixture(autouse=True, scope="module")
20 def _force_cpu_default_device():
21 # _get_mrope_input_positions returns CPU tensors (via torch.from_numpy).
22 # Ensure the default device is CPU so the rest of the test tensors match.
23 original = torch.get_default_device()
24 torch.set_default_device("cpu")
25 yield
26 torch.set_default_device(original)
27
28
29 IMAGE_TOKEN_ID = 999
30 VIDEO_TOKEN_ID = 888
31 VISION_START_TOKEN_ID = 777
32 VISION_END_TOKEN_ID = 778
33
34
35 @dataclass
36 class DummyVisionConfig:
37 spatial_merge_size: int = 1
38
39
40 @dataclass
41 class DummyConfig:
42 image_token_id: int = IMAGE_TOKEN_ID
43 video_token_id: int = VIDEO_TOKEN_ID
44 vision_start_token_id: int = VISION_START_TOKEN_ID
45 vision_end_token_id: int = VISION_END_TOKEN_ID
46 vision_config: DummyVisionConfig = dataclasses.field(
47 default_factory=DummyVisionConfig
48 )
49
50
51 def make_video_embedding(
52 t, h, w, interleave_text_tokens: tuple[int, int], video_pruning_rate: float = 0.0
53 ):
54 """Make a video embedding for a given video size and pruning rate.
55
56 Args:
57 t: Number of frames.
58 h: Number of rows.
59 w: Number of columns.
60 interleave_text_tokens: Tuple of minimum and maximum number of text tokens to
61 interleave with the video.
62 video_pruning_rate: Pruning rate for the video.
63
64 Returns:
65 Tuple of (unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask)
66
67 """
68 unpruned_tokens_sequence = []
69 population = list(range(1, 100))
70
71 for _ in range(t):
72 num_prefix_tokens = random.randint(
73 interleave_text_tokens[0], interleave_text_tokens[1]
74 )
75
76 prefix_tokens = random.choices(population, k=num_prefix_tokens)
77 vision_tokens = (
78 [VISION_START_TOKEN_ID] + [VIDEO_TOKEN_ID] * h * w + [VISION_END_TOKEN_ID]
79 )
80
81 unpruned_tokens_sequence.extend(prefix_tokens)
82 unpruned_tokens_sequence.extend(vision_tokens)
83
84 unpruned_tokens_sequence = torch.tensor(unpruned_tokens_sequence, dtype=torch.long)
85 video_token_mask = unpruned_tokens_sequence == VIDEO_TOKEN_ID
86
87 pruning_mask = torch.bernoulli(video_token_mask.float() * video_pruning_rate).bool() # type: ignore[attr-defined]
88 # Sanity check that we don't prune what should not be pruned.
89 assert not pruning_mask[~video_token_mask].any()
90
91 retention_mask = ~pruning_mask
92 pruned_tokens_sequence = unpruned_tokens_sequence[retention_mask]
93 return unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask
94
95
96 @pytest.mark.parametrize("spatial_merge_size", [1, 2])
97 @pytest.mark.parametrize("grid_thw", [[3, 8, 7], [128, 10, 12]])
98 @pytest.mark.parametrize("num_prefix_tokens", [1, 11])
99 @pytest.mark.parametrize("num_suffix_tokens", [0, 7])
100 @pytest.mark.parametrize("video_pruning_rate", [0, 0.25, 0.75])
101 @pytest.mark.parametrize("interleave_text_tokens", [(0, 0), (1, 4)])
102 def test_match_qwen3vl_mrope_evs_on(
103 spatial_merge_size: int,
104 num_prefix_tokens: int,
105 grid_thw: tuple[int, int, int],
106 num_suffix_tokens: int,
107 video_pruning_rate: float,
108 interleave_text_tokens: tuple[int, int],
109 ):
110 hf_config = DummyConfig()
111 hf_config.vision_config.spatial_merge_size = spatial_merge_size
112
113 t, h, w = grid_thw
114 population = list(range(1, 100))
115 prefix_tokens = random.choices(population, k=num_prefix_tokens)
116 suffix_tokens = random.choices(population, k=num_suffix_tokens)
117
118 video_tokens, video_tokens_pruned, retention_mask = make_video_embedding(
119 t,
120 h // spatial_merge_size,
121 w // spatial_merge_size,
122 interleave_text_tokens=interleave_text_tokens,
123 video_pruning_rate=video_pruning_rate,
124 )
125 assert len(video_tokens) == len(retention_mask)
126
127 input_tokens = prefix_tokens + video_tokens.tolist() + suffix_tokens
128 input_tokens_pruned = prefix_tokens + video_tokens_pruned.tolist() + suffix_tokens
129
130 whole_sequence_retention_mask = torch.cat(
131 [
132 torch.ones(len(prefix_tokens), dtype=torch.bool),
133 retention_mask,
134 torch.ones(len(suffix_tokens), dtype=torch.bool),
135 ],
136 dim=0,
137 )
138
139 # Build the GT mrope for unpruned input.
140 mm_feature = MultiModalFeatureSpec(
141 data=MultiModalKwargsItem(
142 {
143 "video_grid_thw": MultiModalFieldElem(
144 data=torch.tensor(grid_thw),
145 field=None, # HACK.
146 ),
147 }
148 ),
149 modality="video",
150 identifier="DUMMY",
151 mm_position=PlaceholderRange(offset=0, length=len(input_tokens)),
152 )
153 expected_mrope, _ = Qwen3VLForConditionalGeneration._get_mrope_input_positions(
154 input_tokens=input_tokens,
155 mm_features=[mm_feature],
156 config=hf_config,
157 )
158
159 # Compute mrope for a video-only media (unpruned).
160 mm_feature = MultiModalFeatureSpec(
161 data=MultiModalKwargsItem(
162 {
163 "video_grid_thw": MultiModalFieldElem(
164 data=torch.tensor(grid_thw),
165 field=None, # HACK.
166 ),
167 }
168 ),
169 modality="video",
170 identifier="DUMMY",
171 mm_position=PlaceholderRange(offset=0, length=video_tokens.numel()),
172 )
173 video_mrope, _ = Qwen3VLForConditionalGeneration._get_mrope_input_positions(
174 input_tokens=video_tokens.tolist(),
175 mm_features=[mm_feature],
176 config=hf_config,
177 )
178 video_mrope = video_mrope.permute(1, 0) # [N, 3]
179 hidden_size = 16
180
181 is_video_embed = torch.isin(
182 video_tokens_pruned, torch.tensor([VIDEO_TOKEN_ID], dtype=torch.long)
183 )
184
185 expanded_positions = torch.full(
186 (len(video_tokens_pruned), 5),
187 fill_value=-100,
188 device=video_mrope.device,
189 dtype=torch.long,
190 )
191 expanded_positions[is_video_embed, :3] = video_mrope[retention_mask][is_video_embed]
192 expanded_positions[~is_video_embed, :3] = video_mrope[retention_mask][
193 ~is_video_embed
194 ]
195
196 is_vision_start = video_tokens_pruned == VISION_START_TOKEN_ID
197 expanded_positions[..., 3] = is_vision_start
198 expanded_positions[..., 4] = is_video_embed
199
200 # Check that all positions were filled, since we initialized them as negative.
201 assert (expanded_positions >= 0).all()
202
203 video_embeddings = torch.empty(
204 (len(video_tokens_pruned), hidden_size), device=video_mrope.device
205 )
206
207 video_embeddings = torch.cat(
208 [
209 video_embeddings,
210 expanded_positions.float(),
211 ],
212 dim=1,
213 )
214 multimodal_embeddings = [video_embeddings]
215
216 expected_mrope_masked = expected_mrope[:, whole_sequence_retention_mask]
217
218 # Initialize computed_mrope with sequential positions for all prefix tokens
219 computed_mrope = torch.empty((3, len(input_tokens_pruned)), dtype=torch.long)
220 computed_mrope[:, 0 : len(prefix_tokens)] = expected_mrope[
221 :, 0 : len(prefix_tokens)
222 ]
223
224 # Paranoia check that computed_mrope is wrong.
225 assert not torch.equal(computed_mrope, expected_mrope_masked)
226
227 _, actual_mrope, _ = Qwen3VLForConditionalGeneration._recompute_mrope_positions(
228 input_ids=input_tokens_pruned,
229 multimodal_embeddings=multimodal_embeddings,
230 mrope_positions=computed_mrope,
231 num_computed_tokens=len(prefix_tokens),
232 vision_start_token_id=hf_config.vision_start_token_id,
233 image_token_id=hf_config.image_token_id,
234 video_token_id=hf_config.video_token_id,
235 )
236
237 assert torch.equal(actual_mrope, expected_mrope_masked)