flashinfer.attention.prims_ts.batch_decode_with_paged_kv_cache

flashinfer.attention.prims_ts.batch_decode_with_paged_kv_cache(q: Tensor, paged_kv_cache: Tensor | tuple[Tensor, Tensor], paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, *, seq_len_q: int = 1, qo_indptr: Tensor | None = None, max_seq_len_q: int | None = None, mask_type: Literal['dense', 'causal'] = 'dense', window_left: int = -1, kv_layout: Literal['HND'] = 'HND', bmm1_scale: float | None = None, bmm2_scale: float = 1.0, out: Tensor | None = None, out_dtype: dtype | None = None) → Tensor

One-shot fixed or packed-Q native-CSR paged decode.

SQ1 preserves the [B, Hq, D] query/output contract. For fixed seq_len_q>1, query and output are both token-major [B, SQ, Hq, D]. Providing cumulative qo_indptr selects packed [total_q, Hq, D] query/output; the wrapper derives max_seq_len_q once when it is omitted. No transpose is hidden here.

Parameters:

  • q (torch.Tensor) – Fixed or packed query tensor.

  • paged_kv_cache (torch.Tensor or tuple[torch.Tensor, torch.Tensor]) – Combined or separate paged K/V storage.

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

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

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

  • seq_len_q (int) – Fixed query length when qo_indptr is omitted.

  • qo_indptr (torch.Tensor, optional) – Cumulative query offsets selecting packed-query mode.

  • max_seq_len_q (int, optional) – Static packed-query length bound.

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

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

  • kv_layout ({"HND"}) – Layout of the paged K/V cache.

  • bmm1_scale (float, optional) – QK and value/output scaling factors.

  • bmm2_scale (float, optional) – QK and value/output scaling factors.

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

  • out_dtype (torch.dtype, optional) – Output dtype; defaults to out.dtype or the query dtype.