flashinfer.msa_ops.msa_sparse_attention

flashinfer.msa_ops.msa_sparse_attention(q: Tensor, k: Tensor, v: Tensor, q2k_indices: Tensor, cu_seqlens_q: Tensor, cu_seqlens_k: Tensor | None = None, causal: bool = False, softmax_scale: float | None = None, page_table: Tensor | None = None, seqused_k: Tensor | None = None, return_softmax_lse: bool = False, k_scale: Tensor | None = None, v_scale: Tensor | None = None, k_global_scale: float | None = None, v_global_scale: float | None = None, q_offset=None, return_temperature_lse: bool = False, lse_temperature_scale: float = 1.0)

Minimax Sparse Attention forward (prefill) for SM120/SM121.

Each query attends only the top-K KV blocks selected in q2k_indices. Query tokens are processed in tiles: each tile runs one online softmax over the union of the blocks its tokens selected, writing the final output directly. The union metadata is built internally from q2k_indices.

q/k/v are (total_tokens, num_heads, head_dim) with varlen offsets cu_seqlens_q/cu_seqlens_k; q2k_indices is (num_kv_heads, total_q, topk) int32 (ascending, -1 padded). q is bf16/fp16; k/v are bf16/fp16, fp8 (E4M3), or packed NVFP4 (uint8 with k_scale/v_scale), in GQA or MHA layouts, flat or paged.

Notable optional parameters:

page_tabletorch.Tensor, optional

Enables the paged-KV path: k/v are then (num_pages, num_kv_heads, 128, head_dim) and page_table is (batch_size, max_pages) int32 mapping batch-local KV block i to a page. Requires seqused_k.

seqused_ktorch.Tensor, optional

(batch_size,) int32, valid KV length per sequence (paged path). cu_seqlens_k may be omitted when this is given.

return_softmax_lsebool

Also return the natural-log LSE, shape (total_q, num_qo_heads) float32 (-inf for queries with no valid selected blocks).

return_temperature_lsebool

Also return the MSA temperature LSE (exponent scaled by lse_temperature_scale), same shape as the softmax LSE. When set, the return is (out, lse, lse_t).

Returns:

Output (total_q, num_qo_heads, head_dim); plus LSE if return_softmax_lse; plus the temperature LSE if return_temperature_lse.

Return type:

torch.Tensor or tuple of torch.Tensor