[Perf] Read VidCom2 frame budgets once (#55331)

Signed-off-by: levius <2114377220@qq.com>

tests/multimodal/test_vidcom2.py

@@ -56,6 +56,27 @@ def test_retained_count_floors_at_one_token_per_frame() -> None:
     )
 
 
+def test_mask_reads_frame_budgets_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    original_tolist = torch.Tensor.tolist
+    transfers: list[tuple[int, ...]] = []
+
+    def reject_item(tensor: torch.Tensor) -> None:
+        raise AssertionError(f"unexpected scalar read from shape {tuple(tensor.shape)}")
+
+    def record_tolist(tensor: torch.Tensor) -> list[int]:
+        transfers.append(tuple(tensor.shape))
+        return original_tolist(tensor)
+
+    monkeypatch.setattr(torch.Tensor, "item", reject_item)
+    monkeypatch.setattr(torch.Tensor, "tolist", record_tolist)
+    embeds = _fake_video_embeds(num_frames=4, rows=6, cols=8)
+
+    mask = compute_retention_mask(embeds, (4, 12, 16), spatial_merge_size=2, q=0.5)
+
+    assert transfers == [(4,)]
+    assert mask.shape == (4 * 6 * 8,)
+
+
 @pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9])
 @pytest.mark.parametrize("num_frames", [1, 4, 16])
 def test_total_retained_matches_target(q: float, num_frames: int) -> None:

vllm/multimodal/video_prune/vidcom2.py

@@ -88,18 +88,16 @@ def compute_retention_mask(
     ks = (scales * tokens_per_frame).round().long().clamp(min=1, max=tokens_per_frame)
 
     # 4. Retain the smallest-similarity tokens per frame.
+    budgets = ks.tolist()
     mask_2d = torch.zeros(T, tokens_per_frame, dtype=torch.bool, device=device)
-    for i in range(T):
-        k_i = int(ks[i].item())
-        if k_i <= 0:
-            continue
+    for i, k_i in enumerate(budgets):
         _, idx = torch.topk(similarity[i], k=k_i, largest=False, sorted=False)
         mask_2d[i].scatter_(0, idx, True)
 
     # 5. Reconcile rounding/clamp drift to the exact target count by score.
     flat_mask = mask_2d.view(-1)
     flat_sim = similarity.view(-1)
-    current = int(flat_mask.sum().item())
+    current = sum(budgets)
     if current > target_retained:
         drop_n = current - target_retained
         retained_idx = flat_mask.nonzero(as_tuple=False).squeeze(-1)