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

Sparse decode attention for SM100/SM103/SM107 and SM120/SM121 GPUs.

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. On compute capability 10.0/10.3/10.7, topk must be 16 and Q1 through multi-token decode use the direct persistent M16 path.

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). On the paged path, 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/10.7 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.

  • 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). On the paged path, 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/10.7 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.

  • 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, one per 16 elements. The layout is per architecture. On SM120/SM121 they are uint8 bytes in the swizzled 128x4 layout produced by flashinfer.nvfp4_quantize() (rows padded to a multiple of 128), with scale rows following the cache layout: (token, head) order for flat K/V, (page, head, token) for paged. On compute capability 10.0/10.3/10.7 the paged NVFP4 decode route 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 (torch.Tensor, optional) – NVFP4 only: e4m3 block scales, one per 16 elements. The layout is per architecture. On SM120/SM121 they are uint8 bytes in the swizzled 128x4 layout produced by flashinfer.nvfp4_quantize() (rows padded to a multiple of 128), with scale rows following the cache layout: (token, head) order for flat K/V, (page, head, token) for paged. On compute capability 10.0/10.3/10.7 the paged NVFP4 decode route 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).

  • k_global_scale (float, optional) – Global dequant scales. On SM120/SM121, k_global_scale folds into the softmax scale for NVFP4 K and v_global_scale scales the output for any KV dtype. On SM100/SM103/SM107, both are supported only for uniform FP8 Q/K/V decode and for the paged NVFP4 KV decode route, which requires both.

  • v_global_scale (float, optional) – Global dequant scales. On SM120/SM121, k_global_scale folds into the softmax scale for NVFP4 K and v_global_scale scales the output for any KV dtype. On SM100/SM103/SM107, both are supported only for uniform FP8 Q/K/V decode and for the paged NVFP4 KV decode route, which requires both.

  • 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. The compute capability 10.0/10.3/10.7 backend always uses its native float32 split storage and ignores this override.

  • 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). On compute capability 10.0/10.3/10.7 this argument is accepted for API compatibility; the production direct-M16 route does not split.

  • workspace (MSASparseAttentionWorkspace, optional) –

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

    The packed NVFP4 paged-KV route on compute capability 10.0/10.3/10.7 is the one exception: it captures without a workspace, because everything before its single kernel launch is host-side arithmetic over shapes and strides. Passing one is still honoured, including the warm-vs-capture identity check. What that route does require before capture is one eager launch on the device, from flashinfer.msa_ops.msa_decode_nvfp4_specialized_warmup() or from any real call; msa_decode_nvfp4_specialized_stats()["cuda_graph"] states both facts for a caller that wants to check rather than assume.

  • out (torch.Tensor, optional) –

    Destination for the attention output. When given, the kernel writes straight into it and it is what this function returns – no temporary is allocated and nothing is copied. It must be contiguous, BF16, on q’s device, and exactly q.shape; anything else raises rather than being copied into, because a silent copy here is precisely the cost this parameter exists to remove.

    Supported by the packed-NVFP4 paged-KV route on compute capability 10.0/10.3/10.7. Every other route raises NotImplementedError when it is passed – deliberately, so that a caller cannot be handed the copy back without being told.

Returns:

(batch_size * seqlen_q, num_qo_heads, 128) in q’s dtype, except uniform FP8 Q/K/V returns BF16; plus the natural-log LSE if return_softmax_lse. When out is given it IS the returned tensor.

Return type:

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