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.

Parameters:

  • q (torch.Tensor) – Packed query tensor.

  • k_cache (torch.Tensor) – Separate HND key and value page pools.

  • v_cache (torch.Tensor) – Separate HND key and value page pools.

  • qo_indptr (torch.Tensor) – Cumulative packed-query offsets.

  • paged_kv_indptr (torch.Tensor) – FlashInfer CSR page metadata.

  • paged_kv_indices (torch.Tensor) – FlashInfer CSR page metadata.

  • paged_kv_last_page_len (torch.Tensor) – FlashInfer CSR page metadata.

  • page_size (int) – Number of K/V tokens stored in each page.

  • kv_layout ({"HND"}) – Layout of the separate K and V page pools.

  • mask_type ({"dense", "causal"}) – Attention mask mode.

  • window_left (int) – Left sliding-window extent, or -1 to disable the window.

  • sm_scale (float, optional) – Softmax scale; defaults to the inverse square root of head size.

  • output_scale (float) – Scale applied to the attention output.

  • out_dtype (torch.dtype, optional) – Requested output dtype.

  • out (torch.Tensor, optional) – Caller-owned output tensor.