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, workspace: MSASparseAttentionWorkspace | None = None)

Minimax Sparse Attention forward for SM100/SM103 and SM120/SM121 GPUs.

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 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. Packed NVFP4 stores two values per byte, so its K/V last dimension is head_dim // 2 instead of head_dim. q2k_indices is (num_kv_heads, total_q, topk) int32 (ascending, -1 padded). Compute capability 10.0/10.3 supports bf16/fp16 K/V and E4M3 K/V with bf16 Q, plus exactly one packed-NVFP4 route – paged K/V read in place as strided views of a planar page, with k_scale/v_scale and both global scales, causal, without an LSE and with seqused_k rather than cu_seqlens_k (see the k, k_scale, seqused_k and k_global_scale parameters below) – and requires topk == 16. Every other packed-NVFP4 layout remains SM120/SM121-only.

Parameters:

  • q (torch.Tensor) – Query tensor of shape (total_q, num_qo_heads, head_dim) with dtype bf16/fp16.

  • k (torch.Tensor) – Key tensor. For bf16/fp16/fp8, the flat shape is (total_k, num_kv_heads, head_dim) and the paged shape is (num_pages, num_kv_heads, 128, head_dim). Packed NVFP4 uses head_dim // 2 for the last dimension in both layouts.

  • v (torch.Tensor) – Value tensor with the same layout conventions as k.

  • q2k_indices (torch.Tensor) – Int32 tensor of shape (num_kv_heads, total_q, topk) containing the selected KV block indices for each query, sorted ascending and padded with -1.

  • cu_seqlens_q (torch.Tensor) – Int32 cumulative query sequence lengths of shape (batch_size + 1,).

  • cu_seqlens_k (Optional[torch.Tensor], default=None) – Int32 cumulative KV sequence lengths of shape (batch_size + 1,) for the dense path. May be omitted when seqused_k is provided. The compute-capability 10.0/10.3 packed-NVFP4 paged route takes seqused_k as the sole source of per-request KV length and declines a call that also supplies cu_seqlens_k, rather than silently preferring one of two possibly disagreeing sources.

  • causal (bool, default=False) – Whether to apply causal masking.

  • softmax_scale (Optional[float], default=None) – Softmax scale. Defaults to head_dim**-0.5.

  • page_table (Optional[torch.Tensor], default=None) – Enables the paged-KV path. page_table has shape (batch_size, max_pages) and maps batch-local KV block indices to pages. Requires seqused_k. k/v may also be views split from a cache that packs K and V in one 2 * head_dim content dim per token on SM120/SM121 (see supports_packed_kv). Compute capability 10.0/10.3 requires separate contiguous K and V tensors and never copies packed views implicitly, with one exception: packed NVFP4 paged K/V (uint8, (num_pages, 4, 128, 64)) is consumed in place as strided views of a planar [K data | K scale | V data | V scale] page, together with k_scale/v_scale and the two global scales. On that route the block-table width is unconstrained in practice: the per-tile block union is a hash table sized by 8 * topk, not a bitmap indexed by block id, so a 2048-wide table costs the same per block as a 128-wide one.

  • seqused_k (Optional[torch.Tensor], default=None) – Int32 tensor of shape (batch_size,) giving the valid KV length per sequence in the paged path.

  • return_softmax_lse (bool, default=False) – Whether to also return the natural-log LSE with shape (total_q, num_qo_heads).

  • k_scale (Optional[torch.Tensor], default=None) – NVFP4 only: the torch.uint8 byte view (sf.view(torch.uint8)) of the E4M3 block-scale tensor produced by flashinfer.nvfp4_quantize() with sf_vec_size=16 and the swizzled 128x4 layout (one scale per 16 elements, with rows padded to a multiple of 128). Scale rows follow (token, head) order for flat K and (page, head, token) order for paged K. SM120/SM121 only. On compute capability 10.0/10.3 the paged NVFP4 prefill route instead takes the block-scale regions of the packed page as (num_pages, num_kv_heads, page_size, head_dim // 16) views, either uint8 or float8_e4m3fn: K scales linear, V scales (4, 4)-swizzled inside (token, scale index).

  • v_scale (Optional[torch.Tensor], default=None) – NVFP4 block scales for V, with the same dtype, layout, and row-order contract as k_scale.

  • k_global_scale (Optional[float], default=None) – Global dequant scale for K; folds into the softmax scale (NVFP4 K only). On SM100/SM103 it is supported by, and required for, the paged NVFP4 KV prefill route, and by nothing else.

  • v_global_scale (Optional[float], default=None) – Global dequant scale applied to the output on SM120/SM121, for any KV dtype (e.g. an fp8 per-tensor V descale). On SM100/SM103 it is supported by, and required for, the paged NVFP4 KV prefill route, and by nothing else.

  • q_offset (optional) – Optional per-query offset tensor used by specific MSA workflows.

  • return_temperature_lse (bool, default=False) – Whether to also return the MSA temperature LSE. When enabled, the return value is (out, lse, lse_t).

  • lse_temperature_scale (float, default=1.0) – Scale applied to the exponent when computing temperature LSE.

  • workspace (MSASparseAttentionWorkspace, optional) – Caller-owned storage required for CUDA graph capture on compute capability 10.0/10.3. Warm the workspace eagerly with the exact tensors, options, and capture stream before capture. It is not used by the SM120/SM121 backend.

Returns:

out has shape (total_q, num_qo_heads, head_dim). If return_temperature_lse=True, returns (out, lse, lse_t) regardless of return_softmax_lse. Otherwise, setting return_softmax_lse=True returns (out, lse); when both flags are false, returns out. Each returned LSE tensor has shape (total_q, num_qo_heads) and dtype float32.

Return type:

torch.Tensor or tuple of torch.Tensor