[Bugfix][Attention] Size FlashInfer paged-KV buffers for local attention

[thegoldenflow]

Purpose

Fixes #49980.

Under chunked local attention, make_local_attention_virtual_batches turns N requests into M > N virtual batches and passes M to the builder as num_reqs. FlashInferMetadataBuilder sized paged_kv_indptr, paged_kv_last_page_len and paged_kv_indices from max_num_seqs, so they overflow whenever M > max_num_seqs. The reporter (max_num_seqs=1) hit ValueError: provided out is the wrong size for the accumulation from the np.cumsum in _compute_flashinfer_kv_metadata.

This PR sizes those buffers in __init__ from an upper bound on the virtual batch count, min(2 * max_num_seqs + cdiv(max_num_batched_tokens, c), max_num_batched_tokens) (new helper max_local_attention_virtual_batches).

• Detected with getattr(spec, "attention_chunk_size", None). The builder gets a ChunkedLocalAttentionSpec when the hybrid KV cache manager is on (the CUDA default). When it is off (e.g. EAGLE, --disable-hybrid-kv-cache-manager) it gets a FullAttentionSpec(attention_chunk_size=...), which an isinstance check would miss. The test covers both.
• Sized up front, not grown on overflow. When a whole batch runs on trtllm-gen (SM100), build() skips _compute_flashinfer_kv_metadata and runs a torch.cumsum into an out-of-range slice of paged_kv_indptr.gpu. The only signal is a warning, and the kernel gets a cum_seq_lens_kv that is too short. The issue's grow-on-overflow check lives in _compute_flashinfer_kv_metadata, so it would never see this path. Rebinding the buffers is also unsafe: CUDA-graph decode wrappers alias them, and fast_decode_plan does not copy into them.
• Backstop. _ensure_paged_kv_capacity runs at the top of build() and again once the page count is known. If the bound is ever wrong, it grows the buffers, or raises when enable_cuda_graph is set.
• Page count only grows (max, not min). With the hybrid manager off, all Llama-4 layers merge into one KV cache group whose FullAttentionSpec keeps attention_chunk_size. The global-attention layers read that spec but attend over the whole sequence, so capping pages at the chunk size would under-size them (2250 → 512 pages in the reporter's config).
• XQA draft mask (third commit). [Attention] Add FlashInfer XQA decode support on SM12x #49718, merged after this PR was opened, added _decode_mask_cache, also sized from max_num_seqs. With XQA speculative decoding, a verify window split evenly by a chunk boundary can produce more uniform decodes than that. XQA indexes the mask per request without a shape check, so it reads past the end. The cache now uses the same bound. Reaching this needs Llama-4 with FlashInfer on SM90/SM12x and speculative decoding (any method) with an odd k >= 3.

Why the bound holds, and what it allocates

With c = attention_chunk_size, N = max_num_seqs and T = max_num_batched_tokens, each request contributes 1 + cdiv(q_i - f_i, c) virtual batches, where f_i = min(c - ((s_i - q_i) mod c), q_i) >= 1. Summing gives V <= 2N + cdiv(T, c). The 2N term is needed: q_i = 2 with f_i = 1 gives V = 2N. Every virtual batch holds at least one query token, so V <= T. Every virtual batch sees at most c KV tokens, so pages <= V * cdiv(c, block_size). A brute-force check over 2.7M configurations found no violations. Zero-length padded requests only occur in FULL CUDA-graph mode, which chunked-local rules out (AttentionCGSupport.NEVER).

config	paged_kv_indptr	paged_kv_indices
reporter's (max_model_len=36000, T=36000, N=1, c=8192, block 16)	2 → 8	2250 → 3584
max_model_len=131072, T=8192, N=256, c=8192, block 16	257 → 514	unchanged
spec without attention_chunk_size	unchanged	unchanged

Question for reviewers: FullAttentionSpec.merge keeps attention_chunk_size, but its consistency assert only checks fields(AttentionSpec), which excludes it. Is that intended? If not, the max() above can go. I left it as out of scope.

Test Plan

pytest -q tests/v1/attention/test_chunked_local_attention.py             # CPU only
pytest -q tests/v1/attention/test_flashinfer_chunked_local_attention.py  # CUDA + flashinfer, no Llama-4 weights

To see the failures, revert vllm/v1/attention/backends/flashinfer.py to origin/main and re-run the second file. create_vllm_config fetches config.json from the gated meta-llama/Meta-Llama-3-8B, so an HF token is needed.

Test Result

• Rebased onto current main, CPU only. The only conflict was [MRV2] Buffer util simplifications #56888's buffer-type refactor; the fix was ported with its logic unchanged. test_chunked_local_attention.py: 11 passed. test_flashinfer_chunked_local_attention.py: 5 passed, on a copy with the CUDA-only skip removed and a locally cached Qwen config. With flashinfer.py reverted to origin/main: 3 failed. Both paged-KV cases fail with the [Bug]: FlashInfer builder ValueError 'provided out is the wrong size for the accumulation' with Llama-4 chunked local attention when a prefill exceeds attention_chunk_size and max_num_seqs is small #49980 ValueError, and the mask case fails with assert 1 == 2.
• Before the rebase, on an RTX 4090. The FlashInfer test file: 4 passed; with the fix reverted, 2 failed with the [Bug]: FlashInfer builder ValueError 'provided out is the wrong size for the accumulation' with Llama-4 chunked local attention when a prefill exceeds attention_chunk_size and max_num_seqs is small #49980 ValueError. pytest tests/v1/attention/ -k flashinfer: 13 passed, 844 skipped.
• Independent end-to-end. @janbernloehr reported that the original patch, cherry-picked onto 0.26.1rc1.dev528+gf8d03e774, passes their previously fatal 32k-prefix Llama-4 Scout scenario on DGX H100 (ISL avg 34,000.80 tok, 108.23 output tok/s). Thank you! It covers only the [Bug]: FlashInfer builder ValueError 'provided out is the wrong size for the accumulation' with Llama-4 chunked local attention when a prefill exceeds attention_chunk_size and max_num_seqs is small #49980 native-prefill path: that build predates [MRV2] Buffer util simplifications #56888, MRV2 becoming the default and SM90 XQA decode.
• Still open. A GPU run of the rebased branch. The XQA mask fix is CPU-verified only, and no run has exercised the SM100 trtllm-gen path.
• Lint. After the rebase, pre-commit (including manual-stage mypy 3.10–3.13) passes on Windows, except update-dockerfile-graph, which needs /bin/bash.

AGENTS.md

• Not a duplicate. Re-ran the duplicate-work checks on 2026-09-18. This is the only open PR referencing [Bug]: FlashInfer builder ValueError 'provided out is the wrong size for the accumulation' with Llama-4 chunked local attention when a prefill exceeds attention_chunk_size and max_num_seqs is small #49980, and no other open PR resizes these buffers.
• Model evaluation. N/A. Only buffer sizes change; no kernel or attention math is touched.
• AI assistance. Written with AI assistance (Claude). I reviewed every changed line and verified the original fix on GPU myself.

[thegoldenflow]

GPU verification is done. Ran the FlashInfer chunked-local test on an RTX 4090: 4 passed with this patch. Reverting only the two source files this PR touches reproduces #49980 with the predicted "provided out is the wrong size for the accumulation", and restoring them returns to 4 passed. Neighbouring flashinfer tests show no regressions, and full pre-commit passes on Linux. Both the GPU test and Linting sections of the body now carry the real output, so this is no longer a draft.

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

[janbernloehr]

I confirm that cherry-picking this patch into the current upstream vLLM 0.26.1rc1.dev528+gf8d03e774 passes our previously-fatal 32k-prefix Llama-4 Scout scenario on DGX H100 (ISL avg 34,000.80 tok, 108.23 output tok/s).

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @wsyjh8.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[thegoldenflow]

@janbernloehr Thank you for verifying this on DGX H100 with your Llama-4 Scout workload, much appreciated! I've since rebased onto current main and referenced your result in the PR description.