@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 165 lines 5.18 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 from vllm.multimodal.video_prune.vidcom2 import (
8 compute_retained_tokens_count,
9 compute_retention_mask,
10 )
11
12
13 def _fake_video_embeds(
14 num_frames: int,
15 rows: int,
16 cols: int,
17 hidden: int = 64,
18 seed: int = 0,
19 ) -> torch.Tensor:
20 """Deterministic fake ViT output with a distinct mean per frame."""
21 g = torch.Generator().manual_seed(seed)
22 frames = []
23 for f in range(num_frames):
24 base = torch.randn(hidden, generator=g) * (0.1 + 0.05 * f)
25 frames.append(
26 base[None, :].expand(rows * cols, hidden)
27 + 0.01 * torch.randn(rows * cols, hidden, generator=g)
28 )
29 return torch.cat(frames, dim=0)
30
31
32 @pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9])
33 @pytest.mark.parametrize("num_frames", [1, 4, 16])
34 def test_mask_shape_and_dtype(q: float, num_frames: int) -> None:
35 merge = 2
36 rows, cols = 6, 8
37 embeds = _fake_video_embeds(num_frames, rows, cols)
38 mask = compute_retention_mask(
39 embeds,
40 (num_frames, rows * merge, cols * merge),
41 spatial_merge_size=merge,
42 q=q,
43 )
44 assert mask.dtype == torch.bool
45 assert mask.shape == (num_frames * rows * cols,)
46
47
48 def test_retained_count_floors_at_one_token_per_frame() -> None:
49 """The global minimum is one token per frame (not a full first frame)."""
50 assert (
51 compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.999) == 4
52 )
53 assert (
54 compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.0)
55 == 48 * 4
56 )
57
58
59 def test_mask_reads_frame_budgets_once(monkeypatch: pytest.MonkeyPatch) -> None:
60 original_tolist = torch.Tensor.tolist
61 transfers: list[tuple[int, ...]] = []
62
63 def reject_item(tensor: torch.Tensor) -> None:
64 raise AssertionError(f"unexpected scalar read from shape {tuple(tensor.shape)}")
65
66 def record_tolist(tensor: torch.Tensor) -> list[int]:
67 transfers.append(tuple(tensor.shape))
68 return original_tolist(tensor)
69
70 monkeypatch.setattr(torch.Tensor, "item", reject_item)
71 monkeypatch.setattr(torch.Tensor, "tolist", record_tolist)
72 embeds = _fake_video_embeds(num_frames=4, rows=6, cols=8)
73
74 mask = compute_retention_mask(embeds, (4, 12, 16), spatial_merge_size=2, q=0.5)
75
76 assert transfers == [(4,)]
77 assert mask.shape == (4 * 6 * 8,)
78
79
80 @pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9])
81 @pytest.mark.parametrize("num_frames", [1, 4, 16])
82 def test_total_retained_matches_target(q: float, num_frames: int) -> None:
83 """Mask total must equal the placeholder-sizing helper."""
84 merge = 2
85 rows, cols = 6, 8
86 tpf = rows * cols
87 embeds = _fake_video_embeds(num_frames, rows, cols)
88 mask = compute_retention_mask(
89 embeds,
90 (num_frames, rows * merge, cols * merge),
91 spatial_merge_size=merge,
92 q=q,
93 )
94 expected = compute_retained_tokens_count(
95 tokens_per_frame=tpf, num_frames=num_frames, q=q
96 )
97 assert int(mask.sum().item()) == expected
98
99
100 def test_per_frame_min_one_when_budget_allows() -> None:
101 """No frame is fully dropped when the budget allows."""
102 merge = 2
103 rows, cols = 6, 8
104 num_frames = 8
105 embeds = _fake_video_embeds(num_frames, rows, cols)
106 mask = compute_retention_mask(
107 embeds,
108 (num_frames, rows * merge, cols * merge),
109 spatial_merge_size=merge,
110 q=0.25,
111 )
112 per_frame = mask.view(num_frames, rows * cols).sum(dim=1)
113 assert (per_frame >= 1).all(), f"zero-token frame detected: {per_frame.tolist()}"
114
115
116 def test_dynamic_per_frame_budget() -> None:
117 """A distinctive frame gets more retained tokens than bland ones."""
118 merge = 2
119 rows, cols = 6, 8
120 tpf = rows * cols
121 hidden = 64
122 torch.manual_seed(0)
123 bland = 0.01 * torch.randn(tpf, hidden)
124 frames = [torch.randn(tpf, hidden) * 1.0]
125 for _ in range(7):
126 frames.append(bland + 0.001 * torch.randn(tpf, hidden))
127 embeds = torch.cat(frames, dim=0)
128 mask = compute_retention_mask(
129 embeds,
130 (8, rows * merge, cols * merge),
131 spatial_merge_size=merge,
132 q=0.5,
133 )
134 per_frame = mask.view(8, tpf).sum(dim=1)
135 assert per_frame[0].item() > per_frame[1:].float().mean().item()
136
137
138 def test_empty_input_safe() -> None:
139 embeds = torch.zeros(0, 32)
140 mask = compute_retention_mask(embeds, (0, 0, 0), spatial_merge_size=2, q=0.25)
141 assert mask.numel() == 0
142
143
144 @pytest.mark.parametrize("q", [0.0, 0.25, 0.5, 0.75])
145 def test_first_frame_not_privileged(q: float) -> None:
146 """A bland first frame is not force-retained (unlike EVS)."""
147 merge = 2
148 rows, cols = 6, 8
149 tpf = rows * cols
150 torch.manual_seed(1)
151 bland = 0.01 * torch.randn(tpf, 64)
152 frames = [bland]
153 for f in range(7):
154 frames.append(torch.randn(tpf, 64) * (1.0 + 0.1 * f))
155 embeds = torch.cat(frames, dim=0)
156 mask = compute_retention_mask(
157 embeds,
158 (8, rows * merge, cols * merge),
159 spatial_merge_size=merge,
160 q=q,
161 )
162 per_frame = mask.view(8, tpf).sum(dim=1)
163 assert per_frame[0].item() <= tpf
164 if q > 0.0:
165 assert per_frame[0].item() < int(mask.sum().item())