[Perf] Read VidCom2 frame budgets once

[Levius-Fubuki]

Purpose

VidCom2 currently reads each dynamic per-frame Top-K budget with
ks[i].item(), then reads the selected-token count with another scalar
reduction. For T frames this introduces T + 1 host-visible scalar reads in
the token-selection path.

This PR copies the small budget vector to a Python list once, preserves the
original one-dimensional unsorted torch.topk operation for every frame, and
reuses the budget sum during count reconciliation:

before: T per-frame item() reads + 1 mask.sum().item() read
after:  1 ks.tolist() transfer

The budgets are already clamped to at least one, so the previous non-positive
budget branch was unreachable. The change keeps one portable PyTorch path with
no custom kernel, architecture gate, device branch, or device-specific tuning.

This is not duplicating existing work. The following open-PR searches returned
no matching VidCom2 optimization: VidCom2 topk, VidCom2 frame budget, and
video prune synchronization.

Test Plan

Regression test

.venv/bin/python -m pytest -q tests/multimodal/test_vidcom2.py

The new test runs the complete mask function, rejects any per-element
Tensor.item() read, and verifies exactly one vector budget transfer. Existing
tests cover retained counts, dynamic frame budgets, empty input, and first-frame
behavior.

Lint and pre-commit

pre-commit run --files \
  tests/multimodal/test_vidcom2.py \
  vllm/multimodal/video_prune/vidcom2.py

RTX 4090 complete-function benchmark

python bench_vllm_vidcom2_single_sync.py \
  --baseline-module ../vllm-vidcom2-single-sync-baseline/vllm/multimodal/video_prune/vidcom2.py \
  --candidate-module vllm/multimodal/video_prune/vidcom2.py \
  --dtype float16 bfloat16 float32 \
  --frames 1 2 4 8 16 32 64 128 \
  --tokens-per-frame 48 196 576 \
  --q 0.25 0.5 0.75 0.9 \
  --warmups 25 --iterations 100 \
  --output full-matrix.csv

Hardware/software: NVIDIA RTX 4090, driver 580.76.05, CUDA 13.2,
PyTorch 2.13.0+cu132. Baseline commit: c615b1fd.

dtype	T	Tokens/frame	q	Old p20/p50/p80 (us)	New p20/p50/p80 (us)	p50 speedup
BF16	1	48	0.50	781.79 / 784.24 / 790.75	750.84 / 754.55 / 761.13	1.04x
BF16	2	48	0.50	821.88 / 824.39 / 831.62	775.74 / 777.62 / 784.39	1.06x
FP16	8	196	0.75	1134.89 / 1138.19 / 1146.59	1024.67 / 1027.04 / 1034.63	1.11x
BF16	32	196	0.50	1996.89 / 2005.31 / 2022.69	1589.97 / 1595.34 / 1604.75	1.26x
FP32	64	576	0.25	3228.88 / 3359.86 / 3497.51	2202.08 / 2216.91 / 2252.41	1.52x
BF16	128	576	0.75	5158.44 / 5176.38 / 5210.99	3582.24 / 3592.81 / 3604.61	1.44x

Across all 288 configurations, the raw geometric-mean speedup was 1.181x and
the median was 1.142x. All masks matched bit-for-bit. One T=1 sample was a raw
timing outlier; five alternating-order trials with 500 iterations reproduced a
1.041x median speedup (1.037–1.041x), leaving no stable regression beyond 3%.

The profiler confirms that Top-K and scatter counts are unchanged while the
per-frame scalar reads disappear:

T	item calls, old → new	CUDA memcpy events, old → new	Top-K calls, old/new	Scatter calls, old/new
2	60 → 0	60 → 20	60 / 60	40 / 40
8	180 → 0	220 → 60	200 / 200	160 / 160
32	660 → 0	700 → 60	680 / 680	640 / 640

Each profiler row contains 20 complete mask calls. Other reconciliation-path
copies remain; the removed difference is exactly 20 * T per-frame budget
reads.

The standalone benchmark, raw CSV, profiler output, outlier recheck, and paired
model-evaluation output are available in the
benchmark artifact Gist.

Qwen3-VL paired evaluation

The offline evaluation used Qwen/Qwen3-VL-2B-Instruct revision
89644892e4d85e24eaac8bacfd4f463576704203. Baseline and candidate requests
were interleaved in one process, prefix caching was disabled, and both mask
functions were also run on the same visual embeddings.

Source frames	Visual grid	Mask result	Operator old → new	E2E old → new
8	[4, 8, 8]	exact, 0/64 differ	1.061 → 0.836 ms (1.27x)	647.60 → 645.29 ms (1.004x)
32	[16, 4, 4]	exact, 0/256 differ	1.549 → 1.196 ms (1.29x)	644.90 → 644.57 ms (1.001x)

All ten timed baseline/candidate outputs had identical generated token IDs,
text, and cumulative log probability.

Test Result

33 passed

ruff-format, ruff-check, typo, SPDX-header, forbidden-import,
new-torch.cuda API, and git diff --check checks passed locally.

Limitations

The broad performance matrix uses synthetic post-ViT embeddings on one RTX
4090. The two Qwen3-VL cases validate output stability but do not establish an
end-to-end throughput improvement; the optimized helper is a small fraction of
total inference latency. No performance claim is made for other architectures
or backends.

AI assistance

AI assistance was used during implementation, testing, benchmarking, model
evaluation, and PR drafting. The human contributor reviewed and takes
responsibility for the change and evidence.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: abb4f15b-2cc7-44ab-88c5-6dcfcd1a77b4

📥 Commits

Reviewing files that changed from the base of the PR and between 8b6de0e and 361d811.

📒 Files selected for processing (2)

• tests/multimodal/test_vidcom2.py
• vllm/multimodal/video_prune/vidcom2.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Performance

  • Improved video frame budget processing by reading all budgets in a single transfer, reducing repeated scalar reads.
  • Streamlined retention-mask reconciliation while preserving expected mask results.

• Tests

  • Added coverage verifying efficient frame-budget reads and the expected retention-mask shape.

Walkthrough

The retention mask computation now materializes frame budgets with one .tolist() call. A test verifies that no scalar .item() reads occur and that the output mask has the expected shape.

Changes

Frame Budget Optimization

Layer / File(s)	Summary
Materialize frame budgets and validate mask computation vllm/multimodal/video_prune/vidcom2.py, tests/multimodal/test_vidcom2.py	compute_retention_mask iterates over transferred budgets and uses their sum for reconciliation. The test verifies one (4,) transfer, rejects scalar reads, and checks the mask shape.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 361d8

VidCom2 now transfers frame budgets once while preserving retention-mask behavior. Existing validation covers the optimized read path and mask shape, with no remaining merge-blocking risk.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: reading VidCom2 frame budgets once for improved performance.
Description check	✅ Passed	The description is directly related to the changeset and explains the optimization, implementation, tests, benchmarks, and limitations.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Isotr0py]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87262 for commit 361d811cddcb.