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, block_tables: Tensor, seq_lens_kv: 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, block_tables supplies one fixed row-strided page table, and seq_lens_kv supplies logical K/V lengths. Physical page indices need not be identity ordered. D may be 128 or 256; Q, K, and V must share one supported dtype.

This convenience API reads request metadata on the host to derive exact plan bounds and is not CUDA-graph-capturable. Capture-sensitive callers should plan BatchPrefillPagedTSWrapper before capture and bind the same metadata tensors directly in run.

Parameters:

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

  • k_cache (torch.Tensor) – HND key page pool.

  • v_cache (torch.Tensor) – HND value page pool isomorphic to k_cache.

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

  • block_tables (torch.Tensor) – Physical page IDs with row-strided shape [B, C].

  • seq_lens_kv (torch.Tensor) – Logical K/V lengths with shape [B].

  • 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) – Output dtype; defaults to out.dtype when out is provided, otherwise the query dtype.

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

Returns:

The packed attention output.

Return type:

torch.Tensor