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, *, seq_len_kv: int, seq_lens_kv: Tensor | None = None, 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 one-shot entry point mirrors BlockSparsePagedTSWrapper: the logical request-to-page spans come from paged_kv_indptr and optional seq_lens_kv, while live physical page IDs and sparse routes still come from paged_kv_indices plus BSR metadata at run time.

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.

  • seq_len_kv (int) – Exact shared logical KV length, or the maximum logical length when seq_lens_kv is provided.

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

  • kv_valid_bits (torch.Tensor, optional) – Contiguous UInt32 logical-token validity bitmap [B, ceil(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