[Perf][Attention] Vectorize causal-conv metadata offsets

[waizuichougou]

Purpose

compute_causal_conv1d_metadata builds the per-chunk offsets used by causal-convolution prefill kernels. The current implementation creates these offsets with a Python loop that calls extend(range(...)) once for every sequence. This host-side work grows with both the prefill batch size and the total number of 8-token chunks, delaying the following Triton launch.

This change derives the same per-row offsets with NumPy prefix sums and repeated row starts, then keeps the result in the existing pinned int32 staging-tensor path. The metadata contract and kernel inputs are unchanged. The helper is shared by Mamba, GDN, ShortConv, and Kimi K3 attention backends.

Test Plan

Existing kernel regression suite and changed-file checks:

pytest -q tests/kernels/mamba/test_causal_conv1d.py --tb=short
pre-commit run --files vllm/v1/attention/backends/utils.py

The metadata output was also compared with the original implementation over 2,000 randomized CPU cases, including empty batches, zero-length rows, mixed chunk counts, and both int32 and int64 cumulative-length inputs. The resulting metadata was exercised through causal_conv1d_fn on an NVIDIA A100 over four targeted variable-length batches, with kernel output and convolution-state updates bitwise identical to the inline/original metadata path.

End-to-end metadata construction benchmark on an NVIDIA A100 PCIe 40GB, including device-buffer allocation, metadata H2D copies, and CUDA synchronization. Sequence lengths are equal within each row. Results are medians of 100 measured iterations after 10 warmups (30 iterations for the two largest rows):

batch	sequence length	total tokens	8-token chunks	before (µs)	after (µs)	speedup
1	2048	2048	256	222.29	208.39	1.07x
8	256	2048	256	233.52	210.69	1.11x
32	64	2048	256	269.80	208.63	1.29x
128	16	2048	256	412.35	211.10	1.95x
1	8192	8192	1024	281.78	213.45	1.32x
8	1024	8192	1024	285.98	214.37	1.33x
32	256	8192	1024	321.63	213.91	1.50x
128	64	8192	1024	466.03	215.82	2.16x
32	2048	65536	8192	794.17	245.52	3.23x
128	2048	262144	32768	2538.63	329.42	7.71x
128	8192	1048576	131072	13789.55	802.64	17.18x

Test Result

All 164 existing causal-convolution kernel tests passed. The 2,000-case randomized metadata comparison and targeted A100 causal_conv1d_fn integration checks passed, and all production changed-file pre-commit checks passed.

Duplicate-work check

Checked the open PR list for causal_conv1d, causal conv, and offsetlist on 2026-09-05. No open PR was found that changes compute_causal_conv1d_metadata or its host-side offset construction. Related work is separate: #52611 addresses a JIT specialization race in causal_conv1d.py, the merged #52388 optimizes Kimi K3 state metadata, and the open #50226 adds SM103-specific GDN kernels; none changes this shared metadata helper.

AI-assisted contribution

AI assistance was used for this change.

[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.

[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: 95f26d36-19e5-4cff-9349-8c303d574273

📥 Commits

Reviewing files that changed from the base of the PR and between bc96d76 and 2b74fdc.

📒 Files selected for processing (1)

• vllm/v1/attention/backends/utils.py

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

---

Walkthrough

The causal convolution metadata path now computes per-row chunk offsets with vectorized NumPy operations and creates pinned tensors through np_to_pinned_tensor. The unused PIN_MEMORY import was removed.

Changes

Causal convolution metadata

Layer / File(s)	Summary
Vectorized metadata construction vllm/v1/attention/backends/utils.py	compute_causal_conv1d_metadata reuses a NumPy array for chunk totals and row indices, replaces the offset loop with np.cumsum and np.repeat, and uses np_to_pinned_tensor for the result. The unused PIN_MEMORY import was removed.

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

Suggested reviewers: mgoin, njhill

Merge Risk: ⚪ Minimal · up to 2b74f

The metadata construction is vectorized while retaining the pinned staging path and existing contract; supplied tests pass with no merge-blocking risk identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title clearly and concisely describes the main change: vectorizing causal-convolution metadata offset construction for performance.
Description check	✅ Passed	The description directly explains the vectorization, preserved metadata behavior, affected helper, validation, and benchmark results.

---

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.}

[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.