[Perf] Read VidCom2 frame budgets once (#55331)
Signed-off-by: levius <2114377220@qq.com>
Sun committed
Sep 5, 2026 at 04:11 UTC
c81ace18596f75660bbc04efc2075a5f17a76791
2 files changed
+24
-5
tests/multimodal/test_vidcom2.py
+21
@@ -56,6 +56,27 @@ def test_retained_count_floors_at_one_token_per_frame() -> None:
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:
vllm/multimodal/video_prune/vidcom2.py
+3
-5
@@ -88,18 +88,16 @@ def compute_retention_mask(
88
ks = (scales * tokens_per_frame).round().long().clamp(min=1, max=tokens_per_frame)
89
90
# 4. Retain the smallest-similarity tokens per frame.
91
+ budgets = ks.tolist()
92
mask_2d = torch.zeros(T, tokens_per_frame, dtype=torch.bool, device=device)
92
- for i in range(T):
93
- k_i = int(ks[i].item())
94
- if k_i <= 0:
95
- continue
93
+ for i, k_i in enumerate(budgets):
94
_, idx = torch.topk(similarity[i], k=k_i, largest=False, sorted=False)
95
mask_2d[i].scatter_(0, idx, True)
96
97
# 5. Reconcile rounding/clamp drift to the exact target count by score.
98
flat_mask = mask_2d.view(-1)
99
flat_sim = similarity.view(-1)
102
- current = int(flat_mask.sum().item())
100
+ current = sum(budgets)
101
if current > target_retained:
102
drop_n = current - target_retained
103
retained_idx = flat_mask.nonzero(as_tuple=False).squeeze(-1)