flashinfer.cute_dsl.sparse.bsa_attn_sm100_blk64.bsa_attn_sm100_blk64_fwd

flashinfer.cute_dsl.sparse.bsa_attn_sm100_blk64.bsa_attn_sm100_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, kv_splits: int | str = 1, use_clc: bool | None = None, q_scale: Tensor | None = None, k_scale: Tensor | None = None, v_scale: Tensor | None = None) → Tuple[Tensor, Tensor | None]

Forward pass for BSA block-sparse attention using the blk64 CuTe-DSL kernel.

Supports SM100 and SM103. MHA only (num_kv_heads must equal num_heads); head_dim must be 128.

Parameters:

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

• k – Key tensor (batch, seqlen_k, num_heads, head_dim).

• v – Value tensor (batch, seqlen_k, num_heads, 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).

• kv_splits – Number of KV splits (“auto”, or an explicit int in [1, 256]). kv_splits=1 (or an “auto” resolution of 1) reproduces the original single-kernel behavior exactly.

• use_clc – Optional scheduler override (True forces the CLC persistent scheduler, False forces the static scheduler, None uses the shape-based heuristic).

• q_scale – Sage FP8 quantization scales. All three must be provided together to enable the FP8 path (q/k/v must then be float8_e4m3fn); otherwise all three must be None and q/k/v must be bfloat16. The FP8 path additionally requires batch_size == 1, num_head in (4, 8), q2k_block_nums is None, and block_sizes is None (upstream kernel limits, not specific to this integration).

• k_scale – Sage FP8 quantization scales. All three must be provided together to enable the FP8 path (q/k/v must then be float8_e4m3fn); otherwise all three must be None and q/k/v must be bfloat16. The FP8 path additionally requires batch_size == 1, num_head in (4, 8), q2k_block_nums is None, and block_sizes is None (upstream kernel limits, not specific to this integration).

• v_scale – Sage FP8 quantization scales. All three must be provided together to enable the FP8 path (q/k/v must then be float8_e4m3fn); otherwise all three must be None and q/k/v must be bfloat16. The FP8 path additionally requires batch_size == 1, num_head in (4, 8), q2k_block_nums is None, and block_sizes is None (upstream kernel limits, not specific to this integration).

Returns:

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