flashinfer.attention.prims_ts.block_sparse_attention

flashinfer.attention.prims_ts.block_sparse_attention(q: Tensor, k: Tensor, v: Tensor, block_indptr: Tensor, block_indices: Tensor, q_block_size: int, kv_block_size: int, *, 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 compact-BSHD block-sparse attention launch.

This one-shot form synchronously inspects canonical BSR, derives its largest semantic row, creates a capacity-only plan, and passes the original routing tensors to BlockSparseTSWrapper.run(). 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) – Contiguous Int32 BSR row offsets with shape [B, Hkv, ceil(Sq / q_block_size) + 1].

  • block_indices (torch.Tensor) – Contiguous Int32 semantic 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.

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