flashinfer.attention.prims_ts.block_sparse_attention

flashinfer.attention.prims_ts.block_sparse_attention(q: Tensor, k: Tensor, v: Tensor, block_indptr: Tensor | None, block_indices: Tensor | None, q_block_size: int, kv_block_size: int, *, exact_block_bits: Tensor | None = None, k_summary: Tensor | None = None, v_summary: Tensor | None = None, kv_valid_bits: Tensor | None = None, sparse_format: Literal['bsr', 'bitmask'] = 'bsr', use_proxy_routes: bool = False, mask_type: Literal['dense', 'causal'] = 'dense', sm_scale: float | None = None, out: Tensor | None = None) → Tensor

Plan and run one compact-BSHD block-sparse attention launch.

This one-shot form creates a capacity-only plan and passes the original routing tensors to BlockSparseTSWrapper.run(). BSR inputs are synchronously inspected to validate canonical rows and derive their maximum width. Bitmask inputs use the structural KV-block count as a conservative capacity bound. It therefore cannot be invoked inside CUDA Graph capture; plan a wrapper outside capture and capture only run() instead.

Parameters:

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

  • k (torch.Tensor) – Compact key tensor [B, Skv, Hkv, D].

  • v (torch.Tensor) – Compact value tensor with the same shape, dtype, and strides as k.

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

  • block_indices (torch.Tensor, optional) – Contiguous Int32 semantic KV-block IDs referenced by block_indptr. Required in BSR mode.

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

  • exact_block_bits (torch.Tensor, optional) – Compact UInt32 exact-block bitmap required in bitmask mode.

  • k_summary (torch.Tensor, optional) – Per-block mean K tensor required when proxy routes are enabled.

  • v_summary (torch.Tensor, optional) – Per-block summed V tensor required when proxy routes are enabled.

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

  • sparse_format ({"bsr", "bitmask"}, optional) – Runtime sparse representation. Defaults to "bsr".

  • use_proxy_routes (bool, optional) – Represent unselected blocks through K/V summaries. Proxy routes currently require dense masking.

  • 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