flashinfer.cute_dsl.sparse.bsa_attn_fwd

flashinfer.cute_dsl.sparse.bsa_attn_fwd(q: Tensor, k: Tensor, v: Tensor, q2k_block_index: Tensor, block_sparse_num: int, block_sizes: Tensor | None = None, q2k_block_nums: Tensor | None = None, allow_empty_block_nums: bool = True, softmax_scale: float | None = None, pack_gqa: bool | None = None, return_lse: bool = False, out: Tensor | None = None, lse: Tensor | None = None) → Tuple[Tensor, Tensor]

Forward pass for BSA block-sparse attention (SM100 only).

Parameters:

  • q – Query tensor (batch, seqlen_q, num_heads, head_dim)

  • k – Key tensor (batch, seqlen_k, num_heads_kv, head_dim)

  • v – Value tensor (batch, seqlen_k, num_heads_kv, head_dim_v)

  • q2k_block_index – (batch, num_heads, num_q_blocks, max_kv_blocks) int32

  • block_sparse_num – Number of KV blocks per Q block (even, >= 2). Ignored when q2k_block_nums is provided.

  • block_sizes – Actual token count per KV block (num_kv_blocks,) int32. Pass None to skip block-size masking (assumes full blocks).

  • q2k_block_nums – Per-(batch,head,q_block) number of KV blocks, (batch, num_heads, num_q_blocks) int32. Optional.

  • allow_empty_block_nums – Allow q2k_block_nums to contain 0.

  • softmax_scale – Softmax scale (default: 1/sqrt(head_dim)).

  • pack_gqa – Whether to pack GQA heads.

  • return_lse – Whether to return log-sum-exp.

  • out – Pre-allocated output tensor.

  • lse – Pre-allocated LSE tensor.