perf(prefill): reduce with Tensor.max() instead of the Python builtin in plan()

[gf239]

📌 Description

BatchPrefillWithPagedKVCacheWrapper.plan() derives _max_q_len and _max_kv_len with the Python builtin max() over a 1-D int32 host tensor:

self._max_q_len = max(qo_indptr_host[1:] - qo_indptr_host[:-1]).item()
...
self._max_kv_len = max(kv_lens_arr_host).item()

The builtin iterates the tensor element by element and materialises a 0-d tensor per element; Tensor.max() is a single reduction dispatch. This PR switches both lines to the tensor reduction.

Both lines run on every plan() call for any caller that omits max_token_per_sequence / max_sequence_kv -- vLLM's FlashInfer backend among them, which calls plan() twice per engine step under speculative decoding (target verify batch + drafter). Measured on the host (torch 2.14, CPU):

site	requests	before	after	saved
_max_q_len	8	12.7 µs	3.1 µs	9.6 µs
_max_kv_len	8	10.0 µs	1.2 µs	8.8 µs
_max_q_len	1	3.9 µs	3.3 µs	0.6 µs
_max_kv_len	1	1.1 µs	0.8 µs	0.3 µs

About 18 µs per plan() at 8 requests, ~1 µs at a single request. Small, but it sits on the CPU critical path of every step when CUDA graphs are piecewise, and it costs nothing.

Equivalence. On every reachable input the two spellings return the same Python int: the operand is 1-D, int32, numel >= 1 (plan() is only reached for a non-empty batch). They differ only where the builtin already raised -- empty (ValueError vs RuntimeError), 0-d, and 2-D or higher (the builtin iterates rows and compares tensors) -- and int32 admits no NaN, so the max/amax NaN distinction cannot arise. Checked over 1800 random 1-D int32 tensors with numel ∈ {1, 2, 3, 8, 16, 64}: 0 mismatches, identical types.

Test. Adds tests/attention/test_prefill_plan_max_lens.py, which pins both attributes after plan() against the builtin-max computation (spelled the old way) and against the input lengths directly, over batch_size ∈ {1, 3, 8, 17}, page_size ∈ {1, 16} and fp16/bf16 -- including the single-element reduction.

🔍 Related Issues

None.

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Results. Run on an RTX 4090 (sm89), each time twice -- once against the stock package as a control that the environment and JIT-compiled kernels work, once with the two-line patch overlaid first on PYTHONPATH -- so a failure could be attributed:

environment	tests	control (stock)	patched
flashinfer-python 0.6.16.post3 (the vLLM 0.28 production venv on this box; plan() is line-identical to main at both sites)	test_prefill_plan_max_lens.py (16) + tests/attention/test_batch_prefill.py (4)	20 passed	20 passed
flashinfer-python 0.6.18 (version.txt on main)	same	not runnable here	not runnable here

The 0.6.18 row could not be completed on this machine and I am not going to pretend otherwise: 0.6.18's JIT needs a CUDA toolchain newer than the system CUDA 12.0 (nvcc fatal: Unknown option '--compress-mode=size'), and the pip-provided CUDA 13 packages I tried in a fresh venv failed in three successive layers (bundled CCCL rejecting a 13.3 nvcc against 13.0 runtime headers; a leftover 13.3 cicc emitting PTX 9.3 for a 13.0 ptxas; then crt/host_runtime.h vs cudafe stub disagreeing on __cudaLaunch). The control run fails identically to the patched run in every one of those attempts, so none of it is attributable to this change. Since plan() is line-identical between 0.6.16.post3 and main at both edited sites, the 0.6.16.post3 row is the same code path; CI on its supported toolchain covers 0.6.18.

pre-commit run --files flashinfer/prefill.py tests/attention/test_prefill_plan_max_lens.py: all hooks passed (mypy, ruff check, ruff format, whitespace/EOL/tab checks); no file modified.

🔬 Experimental Track

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Improved calculation of maximum query and key-value sequence lengths during paged KV-cache prefill planning, while preserving existing planning behavior.

• Tests

  • Added coverage across batch sizes, page sizes, data types, and single-request scenarios to verify accurate maximum-length calculations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: efe507c5-8b59-4961-ab5a-2542b60a21db

📥 Commits

Reviewing files that changed from the base of the PR and between 62362bb and 524fd42.

📒 Files selected for processing (2)

• flashinfer/prefill.py
• tests/attention/test_prefill_plan_max_lens.py

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

---

📝 Walkthrough

Walkthrough

BatchPrefillWithPagedKVCacheWrapper.plan now computes maximum query and KV lengths with tensor reductions. A parametrized test validates the resulting Python integers across batch sizes, page sizes, and dtypes.

Changes

Prefill maximum length handling

Layer / File(s)	Summary
Tensor maximum reductions and validation flashinfer/prefill.py, tests/attention/test_prefill_plan_max_lens.py	plan() derives _max_q_len and _max_kv_len with tensor .max().item() reductions. The test compares both values with host-side and raw tensor maxima across parameterized inputs.

Priority: ⬇️ Low

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

Merge Risk: ⚪ Minimal · up to 524fd

Prefill planning now computes query and KV length maxima through tensor reductions while retaining the same integer results. Covered inputs preserve prior behavior, 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 2 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 primary change: replacing Python's builtin max() with Tensor.max() in prefill plan().
Description check	✅ Passed	The description explains the change, motivation, measured impact, equivalence, tests, pre-commit results, and experimental status. It satisfies the repository template.
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

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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