flashinfer.sparse¶

Kernels for block sparse flashattention.

class flashinfer.sparse.BlockSparseAttentionWrapper(float_workspace_buffer: Tensor, backend: str = 'auto')¶

Wrapper class for attention computation with a block-sparse matrix as attention mask. The definition of block sparse matrix can be found at bsr_matrix in SciPy.

This API supports any block size (R, C).

Example

>>> import torch
>>> import flashinfer
>>> num_qo_heads = 32
>>> num_kv_heads = 8
>>> head_dim = 128
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> bsr_wrapper = flashinfer.BlockSparseAttentionWrapper(workspace_buffer)
>>> # sparse mask: [[0, 0, 1], [1, 0, 1], [0, 1, 1]]
>>> M = 3
>>> N = 3
>>> indptr = torch.tensor([0, 1, 3, 5], dtype=torch.int32, device="cuda:0")
>>> indices = torch.tensor([2, 0, 2, 1, 2], dtype=torch.int32, device="cuda:0")
>>> bsr_wrapper.plan(
...     indptr,
...     indices,
...     M,
...     N,
...     1, # R(block_rows)=1
...     1, # C(block_columns)=1
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
... )
>>> q = torch.randn((M, num_qo_heads, head_dim), dtype=torch.float16, device="cuda:0")
>>> k = torch.randn((N, num_kv_heads, head_dim), dtype=torch.float16, device="cuda:0")
>>> v = torch.randn((N, num_kv_heads, head_dim), dtype=torch.float16, device="cuda:0")
>>> o = bsr_wrapper.run(q, k, v)
>>> # use dense implementation with attention mask for comparison
>>> mask = torch.tensor([[0, 0, 1], [1, 0, 1], [0, 1, 1]], dtype=torch.bool, device="cuda:0")
>>> o_ref = flashinfer.single_prefill_with_kv_cache(q, k, v, custom_mask=mask)
>>> torch.allclose(o, o_ref)
True

__init__(float_workspace_buffer: Tensor, backend: str = 'auto') → None¶

Constructs of BlockSparseAttentionWrapper.

Parameters:

float_workspace_buffer (torch.Tensor) – The user reserved float workspace buffer used to store intermediate attention results in the split-k algorithm. The recommended size is 128MB, the device of the workspace buffer should be the same as the device of the input tensors.
backend (str) – The implementation backend, could be auto/fa2/fa3 or cake. Defaults to auto. If set to auto, the function will automatically choose the backend based on the device architecture and kernel availability.

Create auxiliary data structures for block sparse attention.

Parameters:

indptr (torch.Tensor, optional) – The block index pointer of the block-sparse matrix on row dimension, shape (MB + 1,), where MB is the number of blocks in the row dimension. Required for all backends except cake, vsa_sm100_blk128, vsa_sm100_blk64, and vsa_sm120_blk64 when block_mask is provided.
indices (torch.Tensor, optional) – The block indices of the block-sparse matrix on column dimension, shape (nnz,), where nnz is the number of non-zero blocks. The elements in indices array should be less then NB: the number of blocks in the column dimension. Required for all backends except cake, vsa_sm100_blk128, vsa_sm100_blk64, and vsa_sm120_blk64 when block_mask is provided.
M (int) – The number of rows of the block-sparse matrix, MB = ceil_div(M, R).
N (int) – The number of columns of the block-sparse matrix, NB = N // C, N should be divisible by C.
R (int) – The number of rows in each block.
C (int) – The number of columns in each block.
num_qo_heads (int) – The number of heads in the query/output tensor.
num_kv_heads (int) – The number of heads in the key/value tensor.
head_dim (int) – The dimension of each head.
mask (torch.Tensor, optional) – The mask tensor with shape (nnz, R, C,), where nnz is the number of non-zero blocks. If every block is full, then we don’t need to provide the mask tensor.
packed_mask (torch.Tensor, optional) – The 1D packed mask tensor, if provided, the custom_mask will be ignored. The packed mask tensor is generated by flashinfer.quantization.packbits().
causal (bool) – Whether to apply causal mask to the attention matrix. This is only effective when custom_mask is not provided in plan().
pos_encoding_mode (str, optional) – The position encoding applied inside attention kernels, could be NONE/ROPE_LLAMA (LLAMA style rotary embedding) /ALIBI. Default is NONE.
use_fp16_qk_reduction (bool) – Whether to use f16 for qk reduction (faster at the cost of slight precision loss).
logits_soft_cap (Optional[float]) – The attention logits soft capping value (used in Gemini, Grok and Gemma-2, etc.), if not provided, will be set to 0. If greater than 0, the logits will be capped according to formula: \(\texttt{logits_soft_cap} \times \mathrm{tanh}(x / \texttt{logits_soft_cap})\), where \(x\) is the input logits.
sm_scale (Optional[float]) – The scale used in softmax, if not provided, will be set to 1.0 / sqrt(head_dim).
rope_scale (Optional[float]) – The scale used in RoPE interpolation, if not provided, will be set to 1.0.
rope_theta (Optional[float]) – The theta used in RoPE, if not provided, will be set to 1e4.
q_data_type (str, optional) – The data type of the query tensor.
kv_data_type (Optional[Union[str, torch.dtype]]) – The data type of the key/value tensor. If None, will be set to q_data_type.
o_data_type (str, optional) – The data type of the output tensor. Default is half. As output dtype cannot be inferred by input dtype in quantization
non_blocking (bool) – Whether to copy the input tensors to the device asynchronously, defaults to True.
block_mask (torch.Tensor, optional) – Per-head block-level attention mask, dtype bool. Shape may be either (num_qo_heads, MB, NB) or (num_kv_heads, MB, NB). block_mask[h, i, j] = True means the Q-block i attends to KV-block j for head h. For GQA (num_qo_heads > num_kv_heads), when providing (num_qo_heads, MB, NB), the first QO-head from each KV-head group is used (sparsity must be the same across QO-heads that share a KV-head). Supported by the cake, vsa_sm100_blk128, vsa_sm100_blk64, and vsa_sm120_blk64 backends. When provided, indptr/indices are not required and will be ignored.
kv_block_lens (torch.Tensor, optional) – Number of valid tokens in every KV block, shape (NB,). Entries must be in [1, C]. Supported by the cake block-64 route; when omitted, every block is treated as having C valid tokens.
q2k_indices (torch.Tensor, optional) – Direct per-head KV-block selections, contiguous int32 with shape (num_qo_heads, MB, topk). Supported by the cake block-64 route and mutually exclusive with block_mask and BSR metadata.
q2k_num (torch.Tensor, optional) – Number of valid entries in each direct selection row, contiguous int32 with shape (num_qo_heads, MB). When omitted, every direct row uses the full topk dimension.
kv_splits (Optional[Union[int, str]]) – Number of KV splits for the split-KV combine path, or "auto" to pick a split count from the sparsity heuristics. Only supported for the vsa_sm100_blk64 backend; must be None for all other backends. None (default) disables splitting, equivalent to passing 1 explicitly. Pass "auto" to select the split count automatically via a sparsity heuristic.
use_clc (Optional[bool]) – Override the SM100 blk64 scheduler: True forces the CLC persistent scheduler, False forces the static scheduler, None (default) uses the shape-based heuristic. Only supported for the vsa_sm100_blk64 backend.
q_scale (torch.Tensor, optional) – Sage FP8 quantization scale for q, shape (1, num_qo_heads, seqlen_q), float32. Only supported for the vsa_sm100_blk64 backend, and only when q/k/v are float8_e4m3fn. Must be provided together with k_scale/v_scale, or not at all. The Sage FP8 path additionally requires batch_size == 1, num_qo_heads in (4, 8), and dense (non-variable) block sparsity – see bsa_attn_sm100_blk64_fwd().
k_scale (torch.Tensor, optional) – Sage FP8 quantization scale for k, shape (1, num_qo_heads, ceil(seqlen_k / 16)), float32. See q_scale.
v_scale (torch.Tensor, optional) – Sage FP8 quantization scale for v, shape (num_qo_heads, head_dim), float32. See q_scale.

:param The plan() method should be called before any run() or: :param run_return_lse() calls: :param auxiliary data structures will be created: :param during this call and cached for multiple kernel runs.: :param The num_qo_heads must be a multiple of num_kv_heads. If num_qo_heads: :param is not equal to num_kv_heads: :param the function will use: :param grouped query attention.: :param .. note::: The vsa_sm100_blk64 backend does not support GQA/MQA: it has no

KV-head mapping and requires num_kv_heads == num_qo_heads.

reset_workspace_buffer(float_workspace_buffer: Tensor, int_workspace_buffer: Tensor) → None¶

Reset the workspace buffer.

Parameters:

float_workspace_buffer (torch.Tensor) – The new float workspace buffer, the device of the new float workspace buffer should be the same as the device of the input tensors.
int_workspace_buffer (torch.Tensor) – The new int workspace buffer, the device of the new int workspace buffer should be the same as the device of the input tensors.

Compute block-sparse attention between Q/K/V tensors.

Parameters:

q (torch.Tensor) – The query tensor with shape (M, num_qo_heads, head_dim).
k (torch.Tensor) – The key tensor with shape (N, num_kv_heads, head_dim).
v (torch.Tensor) – The value tensor with shape (N, num_kv_heads, head_dim).
scale_q (Optional[torch.Tensor]) – The scale tensor for query, per-head quantization with shape: [num_qo_heads]. Used with FP8 Quantization. If not provided, will be set to 1.0.
scale_k (Optional[torch.Tensor]) – The scale tensor for key, per-head quantization with shape: [num_kv_heads]. Used with FP8 Quantization. If not provided, will be set to 1.0.
scale_v (Optional[torch.Tensor]) – The scale tensor for value, per-head quantization with shape: [num_kv_heads]. Used with FP8 Quantization. If not provided, will be set to 1.0.
out (Optional[torch.Tensor]) – The output tensor, if not provided, will be allocated internally.
lse (Optional[torch.Tensor]) – The log-sum-exp of attention logits, if not provided, will be allocated internally.
return_lse (bool) – Whether to return the log-sum-exp of attention logits
enable_pdl (bool) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization Only supported for >= sm90, and currently only for FA2 and CUDA core decode.

Returns:

If return_lse is False, the attention output, shape: [M, num_qo_heads, head_dim]. If return_lse is True, a tuple of two tensors:

The attention output, shape: [M, num_qo_heads, head_dim].
The logsumexp of attention output, shape: [M, num_qo_heads].

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

class flashinfer.sparse.VariableBlockSparseAttentionWrapper(float_workspace_buffer: Tensor, backend: str = 'auto')¶

Wrapper class for attention computation with a block-sparse matrix as attention mask. This API supports variable block sizes provided by block_row_sz and block_col_sz. Besides, each kv_head_idx can specify its own sparse patterns without using the same mask.

Example

>>> import torch
>>> import flashinfer
>>> num_qo_heads = 1
>>> num_kv_heads = 1
>>> head_dim = 128
>>> seq_len = 6 # This corresponds to the `block_row_sz` and `block_col_sz`
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> wrapper = flashinfer.VariableBlockSparseAttentionWrapper(workspace_buffer)
>>> block_mask_map = torch.tensor([[[0, 0, 1], [1, 0, 1], [0, 1, 1]]], dtype=torch.bool, device="cuda:0")
>>> block_row_sz = torch.tensor([[1, 2, 3]], dtype=torch.int32, device="cuda:0")
>>> block_col_sz = torch.tensor([[3, 1, 2]], dtype=torch.int32, device="cuda:0")
>>> wrapper.plan(
...     block_mask_map,
...     block_row_sz,
...     block_col_sz,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
... )
>>> q = torch.randn((num_qo_heads, seq_len, head_dim), dtype=torch.float16, device="cuda:0")
>>> k = torch.randn((num_kv_heads, seq_len, head_dim), dtype=torch.float16, device="cuda:0")
>>> v = torch.randn((num_kv_heads, seq_len, head_dim), dtype=torch.float16, device="cuda:0")
>>> o = wrapper.run(q, k, v)

variable block sparse attention plan function diagram

__init__(float_workspace_buffer: Tensor, backend: str = 'auto') → None¶

Constructs of VariableBlockSparseAttentionWrapper.

Parameters:

float_workspace_buffer (torch.Tensor) – The user reserved float workspace buffer used to store intermediate attention results in the split-k algorithm. The recommended size is 128MB, the device of the workspace buffer should be the same as the device of the input tensors.
backend (str) – The implementation backend, could be auto/fa2 or fa3. Defaults to auto. If set to auto, the function will automatically choose the backend based on the device architecture and kernel availability.

plan(block_mask_map: Tensor, block_row_sz: Tensor, block_col_sz: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim: int, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, logits_soft_cap: float | None = None, sm_scale: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, non_blocking: bool = True, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None) → None¶

Create auxiliary data structures for block sparse attention.

Parameters:

block_mask_map (torch.Tensor) – The block mask map (boolean), shape (num_kv_heads, MB, NB), where MB is the number of blocks in the row dimension, NB is the number of blocks in the column dimension.
block_row_sz (torch.Tensor) – The block row size, shape (num_kv_heads, MB,).
block_col_sz (torch.Tensor) – The block column size, shape (num_kv_heads, NB,).
num_qo_heads (int) – The number of heads in the query/output tensor.
num_kv_heads (int) – The number of heads in the key/value tensor. Note that a group of qo_heads shares the same sparse pattern of kv_heads.
head_dim (int) – The dimension of each head.
causal (bool) – Whether to apply causal mask to the attention matrix.
pos_encoding_mode (str, optional) – The position encoding applied inside attention kernels, could be NONE/ROPE_LLAMA (LLAMA style rotary embedding) /ALIBI. Default is NONE.
use_fp16_qk_reduction (bool) – Whether to use f16 for qk reduction (faster at the cost of slight precision loss).
logits_soft_cap (Optional[float]) – The attention logits soft capping value (used in Gemini, Grok and Gemma-2, etc.), if not provided, will be set to 0. If greater than 0, the logits will be capped according to formula: \(\texttt{logits_soft_cap} \times \mathrm{tanh}(x / \texttt{logits_soft_cap})\), where \(x\) is the input logits.
sm_scale (Optional[float]) – The scale used in softmax, if not provided, will be set to 1.0 / sqrt(head_dim).
rope_scale (Optional[float]) – The scale used in RoPE interpolation, if not provided, will be set to 1.0.
rope_theta (Optional[float]) – The theta used in RoPE, if not provided, will be set to 1e4.
non_blocking (bool) – Whether to copy the input tensors to the device asynchronously, defaults to True.
q_data_type (Union[str, torch.dtype]) – Dtype of the query tensor. Used to specialize the JIT-compiled kernel. Defaults to "float16".
kv_data_type (Optional[Union[str, torch.dtype]]) – Dtype of the key/value tensors. When None, defaults to q_data_type.

The plan() method should be called before any run() or run_return_lse() calls, auxiliary data structures will be created during this call and cached for multiple kernel runs.

The num_qo_heads must be a multiple of num_kv_heads. If num_qo_heads is not equal to num_kv_heads, the function will use grouped query attention.

reset_workspace_buffer(float_workspace_buffer: Tensor, int_workspace_buffer: Tensor) → None¶

Reset the workspace buffer.

Parameters:

float_workspace_buffer (torch.Tensor) – The new float workspace buffer, the device of the new float workspace buffer should be the same as the device of the input tensors.
int_workspace_buffer (torch.Tensor) – The new int workspace buffer, the device of the new int workspace buffer should be the same as the device of the input tensors.

run(q: Tensor, k: Tensor, v: Tensor, out: Tensor | None = None, lse: Tensor | None = None, return_lse: bool = False, enable_pdl: bool | None = None) → Tensor | Tuple[Tensor, Tensor]¶

Compute block-sparse attention between Q/K/V tensors.

Parameters:

q (torch.Tensor) – The query tensor with shape (num_qo_heads, qo_len, head_dim).
k (torch.Tensor) – The key tensor with shape (num_kv_heads, kv_len, head_dim).
v (torch.Tensor) – The value tensor with shape (num_kv_heads, kv_len, head_dim).
out (Optional[torch.Tensor]) – The output tensor, if not provided, will be allocated internally.
lse (Optional[torch.Tensor]) – The log-sum-exp of attention logits, if not provided, will be allocated internally.
return_lse (bool) – Whether to return the log-sum-exp of attention logits
enable_pdl (bool) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization Only supported for >= sm90, and currently only for FA2 and CUDA core decode.

Returns:

If return_lse is False, the attention output, shape: [M, num_qo_heads, head_dim]. If return_lse is True, a tuple of two tensors:

The attention output, shape: [M, num_qo_heads, head_dim].
The logsumexp of attention output, shape: [M, num_qo_heads].

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

Cake VSA¶

`plan_cake_vsa`(indptr, indices, block_mask, ...)	Create stable metadata and workspaces for the source-level backend.
`run_cake_vsa`(plan, q, k, v, *, out, lse, ...)	Run one explicit source-level route; no external fallback is available.

flashinfer.msa_ops¶

Minimax Sparse Attention (MSA) sparse prefill, sparse decode, and top-k selection dispatch on compute capability 10.0/10.3 (SM100/SM103) and SM120/SM121 Blackwell GPUs. The proxy-score operations remain SM120/SM121 only. NVFP4 K/V and views split from a packed paged K/V cache are also SM120/SM121-only; the compute capability 10.0/10.3 attention backend requires separate contiguous K and V tensors and does not make implicit copies. The compute capability 10.0/10.3 backend uses TopK16 as its generic contract and additionally retains four shape-exact routes: paged BF16 decode at B64/Q8/KV65536/TopK32, 512-thread paged BF16 decode at B2/Q1/KV257/TopK4, flat BF16-query/FP8-KV prefill at B3/Q1024/KV8192/TopK8, and paged BF16 prefill at B3/Q4096/KV8192/TopK4. Neighboring non-TopK16 shapes fail closed instead of entering a generic kernel. The decode path uses direct persistent M16 ownership for both Q1 and multi-token decode; it does not route BF16 decode through prefill or split-K. Frozen BF16-query/FP8-KV Q1 serving shapes use exact or transformed direct kernels, while paged uniform FP8 Q/K/V supports Q1 through Q32 and returns BF16 output. Long batch-one BF16 causal prefill uses a selected-block reverse producer and deterministic reduction once the query reaches 8192 tokens. Call flashinfer.msa_ops.supports_packed_kv() with the active device when integrating a cache manager across these architectures; the legacy aggregate SUPPORTS_PACKED_KV flag describes the SM120/SM121 backend. Per-token tensor num_valid_pages for flashinfer.msa_ops.msa_topk_select() is likewise SM120/SM121-only; compute capability 10.0/10.3 requires a scalar value or None and rejects the tensor form before backend dispatch.

CUDA graph capture of sparse prefill or decode on compute capability 10.0/10.3 requires a caller-owned flashinfer.msa_ops.MSASparseAttentionWorkspace. Warm the workspace eagerly with the exact tensors, options, and capture stream before capture. The exact decode overrides are eager-only. The exact TopK8 reverse-prefill route is also eager-only because its reducer uses a host-owned monotonic launch generation. The exact paged TopK4 route normally captures its producer and reducer into an internal two-node CUDA graph; while an outer CUDA graph is being captured it emits the two kernel nodes directly.

Normal callers should leave the Blackwell schedule environment variables unset. For advanced diagnostics and benchmarking, FLASHINFER_MSA_PREFILL_SCHEDULE=m64 forces the eligible M64 prefill schedule. FLASHINFER_MSA_FP8_Q1_SCHEDULE can select batch_attention, q1_exact, q1_flat_xform2, q1_paged_xform2, or paged_uniform_fp8 for an eligible FP8 Q1 route. Unsupported values, or a schedule incompatible with the input layout and dtypes, raise ValueError.

`msa_proxy_score`(q, k, cu_seqlens_q[, ...])	MSA dense proxy pass for SM120/SM121: per-KV-block max attention logits.
`msa_proxy_score_fp4`(q_fp4, k_fp4, q_scale, ...)	NVFP4 MSA dense proxy pass for SM120/SM121 (the FP4 counterpart of `msa_proxy_score()`).
`MSASparseAttentionWorkspace`(device)	Caller-owned storage for SM100/SM103 MSA CUDA graph capture.
`supports_packed_kv`(device)	Return whether MSA accepts packed paged K/V views on `device`.
`msa_sparse_attention`(q, k, v, q2k_indices, ...)	Minimax Sparse Attention forward for SM100/SM103 and SM120/SM121 GPUs.
`msa_sparse_decode_attention`(q, k, v, ...[, ...])	Sparse decode attention for SM100/SM103 and SM120/SM121 GPUs.
`msa_topk_select`(max_score, topk[, ...])	Select the top-K KV blocks per query token based on attention scores.