flashinfer.msa_ops.msa_sparse_decode_attention

flashinfer.msa_ops.msa_sparse_decode_attention(q: Tensor, k: Tensor, v: Tensor, q2k_indices: Tensor, *, page_table: Tensor | None = None, seqused_k: Tensor | None = None, cu_seqlens_k: Tensor | None = None, seqlen_q: int = 1, causal: bool = True, softmax_scale: float | 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, partial_dtype: dtype | None = None, force_fused: bool | None = None)

Sparse decode attention for SM120/SM121.

Computes attention for a decode step: each request contributes seqlen_q query tokens (uniform across the batch) attending only the KV blocks selected in q2k_indices. Decode tokens are right-aligned: token i of a request sits at position seqlen_k - seqlen_q + i.

Parameters:

  • q (torch.Tensor) – (batch_size * seqlen_q, num_qo_heads, 128), bf16 or fp16.

  • k (torch.Tensor) – Paged: (num_pages, num_kv_heads, 128, 128) with page_table and seqused_k; or flat varlen (total_k, num_kv_heads, 128) with cu_seqlens_k. May be fp8 E4M3 (upconverted in-kernel).

  • v (torch.Tensor) – Paged: (num_pages, num_kv_heads, 128, 128) with page_table and seqused_k; or flat varlen (total_k, num_kv_heads, 128) with cu_seqlens_k. May be fp8 E4M3 (upconverted in-kernel).

  • q2k_indices (torch.Tensor) – (num_kv_heads, batch_size * seqlen_q, topk) int32, ascending, -1 tail-padded (the format produced by msa_topk_select()).

  • seqlen_q (int) – Uniform query length per request (e.g. 1, or >1 for speculative decoding).

  • page_table (torch.Tensor, optional) – Page-table mapping for paged KV layout.

  • seqused_k (torch.Tensor, optional) – Per-sequence valid KV-token counts for paged KV layout.

  • cu_seqlens_k (torch.Tensor, optional) – (batch_size + 1,) int32 cumulative KV lengths for ragged KV layout.

  • causal (bool) – Right-aligned causal masking (default True for decode).

  • softmax_scale (float, optional) – Overrides the default attention scaling factor.

  • return_softmax_lse (bool) – If True, also return per-query log-sum-exp values.

  • k_scale (torch.Tensor, optional) – NVFP4 only: e4m3 block scales as uint8 bytes in the swizzled 128x4 layout produced by flashinfer.nvfp4_quantize() (one scale per 16 elements, rows padded to a multiple of 128). Scale rows follow the cache layout: (token, head) order for flat K/V, (page, head, token) for paged.

  • v_scale (torch.Tensor, optional) – NVFP4 only: e4m3 block scales as uint8 bytes in the swizzled 128x4 layout produced by flashinfer.nvfp4_quantize() (one scale per 16 elements, rows padded to a multiple of 128). Scale rows follow the cache layout: (token, head) order for flat K/V, (page, head, token) for paged.

  • k_global_scale (float, optional) – NVFP4 global dequant scales; folded into the softmax scale and the output scale respectively, so the kernel applies only block scales.

  • v_global_scale (float, optional) – NVFP4 global dequant scales; folded into the softmax scale and the output scale respectively, so the kernel applies only block scales.

  • q_offset (int or torch.Tensor, optional) – Optional query-position offset used by causal alignment.

  • partial_dtype (torch.dtype, optional) – Accumulator / partial-result dtype override for supported kernels.

  • force_fused (bool, optional) – Override the adaptive split-K decision. By default each token’s selected list is split into chunks (one CTA per chunk online-softmaxes its blocks and writes a partial; the combine kernel reduces them), with the chunk count adapting to fill the SMs: one block per chunk at low batch, a single chunk at high batch. A single chunk is the fused path: one CTA per (token, kv-head) writes the final output directly (no GMEM partials, no combine). True/False force fused/split on; None (default) adapts. NVFP4 KV defaults to the per-block split at every batch size (the in-kernel dequant favors the extra parallelism).

Returns:

(batch_size * seqlen_q, num_qo_heads, 128) in q’s dtype; plus the natural-log LSE if return_softmax_lse.

Return type:

torch.Tensor or (torch.Tensor, torch.Tensor)