flashinfer.cute_dsl.sparse.bsa_attn_blk64_fwd

flashinfer.cute_dsl.sparse.bsa_attn_blk64_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, softmax_scale: float | None = None, return_lse: bool = False, out: Tensor | None = None, lse: Tensor | None = None) → Tuple[Tensor, Tensor | None]

Forward pass for BSA block-sparse attention using the blk64 CUDA C++ kernel (SM100 only).

Block granularity is 64 tokens (kSparseBlockSize=64, kRows=64). Only bfloat16 inputs are supported and head_dim must be 128.

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

• q2k_block_index – Block index tensor (batch, num_heads, num_q_blocks, max_kv_blocks), int32.

• block_sparse_num – Number of KV blocks each Q block attends to (>= 1). Ignored when q2k_block_nums is provided.

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

• q2k_block_nums – Per-(batch, head, q_block) number of KV blocks to attend to, (batch, num_heads, num_q_blocks) int32. When None, uses fixed block_sparse_num.

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

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

• out – Pre-allocated output tensor (batch, seqlen_q, num_heads, head_dim).

• lse – Pre-allocated LSE tensor (batch, num_heads, seqlen_q).

Returns:

(out, lse) where lse is None if return_lse is False.