flashinfer.attention.prims_ts.block_sparse_attention_with_paged_kv_cache

flashinfer.attention.prims_ts.block_sparse_attention_with_paged_kv_cache(q: Tensor, paged_kv_cache: Tensor | tuple[Tensor, Tensor], paged_kv_indptr: Tensor, paged_kv_indices: Tensor, block_indptr: Tensor, block_indices: Tensor, q_block_size: int, kv_block_size: int, *, max_seq_len_kv: int, seq_lens_kv: Tensor, kv_valid_bits: Tensor | None = None, mask_type: Literal['dense', 'causal'] = 'dense', sm_scale: float | None = None, out: Tensor | None = None) → Tensor

Plan and run one fixed-Q paged block-sparse attention launch.

This convenience entry point synchronously validates live page and sparse metadata, including the complete live physical-page-ID prefix, creates a capacity-only temporary plan, then forwards the inspected tensors through the trusted live run API. It cannot run during CUDA Graph capture; plan a wrapper outside capture and capture only BlockSparsePagedTSWrapper.run() instead.

Parameters:

  • q (torch.Tensor) – Compact query tensor [B, Sq, Hq, D].

  • paged_kv_cache (PagedKVCache) – Either a combined cache [P, 2, Hkv, page_size, D] or a (K, V) tuple whose tensors are [P, Hkv, page_size, D].

  • paged_kv_indptr (torch.Tensor) – Contiguous Int32 request offsets into paged_kv_indices, with shape [B + 1].

  • paged_kv_indices (torch.Tensor) – Contiguous Int32 physical page IDs referenced by paged_kv_indptr.

  • block_indptr (torch.Tensor) – Contiguous Int32 BSR row offsets with shape [B, Hkv, ceil(Sq / q_block_size) + 1].

  • block_indices (torch.Tensor) – Contiguous Int32 logical KV-block IDs referenced by block_indptr.

  • q_block_size (int) – Positive number of logical query tokens represented by one BSR row. The product with Hq / Hkv must be divisible by 8 so a physical Q tile does not cross row boundaries.

  • kv_block_size (int) – Number of logical KV tokens represented by one BSR block ID; it must be 8, 16, 32, or a positive multiple of 64.

  • max_seq_len_kv (int) – Static maximum logical K/V length used for planning.

  • seq_lens_kv (torch.Tensor) – Contiguous Int32 per-request logical KV lengths with shape [B].

  • kv_valid_bits (torch.Tensor, optional) – Contiguous UInt32 logical-token validity bitmap [B, ceil(max_seq_len_kv / 32)].

  • mask_type ({"dense", "causal"}, optional) – Attention mask applied inside each selected sparse block.

  • sm_scale (float, optional) – Softmax scale. Defaults to 1 / sqrt(D).

  • out (torch.Tensor, optional) – Caller-owned compact output buffer [B, Sq, Hq, D].

Returns:

The compact output tensor; identical to out when provided.

Return type:

torch.Tensor