flashinfer.attention.prims_ts.batch_prefill_with_paged_kv_cache

flashinfer.attention.prims_ts.batch_prefill_with_paged_kv_cache(q: Tensor, k_cache: Tensor, v_cache: Tensor, qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, *, page_size: int = 32, kv_layout: Literal['HND'] = 'HND', mask_type: Literal['dense', 'causal'] = 'dense', window_left: int = -1, sm_scale: float | None = None, output_scale: float = 1.0, out_dtype: dtype | None = None, out: Tensor | None = None) → Tensor

Run one-shot packed-Q context attention over separate HND page pools.

Q/output use [total_q, Hq, D] storage and K/V each use [num_pages, Hkv, page_size, D] with page size 16, 32, 64, or 128. qo_indptr describes Q rows while the three paged-KV metadata tensors use FlashInfer’s CSR representation. Physical page indices need not be identity ordered. D may be 128 or 256; Q, K, and V must share one supported dtype.