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], block_tables: Tensor, seq_lens_kv: 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 paged decode from fixed page tables.

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.

The one-shot planner reads seq_lens_kv on the host to derive exact plan bounds, so this convenience API is not CUDA-graph-capturable. Capture callers should plan BatchDecodePagedTSWrapper before capture and either bind a stable runtime seq_lens tensor or let the plan retain fixed sequence lengths before replay.

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.

  • block_tables (torch.Tensor) – Fixed row-strided [B, C] page table. Rows may have padding between them, but each row must be contiguous.

  • seq_lens_kv (torch.Tensor) – Per-request K/V sequence lengths with shape [B].

  • seq_len_q (int) – Fixed query length when qo_indptr is omitted. In packed-query mode, a non-default value is a backward-compatible alias for max_seq_len_q and must agree with it when both are provided.

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

  • max_seq_len_q (int, optional) – Per-request packed-query length capacity. When omitted for packed Q, it is derived from qo_indptr unless a non-default seq_len_q supplies the bound. In fixed-query mode, it must equal seq_len_q.

  • 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.

Returns:

The fixed or packed attention output.

Return type:

torch.Tensor