fix(attention): accept non-contiguous index tensors in plan() - #5385

📌 Description

plan() passes user-supplied index tensors (indptr / indices / last_page_len) to the C++ planner and the kernels as raw data pointers without checking contiguity. A strided view (e.g. indptr_full[::2]) is therefore misread, in two ways:

• CUDA inputs: in non-CUDA-graph mode, self._*_buf = indptr.to(self.device) is a no-op for a strided CUDA view, so the kernel buffers stay strided. indptr.to("cpu"), on the other hand, returns a contiguous host copy for the planner. The plan and the kernel then disagree about the layout.
• CPU inputs: .to("cpu") is a no-op, so the strided host tensor goes straight into the C++ planner.

Measured on H200 (SM90) with stride-2 views holding identical values (fp16, 8/2 heads, head_dim 128, page_size 16, 4 requests):

Fix: call .contiguous() on the index tensors at the top of plan() in:

• BatchPrefillWithPagedKVCacheWrapper.plan: qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len
• BatchPrefillWithRaggedKVCacheWrapper.plan: qo_indptr, kv_indptr, and the optional o_indptr / v_indptr
• BatchDecodeWithPagedKVCacheWrapper._plan_impl and fast_decode_plan: indptr, indices, last_page_len
• BlockSparseAttentionWrapper.plan: indptr, indices
• seq_lens (when given) in the paged prefill plan() and decode _plan_impl: it is copied to the host with seq_lens.cpu(), which is a no-op for a strided CPU view, and the host copy goes to the native planner. With fa3 paged prefill a strided CPU seq_lens was an illegal memory access.
• the caller-supplied host copies, which go straight to the native planner: _plan_indptr_host (the ragged wrapper's qo_indptr_cpu / kv_indptr_cpu, which it validated for device/shape/dtype but not contiguity) and fast_decode_plan's global_override_indptr_cpu. These parameters were added upstream after the first version of this PR.
• workspace_size() on the paged prefill and decode wrappers, which mirrors plan()'s host copies: the same tensors, including seq_lens. Without it, strided CPU indptrs failed the native size query's PrefillSplitQOKVIndptr check (tensor-core decode, fa2 prefill). Thanks to the CodeRabbit review for flagging this one.

These tensors are tiny, and plan() already copies them between host and device. .contiguous() returns self for contiguous inputs, so the usual path is unchanged. The CUDA-graph paths copy_() into preallocated buffers and were already correct.

Out of scope:

• The MLA planner already rejects non-contiguous metadata (mla/_batch_mla/_planning.py).
• VariableBlockSparseAttentionWrapper builds its own indptr via cat/cumsum, so it is always contiguous.

🔍 Related Issues

Fixes #5260

🚀 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.).

New file tests/attention/test_plan_noncontiguous_metadata.py (32 cases): ragged and paged prefill (fa2/fa3), paged decode (with and without tensor cores), fast_decode_plan, and block-sparse, with metadata on CUDA and on CPU, and with and without seq_lens for paged prefill and decode. Each case plans the same problem with contiguous metadata and with stride-2 views of the same values, then compares outputs. A workspace_size() case checks that strided and contiguous metadata give the same sizes.

Verified on H200 (SM90, CUDA 12.8, torch 2.14, flashinfer at main 9cab083 + this branch), each case in its own process so an illegal access cannot mask later cases:

The three cases added in the last round (qo_indptr_cpu / kv_indptr_cpu and global_override_indptr_cpu) also fail with only those two normalizations reverted on top of this branch, so they cover a distinct bug rather than the plan() one.

Regression check on H200, sampled from test_batch_prefill_kernels.py, test_batch_decode_kernels.py, test_block_sparse.py, test_workspace_size.py, test_lse_base.py and this new file (all cases of functions with <=300 parametrizations, every 25th case of the larger sweeps, 1,498 ids over 3 GPUs): 1,140 passed, 358 skipped, 0 failed. pre-commit run --files <changed> passes (ruff, ruff-format, mypy and the whitespace hooks).

Reviewer Notes

• Normalizing vs rejecting. I chose to normalize so callers that pass strided views keep working (ragged fa3 with CUDA metadata happened to be correct before this change). There is precedent both ways: the MLA planner raises ValueError for non-contiguous metadata, and attention/_core.py rejects a non-contiguous bidirectional_ranges with the rationale that copying "would hand the kernel a fresh pointer on every call, which a captured CUDA graph would then replay against freed memory". That hazard does not apply here: .contiguous() returns the same object for an already-contiguous tensor, so the common path is unchanged, and a pointer only changes for input that was previously misread. Where a copy can happen, the wrapper retains it — the non-graph paths assign it to self._*_buf, and the CUDA-graph paths copy_() into their preallocated buffers instead of rebinding (fast_decode_plan rebinds only in its non-graph branch). I am happy to switch to rejecting if maintainers prefer the strict contract; the tests would just invert.
• BatchAttentionWithAttentionSinkWrapper subclasses BatchPrefillWithPagedKVCacheWrapper, so it inherits the fix and needs no separate change.
• The C++ launchers still don't check contiguity (CHECK_INPUT) on indptr. This PR keeps the change at the Python boundary. Adding C++ checks would also cover direct callers of the low-level modules, and could be a follow-up.
• Rebased onto current main after the first review round; seq_lens and workspace_size() coverage came from the earlier CodeRabbit comment, and fast_decode_plan coverage was added in this round.

Summary by CodeRabbit

• Bug Fixes
  • Improved handling of non-contiguous metadata during attention planning and workspace sizing, helping ensure correct results and consistent workspace requirements.
  • Applied across decode, prefill, and block-sparse attention workflows, including planning inputs and optional metadata.
• Tests
  • Added regression coverage comparing contiguous and non-contiguous metadata behavior across decode, prefill, and block-sparse attention, including CPU and CUDA cases where applicable.