FlashInfer Attention Kernels¶

Experimental Task-Scheduled Attention¶

The experimental Blackwell task-scheduled FMHA context, FMHA decode, block-sparse FMHA, and MLA decode APIs are imported from flashinfer.attention.prims_ts. Scheduling, tile selection, and split-KV reduction are automatic implementation details; there are no public tuning knobs.

See the PrimTS guide index for the public entry points, supported contracts, and examples. Current accuracy and performance signoff is on SM100a/B200; SM103a/B300 is architecture-gated but not yet signoff-qualified.

Calling these APIs is an explicit opt-in and emits an ExperimentalWarning once per decorated function. They provide no API compatibility guarantee; generated stable API reference entries are deferred until graduation. Logging and existing fi_trace bindings remain available.

QToken-KvBlock-Sparse-Attention¶

QToken-KvBlock-Sparse-Attention consumes per-query indexer_block_ids[total_q, block_topk] and a dense physical block_table. Packed prefill uses [total_q, Hq, D] with qo_indptr; fixed MTP decode uses [B, Nq, G, Hq, D]. kv_block_size is the semantic sparse K/V atom and currently supports only four tokens. The wrapper plans capacity outside CUDA Graph capture and runs live route metadata on the hot path.

flashinfer.decode¶

Single Request Decoding¶

`single_decode_with_kv_cache`()	Decode attention with a KV cache for a single request and return the attention output.
`single_decode_with_kv_cache_with_jit_module`(...)	Single-request decode using a pre-compiled JIT module.

Batch Decoding¶

`cudnn_batch_decode_with_kv_cache`(q, k_cache, ...)	Batched decode attention with paged KV cache, backed by cuDNN SDPA.
`trtllm_batch_decode_with_kv_cache`(query, ...)
`xqa_batch_decode_with_kv_cache`(query, ...[, ...])

DCP Speculative Decode Workspace¶

The native Cake FMHA DCP speculative route of flashinfer.decode.trtllm_batch_decode_with_kv_cache() uses caller-owned scratch buffers so a prewarmed invocation can be captured in a CUDA Graph. It is also reachable through flashinfer.cake_fmha.cake_batch_decode_with_kv_cache(); the non-null causal_seqlens_kv_global argument is the explicit add-on selection key. On SM103, the same Cake entrypoint accepts a device request_order tensor for BF16-query, FP8-E4M3 paged decode with head dimension 256. Precompute an optional immutable length-aware schedule with flashinfer.plan_cake_fmha_request_ordered_paged_decode() before graph capture. Page-table rows must be padded to 4 * ceil(max_seq_len / 256) entries, and the exact tensor/workspace binding must be invoked once eagerly to initialize its TMA descriptors before capture. After that prewarm, changing only the order tensor contents does not require recapture.

`get_dcp_spec_workspace_size_bytes`(...[, ...])	Bytes for Cake FMHA Split-KV BF16 partial-O and FP32 partial-LSE scratch.
`get_dcp_spec_counter_bytes`(batch_size, ...)	Bytes for the v4 completion tickets, zeroed once then self-reset.
`plan_cake_fmha_request_ordered_paged_decode`(...)	Select an exported schedule from host-visible immutable metadata.
`CakeFmhaRequestOrderedDecodePlan`(...)	Immutable host-selected launch plan for graph-safe request ordering.

class flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, use_tensor_cores: bool = False, paged_kv_indptr_buffer: Tensor | None = None, paged_kv_indices_buffer: Tensor | None = None, paged_kv_last_page_len_buffer: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None)¶

Wrapper class for decode attention with a paged KV cache (first proposed in vLLM) for a batch of requests.

Check our tutorial for page table layout.

Examples

>>> import torch
>>> import flashinfer
>>> num_layers = 32
>>> num_qo_heads = 64
>>> num_kv_heads = 8
>>> head_dim = 128
>>> max_num_pages = 128
>>> page_size = 16
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> decode_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
...     workspace_buffer, "NHD"
... )
>>> batch_size = 7
>>> kv_page_indices = torch.arange(max_num_pages).int().to("cuda:0")
>>> kv_page_indptr = torch.tensor(
...     [0, 17, 29, 44, 48, 66, 100, 128], dtype=torch.int32, device="cuda:0"
... )
>>> # 1 <= kv_last_page_len <= page_size
>>> kv_last_page_len = torch.tensor(
...     [1, 7, 14, 4, 3, 1, 16], dtype=torch.int32, device="cuda:0"
... )
>>> kv_cache_at_layer = [
...     torch.randn(
...         max_num_pages, 2, page_size, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
...     ) for _ in range(num_layers)
... ]
>>> # create auxiliary data structures for batch decode attention
>>> decode_wrapper.plan(
...     kv_page_indptr,
...     kv_page_indices,
...     kv_last_page_len,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
...     page_size,
...     pos_encoding_mode="NONE",
...     data_type=torch.float16
... )
>>> outputs = []
>>> for i in range(num_layers):
...     q = torch.randn(batch_size, num_qo_heads, head_dim).half().to("cuda:0")
...     kv_cache = kv_cache_at_layer[i]
...     # compute batch decode attention, reuse auxiliary data structures for all layers
...     o = decode_wrapper.run(q, kv_cache)
...     outputs.append(o)
...
>>> outputs[0].shape
torch.Size([7, 64, 128])

Note

To accelerate computation, FlashInfer’s batch decode attention creates some auxiliary data structures, these data structures can be reused across multiple batch decode attention calls (e.g. different Transformer layers). This wrapper class manages the lifecycle of these data structures.

__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, use_tensor_cores: bool = False, paged_kv_indptr_buffer: Tensor | None = None, paged_kv_indices_buffer: Tensor | None = None, paged_kv_last_page_len_buffer: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None) → None¶

Constructor of BatchDecodeWithPagedKVCacheWrapper.

Parameters:

float_workspace_buffer (torch.Tensor. Must be initialized to 0 for its first use.) – 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. The buffer must be 16-byte aligned; tensors created by torch.empty satisfy this on supported devices.
kv_layout (str) – The layout of the input k/v tensors, could be either NHD or HND.
use_cuda_graph (bool) – Whether to enable CUDAGraph for batch decode attention, if enabled, the auxiliary data structures will be stored as the provided buffers. The batch_size cannot change during the lifecycle of this wrapper when CUDAGraph is enabled.
use_tensor_cores (bool) – Whether to use tensor cores for the computation. Will be faster for large group size in grouped query attention. Defaults to False.
paged_kv_indptr_buffer (Optional[torch.Tensor]) – The user reserved buffer on GPU to store the indptr of the paged kv cache, the size of the buffer should be [batch_size + 1]. Only needed when use_cuda_graph is True.
paged_kv_indices_buffer (Optional[torch.Tensor]) – The user reserved buffer on GPU to store the page indices of the paged kv cache, should be large enough to store the maximum number of page indices (max_num_pages) during the lifecycle of this wrapper. Only needed when use_cuda_graph is True.
paged_kv_last_page_len_buffer (Optional[torch.Tensor]) – The user reserved buffer on GPU to store the number of entries in the last page, the size of the buffer should be [batch_size]. Only needed when use_cuda_graph is True.
backend (str) – The implementation backend, could be auto/fa2/fa3/trtllm-gen/ cute-dsl/cudnn. Defaults to auto. If set to auto, the wrapper will automatically choose the backend based on the device architecture and kernel availability. On SM100 / SM103 with cudnn-frontend 1.30+ installed, auto resolves to cudnn for multi-token decode rows (2 <= q_len_per_req <= 4, speculative / MTP verification) of fp16 / bf16 head_dim-128 GQA models whose group divides 128, without RoPE / soft-cap / window, when cuDNN’s default SM100 engine serves them with its FROST decode tile and that tile has enough work: q_len_per_req * group in [32, 128] rows per CTA and batch_size * num_kv_heads >= 64 CTAs, where it measures at 0.35-0.95x fa2 (see _auto_decode_prefers_cudnn); single-token decode stays on fa2, except head_dim-256 GQA models (groups 4, 8, 16) with batch_size * num_kv_heads in [96, 256], where cuDNN’s d256 decode tile measures at 0.5-0.95x fa2 (from 64 under use_cuda_graph=True on cudnn-frontend 1.31+, which is told the graph is replayed and splits KV). Multi-token rows without tensor cores have no fa2 kernel, so there auto takes cudnn for every plan its decode path can run (such a plan used to raise). A caller-owned block_tables narrower than 2048 tokens keeps fa2 on cudnn-frontend 1.30 (its placement puts the backend’s prefill-class engine first below that; 1.31+ lifts the limit). auto takes cudnn only where the frontend’s FROST engines can run (its own CuTe-DSL check: nvidia-cutlass-dsl 4.7.0 or newer, above the 4.6.2 the packages’ dependency metadata admit), and callers that pin fa2’s KV split (fixed_split_size / disable_split_kv, the batch-invariance controls) keep fa2 wherever it can serve the rows. resolved_backend reports the choice after plan(). FLASHINFER_DECODE_AUTO_CUDNN=0 keeps fa2, =1 applies the rule on an older frontend too. Under CUDA graphs (use_cuda_graph=True) auto takes cudnn only when plan() receives a caller-owned block_tables (the auto-built table is captured at its first width and cannot grow) or when fa2 cannot serve the rows (q_len_per_req > 1 without tensor cores); the resolution then becomes part of the frozen graph shape and a later plan() may not change it. The cute-dsl backend uses the CuTe DSL GQA decode kernel for Blackwell (SM100+) and only supports a subset of features (equal head_dim_qk/vo, no RoPE/ALiBi/soft-cap). The cudnn backend runs cuDNN’s paged SDPA decode graph (flashinfer.decode.cudnn_batch_decode_with_kv_cache()). It supports fp16/bf16 q/kv/o of one dtype, GQA, return_lse, CUDA graphs, multi-token decode (q_len_per_req > 1, bottom-right causal), a left sliding window (window_left) and attention sinks, but not RoPE, soft-cap, fp8 or NVFP4 KV. A sink at q_len_per_req == 1 is served when the cuDNN stack’s SDPA engines accept it (cudnn-frontend 1.30+ with the FROST engines enabled); the backend engine raises a not-supported error at the first run. Pass a fixed block_tables to plan() when capturing CUDA graphs so the captured block-table buffer is stable across re-plans. On SM90, tensor-core auto graph-mode wrappers may select XQA MTP for FP16 NHD, page16, head128, 32Q/8KV-head plans with uniform KV lengths: q8/KV8192 at batch4/16/64, or q4/q8/KV1024 at batch16. This provisional rule excludes ordinary wrappers, explicit backends, q2 and unmeasured shapes. It is based on selected historical H100 graph results, not a guarantee of gains on every SM90 GPU. The preferred route is frozen for graph-mode wrappers; changes in eligibility require a separate wrapper and graph capture. Eager warmups of such wrappers use their selected graph-mode route. workspace_size is unavailable for XQA-selected shapes.
jit_args (Optional[List[Any]]) – If provided, the wrapper will use the provided arguments to create the JIT module, otherwise, the wrapper will use default attention implementation.
jit_kwargs (Optional[Dict[str, Any]]) – The keyword arguments to create the JIT module, defaults to None.

plan(indptr: Tensor, indices: Tensor, last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim: int, page_size: int, *deprecated_positional_args: Any, **kwargs: Any) → None¶

Plan batch decode for given problem specification.

Parameters:

indptr (torch.Tensor) – The indptr of the paged kv cache, shape: [batch_size + 1], dtype: torch.int32
indices (torch.Tensor) – The page indices of the paged kv cache, shape: [indptr[-1]], dtype: torch.int32
last_page_len (torch.Tensor) – The number of entries in the last page of each request in the paged kv cache, shape: [batch_size], dtype: torch.int32
num_qo_heads (int) – The number of query/output heads
num_kv_heads (int) – The number of key/value heads
head_dim (int) – The dimension of the heads
page_size (int) – The page size of the paged kv cache
pos_encoding_mode (str) – The position encoding applied inside attention kernels, could be NONE/ROPE_LLAMA (LLAMA style rotary embedding) /ALIBI. Defaults to NONE.
window_left (int) – The left (inclusive) window size for the attention window, when set to -1, the window size will be set to the full length of the sequence. Defaults to -1.
window_right (int) – The right (inclusive) window size for the attention window. -1 disables the right bound. Defaults to 0. Currently window_right != 0 only supported by the cute-dsl backend.
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.
q_data_type (Optional[Union[str, torch.dtype]]) – The data type of the query tensor, defaults torch.float16.
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. Defaults to None.
o_data_type (Optional[Union[str, torch.dtype]]) – The data type of the output tensor. If None, will be set to q_data_type. For FP8 inputs, this should typically be set to torch.float16 or torch.bfloat16.
data_type (Optional[Union[str, torch.dtype]]) – The data type of both the query and key/value tensors. Defaults to torch.float16. data_type is deprecated, please use q_data_type and kv_data_type instead.
sm_scale (Optional[float]) – Softmax scale. If None, defaults to 1 / sqrt(head_dim). Cached on the wrapper and reused at run() time.
rope_scale (Optional[float]) – Scale factor applied during RoPE interpolation. Only consulted when pos_encoding_mode != "NONE". Defaults to 1.0 when None.
rope_theta (Optional[float]) – Base value for the RoPE frequencies. Only consulted when pos_encoding_mode != "NONE". Defaults to 1e4 when None.
non_blocking (bool) – Whether to copy the input tensors to the device asynchronously, defaults to True.
seq_lens (Optional[torch.Tensor]) – A uint32 1D tensor indicating the kv sequence length of each prompt. shape: [batch_size].
block_tables (Optional[torch.Tensor]) – A uint32 2D tensor indicating the block table of each prompt. shape: [batch_size, max_num_blocks_per_seq].
fixed_split_size (Optional[int],) – The fixed split size for FA2 split-kv decode, in pages. Only supported by tensor core decode for now. Recommend setting to the average sequence length of your workload. When enabled for FA2, will lead to deterministic softmax score reduction in the merge_states kernel, and therefore batch-size invariant outputs. See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ Note that compatibility with CUDA graph is NOT guaranteed, as even when bs is fixed, kv seq len can change and lead to a varied number of launched CTAs.
disable_split_kv (bool,) – Whether to disable the split-kv for determinism in CUDA Graph, defaults to False.
q_len_per_req (int) – The number of query tokens per request. Defaults to 1. This must be set at plan() time to enable the XQA MTP decode fast path; setting it only at run() time will not enable that path. q_len_per_req > 1 is currently supported on the fa2 tensor-core backend (and natively by trtllm-gen/cute-dsl). In CUDA graph mode, an auto-built XQA block table keeps its initial capacity. Pass block_tables sized for the longest sequence the graph will see if the page count will grow. Under use_cuda_graph, this value is part of the frozen shape (like the batch size): once the wrapper has been planned, re-planning with a different value raises.
is_causal (Optional[bool]) – Whether the mask is causal within each request block. Defaults to None, which derives it from q_len_per_req > 1. Values that differ from that default are not supported by this wrapper.
use_token_head_sf (bool) – Whether the FP8 KV cache uses the per-(token, head) scale layout: a separate float32 scale tensor (passed via kv_cache_sf in run()) dequantizes the KV cache inside the kernel. When enabled, the KV cache is FP8 (e4m3/ e5m2) with standard layout (last dim head_dim), q must be fp16 (or bf16 on SM80+), and the fa2 tensor-core backend is forced. Defaults to False.

Note

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 run calls.

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.

The plan() method cannot be used in Cuda Graph or in torch.compile.

Optional arguments after page_size are accepted positionally for backward compatibility, but that calling convention is deprecated and scheduled for removal in a future release. Pass them by keyword instead. window_right is keyword-only.

prewarm_paged_kv_stride_variant(variant: str = 'independent') → None¶

Load a lazy tensor-core paged-KV-stride variant after plan().

Call this method before CUDA graph capture when a planned standard FA2 tensor-core decode wrapper will receive K and V tensors with different data strides.

Parameters:: variant (str) – The paged-KV-stride variant to prewarm. The only supported value is "independent" (the default), for K and V with different data strides.

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.

property resolved_backend: str¶

"cudnn" when the auto rule picked cuDNN, the tensor-core pick ("fa2" / "fa3") otherwise, "fa2" for the CUDA-core decode kernel. An explicit backend reports itself; a tensor-core auto wrapper reports "auto" before its first plan(). An auto-selected XQA MTP plan reports "xqa"; unsupported per-call options can still use the fallback.

Type:: What the last plan() resolved backend="auto" to

Compute batch decode attention between query and paged kv cache.

Parameters:

q (torch.Tensor) – The query tensor, shape: [batch_size * q_len_per_req, num_qo_heads, head_dim] On the trtllm-gen and cute-dsl backends the q_len_per_req implied by q.shape[0] may differ from the planned value; the fa2/fa3 tensor-core path requires it to match plan(). XQA MTP is selected only for the documented SM90 graph-mode auto shapes; other plans retain the existing route.
paged_kv_cache (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) –
The paged KV-Cache stored as a tuple of tensors or a single tensor:
- a tuple (k_cache, v_cache) of 4-D tensors, each with shape: [max_num_pages, page_size, num_kv_heads, head_dim] if kv_layout is NHD, and [max_num_pages, num_kv_heads, page_size, head_dim] if kv_layout is HND.
- a single 5-D tensor with shape: [max_num_pages, 2, page_size, num_kv_heads, head_dim] if kv_layout is NHD, and [max_num_pages, 2, num_kv_heads, page_size, head_dim] if kv_layout is HND. Where paged_kv_cache[:, 0] is the key-cache and paged_kv_cache[:, 1] is the value-cache.
For NVFP4 torch.uint8 caches, the last dimension is head_dim // 2 because each byte stores two FP4 values.
*args – Additional arguments for the custom kernel.
q_scale (Optional[float]) – The calibration scale of query for fp8 input, if not provided, will be set to 1.0.
k_scale (Optional[float]) – The calibration scale of key for fp8 or nvfp4 input, if not provided, will be set to 1.0.
v_scale (Optional[float]) – The calibration scale of value for fp8 or nvfp4 input, if not provided, will be set to 1.0.
out (Optional[torch.Tensor]) – The output tensor, if not provided, will be allocated internally. Must be zero-init for cute-dsl backend.
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 logsumexp of attention scores, defaults to False.
enable_pdl (bool) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/programmatic-dependent-launch.html Only supported for >= sm90, and currently only for FA2 and CUDA core decode.
window_left (Optional[int]) – Per-call sliding-window bound. Must either be None (default, inherit the value from plan()) or equal the value passed to plan() (the kernel asserts this). Passing -1 to plan() disables the sliding window for the entire batch.
sinks (Optional[torch.Tensor]) – Per-head attention sink logits, shape [num_qo_heads]. When provided, sinks[head_idx] is appended to each row of the softmax denominator (Streaming-LLM / Attention-Sinks). The dtype requirement is backend-specific and validated by the underlying kernel; pass None to disable.
q_len_per_req (Optional[int]) – DEPRECATED — pass to plan() instead. When provided here, emits a DeprecationWarning and is used to validate the run-time value inferred from q.size(0). Scheduled for removal in a future release.
skip_softmax_threshold_scale_factor (Optional[float] = None) – threshold scale factor for skipping softmax operations. Providing a value for this parameter enables skip-softmax sparsity as described in: https://arxiv.org/abs/2512.12087 If no value is provided, then standard attention is used. Setting the threshold to a higher value generally increases kernel performance at the cost of accuracy degradation. The actual threshold value equals the provided threshold_scale_factor divided by the context length.
kv_cache_sf (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]) –
Per-block scale factors for NVFP4 KV cache. Accepts the same formats as paged_kv_cache:
- a tuple (k_scales, v_scales) of 4-D tensors, each with shape: [num_pages, page_size, num_kv_heads, head_dim // 16] if kv_layout is NHD, and [num_pages, num_kv_heads, page_size, head_dim // 16] if kv_layout is HND.
- a single 5-D tensor with shape: [num_pages, 2, page_size, num_kv_heads, head_dim // 16] if kv_layout is NHD, and [num_pages, 2, num_kv_heads, page_size, head_dim // 16] if kv_layout is HND, where dim 1 holds k (index 0) and v (index 1) scales.
Both tensors have dtype torch.float8_e4m3fn.

Currently, NVFP4 KV supports the fa2 and trtllm-gen backends; fa2 requires use_tensor_cores=True.

With use_token_head_sf=True (plan time), kv_cache_sf must be a (k_sf, v_sf) tuple of float32 tensors mirroring the KV layout minus head_dim: [num_pages, page_size, num_kv_heads] if kv_layout is NHD, and [num_pages, num_kv_heads, page_size] if HND. The tensors may be non-contiguous views (e.g. sliced out of an inline slot layout).

Returns:

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

attention output, shape: [batch_size * q_len_per_req, num_qo_heads, head_dim]
logsumexp of attention scores, shape: [batch_size * q_len_per_req, num_qo_heads].

Return type:

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

workspace_size(indptr: Tensor, indices: Tensor, last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim: int, page_size: int, pos_encoding_mode: str = 'NONE', window_left: int = -1, logits_soft_cap: float | None = None, q_data_type: str | dtype | None = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, data_type: str | dtype | None = None, sm_scale: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, block_tables: Tensor | None = None, seq_lens: Tensor | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False, q_len_per_req: int = 1) → Tuple[int, int]¶

Return the caller-owned workspace size required by plan().

The returned tuple is (float_workspace_size, int_workspace_size) in bytes. The inputs follow plan(); host-side planning tensors such as indptr and last_page_len are copied to CPU in the same way as plan(). The wrapper’s float workspace buffer is only used as the device/stream selector for the underlying module query. This method does not allocate buffers and does not mutate cached plan state.

Parameters:

indptr (torch.Tensor) – The indptr of the paged kv cache, shape: [batch_size + 1], dtype: torch.int32
indices (torch.Tensor) – The page indices of the paged kv cache, shape: [indptr[-1]], dtype: torch.int32
last_page_len (torch.Tensor) – The number of entries in the last page of each request in the paged kv cache, shape: [batch_size], dtype: torch.int32
num_qo_heads (int) – The number of query/output heads.
num_kv_heads (int) – The number of key/value heads.
head_dim (int) – The dimension of the heads.
page_size (int) – The page size of the paged kv cache.
pos_encoding_mode (str) – The position encoding applied inside attention kernels, could be NONE/ROPE_LLAMA (LLAMA style rotary embedding) /ALIBI. Defaults to NONE.
window_left (int) – The left (inclusive) window size for the attention window, when set to -1, the window size will be set to the full length of the sequence. Defaults to -1.
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.
q_data_type (Optional[Union[str, torch.dtype]]) – The data type of the query tensor, defaults to torch.float16.
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 (Optional[Union[str, torch.dtype]]) – The data type of the output tensor. If None, will be set to q_data_type.
data_type (Optional[Union[str, torch.dtype]]) – Deprecated alias — sets both q_data_type and kv_data_type when they are not provided explicitly.
sm_scale (Optional[float]) – Softmax scale. If None, defaults to 1 / sqrt(head_dim).
rope_scale (Optional[float]) – Scale factor applied during RoPE interpolation. Defaults to 1.0 when None.
rope_theta (Optional[float]) – Base value for the RoPE frequencies. Defaults to 1e4 when None.
block_tables (Optional[torch.Tensor]) – Unused by this method; accepted for signature compatibility with plan().
seq_lens (Optional[torch.Tensor]) – A uint32 1D tensor indicating the kv sequence length of each prompt, shape: [batch_size].
fixed_split_size (Optional[int]) – The fixed split size for FA2 split-kv decode, in pages.
disable_split_kv (bool) – Whether to disable the split-kv for determinism in CUDA Graph. Defaults to False.
q_len_per_req (int) – The number of query tokens per request. Defaults to 1.

Returns:

(float_workspace_size, int_workspace_size) in bytes.

Return type:

Tuple[int, int]

Example

>>> float_bytes, int_bytes = wrapper.workspace_size(...)
>>> wrapper.reset_workspace_buffer(
...     torch.empty(float_bytes, dtype=torch.uint8, device="cuda"),
...     torch.empty(int_bytes, dtype=torch.uint8, device="cuda"),
... )
>>> wrapper.plan(...)

class flashinfer.decode.CUDAGraphBatchDecodeWithPagedKVCacheWrapper(workspace_buffer: Tensor, indptr_buffer: Tensor, indices_buffer: Tensor, last_page_len_buffer: Tensor, kv_layout: str = 'NHD', use_tensor_cores: bool = False)¶

CUDAGraph-compatible wrapper class for decode attention with a paged KV cache (first proposed in vLLM) for a batch of requests.

Note that this wrapper may not be as efficient as BatchDecodeWithPagedKVCacheWrapper because we won’t dispatch to different kernels for different batch sizes/sequence lengths/etc to accommodate the CUDAGraph requirement.

Check our tutorial for page table layout.

Note

The plan() method could not be captured by CUDAGraph.

XQA¶

`xqa`(q, k_cache, v_cache, page_table, ...[, ...])	Apply attention with paged KV cache using XQA kernel.
`xqa_mla`(q, k_cache, v_cache, page_table, ...)	Apply attention with paged KV cache using XQA MLA (Multi-Head Latent Attention) kernel.

flashinfer.prefill¶

Attention kernels for prefill & append attention in both single request and batch serving setting.

Single Request Prefill/Append Attention¶

`single_prefill_with_kv_cache`()	Prefill/append attention with a KV cache for a single request and return the attention output.
`single_prefill_with_kv_cache_return_lse`(q, k, v)	Convenience wrapper for `single_prefill_with_kv_cache()` that always returns LSE.
`single_prefill_with_kv_cache_with_jit_module`(...)	Single-request prefill / append attention using a pre-compiled JIT module.

Batch Prefill/Append Attention¶

`cudnn_batch_prefill_with_kv_cache`(q, ...[, ...])	Batched prefill attention with paged KV cache, backed by cuDNN SDPA.
`trtllm_batch_context_with_kv_cache`(query, ...)
`trtllm_ragged_attention_deepseek`(query, key, ...)
`fmha_v2_prefill_deepseek`(query, key, value, ...)	Run causal SM12x FMHA v2 for MLA context or 128-dimension GQA.
`trtllm_fmha_v2_prefill`(qkv, input_layout, ...)	TRT-LLM FMHAv2 prefill attention.
`fmha_v2_prefill_sm120`(query, key, value, ...)	Run SM120 FMHA v2 with contiguous, separate Q/K/V tensors.

class flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None, variant_owns_mask: bool = False, max_total_num_rows: int | None = None)¶

Wrapper class for prefill/append attention with a paged KV cache for a batch of requests.

Check our tutorial for page table layout.

Example

>>> import torch
>>> import flashinfer
>>> num_layers = 32
>>> num_qo_heads = 64
>>> num_kv_heads = 16
>>> head_dim = 128
>>> max_num_pages = 128
>>> page_size = 16
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> prefill_wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
...     workspace_buffer, "NHD"
... )
>>> batch_size = 7
>>> nnz_qo = 100
>>> qo_indptr = torch.tensor(
...     [0, 33, 44, 55, 66, 77, 88, nnz_qo], dtype=torch.int32, device="cuda:0"
... )
>>> paged_kv_indices = torch.arange(max_num_pages).int().to("cuda:0")
>>> paged_kv_indptr = torch.tensor(
...     [0, 17, 29, 44, 48, 66, 100, 128], dtype=torch.int32, device="cuda:0"
... )
>>> # 1 <= paged_kv_last_page_len <= page_size
>>> paged_kv_last_page_len = torch.tensor(
...     [1, 7, 14, 4, 3, 1, 16], dtype=torch.int32, device="cuda:0"
... )
>>> q_at_layer = torch.randn(num_layers, nnz_qo, num_qo_heads, head_dim).half().to("cuda:0")
>>> kv_cache_at_layer = torch.randn(
...     num_layers, max_num_pages, 2, page_size, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
... )
>>> # create auxiliary data structures for batch prefill attention
>>> prefill_wrapper.plan(
...     qo_indptr,
...     paged_kv_indptr,
...     paged_kv_indices,
...     paged_kv_last_page_len,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
...     page_size,
...     causal=True,
... )
>>> outputs = []
>>> for i in range(num_layers):
...     q = q_at_layer[i]
...     kv_cache = kv_cache_at_layer[i]
...     # compute batch prefill attention, reuse auxiliary data structures
...     o = prefill_wrapper.run(q, kv_cache)
...     outputs.append(o)
...
>>> outputs[0].shape
torch.Size([100, 64, 128])
>>>
>>> # below is another example of creating custom mask for batch prefill attention
>>> mask_arr = []
>>> qo_len = (qo_indptr[1:] - qo_indptr[:-1]).cpu().tolist()
>>> kv_len = (page_size * (paged_kv_indptr[1:] - paged_kv_indptr[:-1] - 1) + paged_kv_last_page_len).cpu().tolist()
>>> for i in range(batch_size):
...     mask_i = torch.tril(
...         torch.full((qo_len[i], kv_len[i]), True, device="cuda:0"),
...         diagonal=(kv_len[i] - qo_len[i]),
...     )
...     mask_arr.append(mask_i.flatten())
...
>>> mask = torch.cat(mask_arr, dim=0)
>>> prefill_wrapper.plan(
...     qo_indptr,
...     paged_kv_indptr,
...     paged_kv_indices,
...     paged_kv_last_page_len,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
...     page_size,
...     custom_mask=mask,
... )
>>> for i in range(num_layers):
...     q = q_at_layer[i]
...     kv_cache = kv_cache_at_layer[i]
...     # compute batch prefill attention, reuse auxiliary data structures
...     o_custom = prefill_wrapper.run(q, kv_cache)
...     assert torch.allclose(o_custom, outputs[i], rtol=1e-3, atol=1e-3)
...

Note

To accelerate computation, FlashInfer’s batch prefill/append attention operators create some auxiliary data structures, these data structures can be reused across multiple prefill/append attention calls (e.g. different Transformer layers). This wrapper class manages the lifecycle of these data structures.

__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None, variant_owns_mask: bool = False, max_total_num_rows: int | None = None) → None¶

Constructor of BatchPrefillWithPagedKVCacheWrapper.

Parameters:

float_workspace_buffer (torch.Tensor) – The user reserved workspace buffer used to store intermediate attention results in 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. The buffer must be 16-byte aligned; tensors created by torch.empty satisfy this on supported devices.
kv_layout (str) – The layout of the input k/v tensors, could be either NHD or HND.
use_cuda_graph (bool) – Whether to enable CUDA graph capture for the prefill kernels, if enabled, the auxiliary data structures will be stored in provided buffers. The batch_size cannot change during the lifecycle of this wrapper when CUDAGraph is enabled. With backend="auto", the first successful FA2/FA3 plan fixes the backend for subsequent plans so captured scheduling metadata stays compatible. Inputs requiring a different backend need a new wrapper.
qo_indptr_buf (Optional[torch.Tensor]) – The user reserved buffer to store the qo_indptr array, the size of the buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True.
paged_kv_indptr_buf (Optional[torch.Tensor]) – The user reserved buffer to store the paged_kv_indptr array, the size of this buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True.
paged_kv_indices_buf (Optional[torch.Tensor]) – The user reserved buffer to store the paged_kv_indices array, should be large enough to store the maximum possible size of the paged_kv_indices array during the lifetime of the wrapper. This argument is only effective when use_cuda_graph is True.
paged_kv_last_page_len_buf (Optional[torch.Tensor]) – The user reserved buffer to store the paged_kv_last_page_len array, the size of the buffer should be [batch_size]. This argument is only effective when use_cuda_graph is True.
custom_mask_buf (Optional[torch.Tensor]) – The user reserved buffer to store the custom mask tensor, should be large enough to store the maximum possible size of the packed custom mask tensor during the lifetime of the wrapper. This argument is only effective when use_cuda_graph is set to True and the custom mask will be used in attention computation.
mask_indptr_buf (Optional[torch.Tensor]) – The user reserved buffer to store the mask_indptr array, the size of the buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True and the custom mask will be used in attention computation.
backend (str) – The implementation backend, could be auto/fa2/fa3/cudnn/trtllm-gen cute-dsl/cute-dsl-prims. Defaults to auto. If set to auto, the wrapper will automatically choose the backend based on the device architecture and kernel availability. In CUDA-graph mode, once FA2 has been planned, later plans keep FA2 so its captured scheduling metadata is refreshed. Use a separate wrapper to select cuDNN for a new graph after planning FA2. The cute-dsl backend uses the CuTe DSL attention kernel for Blackwell (SM100+). cute-dsl-prims backend is a SM120-only FP8 attention backend implemented using CUTLASS primitives.
jit_args (Optional[List[Any]]) – If provided, the wrapper will use the provided arguments to create the JIT module, otherwise, the wrapper will use default attention implementation.
jit_kwargs (Optional[Dict[str, Any]]) – The keyword arguments to create the JIT module, defaults to None.
variant_owns_mask (bool) – If True, the attention variant supplied through jit_args computes the complete attention mask in its LogitsMask hook, and MaskMode.CUSTOM is selected without a mask tensor: the kernel evaluates LogitsMask on every KV tile instead of only on the causal/window boundary tiles, and no custom_mask/packed_custom_mask needs to be passed to plan(). Requires jit_args, because the default attention variant dereferences the custom mask buffer under MaskMode.CUSTOM. Only supported with backend="fa2" (the SM90 batch prefill kernels reject MaskMode.CUSTOM), and incompatible with prefix_len_ptr (multi-item scoring), which selects a different mask mode. Defaults to False.
max_total_num_rows (Optional[int]) – Maximum total number of query tokens across requests in CUDA graph mode. Declaring this capacity lets a later plan() accept more query tokens than the first call while retaining the same metadata buffers. If omitted, the first plan() fixes the capacity. Only valid with use_cuda_graph=True. Every plan must stay within this capacity; each captured graph still requires compatible query/output buffers.

plan(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, page_size: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, sm_scale: float | None = None, window_left: int = -1, logits_soft_cap: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, seq_lens: Tensor | None = None, seq_lens_q: Tensor | None = None, block_tables: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False, uniform_q_len: int | None = None, use_token_head_sf: bool = False) → None¶

Plan batch prefill/append attention on Paged KV-Cache for given problem specification.

Parameters:

qo_indptr (torch.Tensor) – The indptr of the query/output tensor, shape: [batch_size + 1]. Token units for every backend, including cudnn (it previously took element-unit offsets; that contract has been normalized).
paged_kv_indptr (torch.Tensor) – The indptr of the paged kv-cache, shape: [batch_size + 1].
paged_kv_indices (torch.Tensor) – The page indices of the paged kv-cache, shape: [paged_kv_indptr[-1]].
paged_kv_last_page_len (torch.Tensor) – The number of entries in the last page of each request in the paged kv-cache, shape: [batch_size].
num_qo_heads (int) – The number of query/output heads.
num_kv_heads (int) – The number of key/value heads.
head_dim_qk (int) – The dimension of the query/key heads.
page_size (int) – The size of each page in the paged kv-cache.
head_dim_vo (Optional[int]) – The dimension of the value/output heads, if not provided, will be set to head_dim_qk.
custom_mask (Optional[torch.Tensor]) –
The flattened boolean mask tensor, shape: (sum(q_len[i] * k_len[i] for i in range(batch_size)). The elements in the mask tensor should be either True or False, where False means the corresponding element in the attention matrix will be masked out.

Please refer to the mask layout for more details about flattened layout of mask tensor.

When custom_mask is provided, and packed_custom_mask is not, the function will pack the custom mask tensor into a 1D packed mask tensor, which introduces additional overhead.
packed_custom_mask (Optional[torch.Tensor]) – The 1D packed uint8 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) – 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).
window_left (int) – The left (inclusive) window size for the attention window, when set to -1, the window size will be set to the full length of the sequence. Defaults to -1. For the trtllm-gen backend, an unsupported non-causal sliding window is rejected during plan() (rather than being deferred until run()).
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 (Union[str, torch.dtype]) – The data type of the query tensor, defaults torch.float16.
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 (Optional[Union[str, torch.dtype]]) – The data type of the output tensor. If None, will be set to q_data_type. For FP8 inputs, this should typically be set to torch.float16 or torch.bfloat16.
non_blocking (bool) – Whether to copy the input tensors to the device asynchronously, defaults to True.
prefix_len_ptr (Optional[torch.Tensor]) – prefix length. A uint32 1D tensor indicating the prefix length of each prompt. The tensor size is equal to the batch size.
token_pos_in_items_ptr (Optional[torch.Tensor]) – A uint16 1D tensor (it will be converted to uint16 in flashinfer) indicating the token position of each item and started from 0 (delimiter) for each item. E.g., if we have 3 items of length 3, 2, 4 respectively for this member. This vector will be looking like [0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0] with 4 delimiters indexed as 0. For batch size > 1, we will concat them as 1D with zero paddings to make sure each has the same length, the padding length is defined by token_pos_in_items_len - length of the raw token_pos_in_items_ptr for each prompt.
token_pos_in_items_len (int) – zero padding length for token_pos_in_items_ptr to better handle the bsz > 1 case. Still using the above 3,2,4 example. If we set token_pos_in_items_len to be 20, it will be [0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0] with 7 padded zeros. (note there’re 8 zeros in the end where the first one is the delimiter token 0 in the end of the prompt)
max_item_len_ptr (Optional[torch.Tensor]) – a uint16 vector contains the max token length of all items for each prompt
seq_lens (Optional[torch.Tensor]) – A uint32 1D tensor indicating the kv sequence length of each prompt. shape: [batch_size].
seq_lens_q (Optional[torch.Tensor]) – A uint32 1D tensor indicating the q sequence length of each prompt. shape: [batch_size]. For cuDNN, defaults to the differences in qo_indptr.
block_tables (Optional[torch.Tensor]) – A uint32 2D tensor indicating the block table of each prompt. shape: [batch_size, max_num_blocks_per_seq].
max_token_per_sequence (Optional[int],) – Maximum query length for cuDNN; inferred from qo_indptr when omitted.
max_sequence_kv (Optional[int],) – Maximum KV length for cuDNN; inferred from the paged metadata when omitted.
fixed_split_size (Optional[int],) – The fixed split size for FA2 split-kv prefill/decode in pages. Recommend setting to the average sequence length of your workload. When enabled, will lead to deterministic softmax score reduction in the merge_states kernel, and therefore batch-size invariant outputs. See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ Note that compatibility with CUDA graph is NOT guaranteed, as even when bs is fixed, kv seq len can change and lead to a varied number of launched CTAs.
disable_split_kv (bool,) – Whether to disable the split-kv for determinism in CUDA Graph, defaults to False.
uniform_q_len (Optional[int]) – Number of query rows of every request in the batch, when the caller can guarantee it for every plan replayed through this wrapper (speculative-decoding target verify: each request carries exactly draft_token_num query rows). Only the fa2 backend in CUDA-graph mode uses it: the scheduler then sizes the query tile for uniform_q_len * gqa_group_size packed rows instead of the ragged worst case total_num_rows - batch_size + 1. The hint is checked against qo_indptr on the host and dropped, with one warning, when the batch is not uniform. Defaults to None (no hint).
use_token_head_sf (bool) – Whether the KV cache uses the per-(token, head) FP8 scale layout: the KV cache is an FP8 (e4m3/e5m2) tensor with the standard layout (last dim head_dim), and a separate float32 scale tensor (mirroring the KV layout minus head_dim) is passed via kv_cache_sf at run() time. When enabled, the query must be fp16/bf16 and the backend is forced to fa2. Defaults to False.

Note

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.

The plan() method cannot be used in Cuda Graph or in torch.compile.

prewarm_paged_kv_stride_variant(variant: str = 'independent') → None¶

Load a lazy paged-KV-stride variant after plan().

Call this method before CUDA graph capture when a planned standard FA2 wrapper will receive K and V tensors with different data strides.

Parameters:: variant (str) – The paged-KV-stride variant to prewarm. The only supported value is "independent" (the default), for K and V with different data strides.

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 batch prefill/append attention between query and paged kv-cache.

Parameters:

q (torch.Tensor) – The query tensor, shape: [qo_indptr[-1], num_qo_heads, head_dim_qk]
paged_kv_cache (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) –
The paged KV-Cache stored as a tuple of tensors or a single tensor:
- a tuple (k_cache, v_cache) of 4-D tensors, each with shape: [max_num_pages, page_size, num_kv_heads, head_dim] if kv_layout is NHD, and [max_num_pages, num_kv_heads, page_size, head_dim] if kv_layout is HND.
- a single 5-D tensor with shape: [max_num_pages, 2, page_size, num_kv_heads, head_dim] if kv_layout is NHD, and [max_num_pages, 2, num_kv_heads, page_size, head_dim] if kv_layout is HND. Where paged_kv_cache[:, 0] is the key-cache and paged_kv_cache[:, 1] is the value-cache.
In these shapes, head_dim is head_dim_qk for the key cache and head_dim_vo for the value cache. For NVFP4 torch.uint8 caches, the stored last dimensions are head_dim_qk // 2 and head_dim_vo // 2, respectively. A single 5-D tensor requires equal key and value storage widths.
*args – Additional arguments for custom kernels.
q_scale (Optional[Union[float, torch.Tensor]]) – The calibration scale of query for fp8 input, if not provided, will be set to 1.0.
k_scale (Optional[Union[float, torch.Tensor]]) – The calibration scale of key for fp8 or nvfp4 input, if not provided, will be set to 1.0.
v_scale (Optional[Union[float, torch.Tensor]]) – The calibration scale of value for fp8 or nvfp4 input, 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 logsumexp of attention output
lse_base (str) – Base of the returned lse when return_lse is True: "log2" (default, FlashInfer’s convention, max + log2(sum(exp2(...)))) or "ln" (natural log, the torch.logsumexp convention that merge kernels written with expf consume). "ln" lets the cuDNN backend hand out its native stats without a conversion kernel; the other backends rescale.
lse_layout (str) – Layout of the returned lse: "NH" (default) is [qo_indptr[-1], num_qo_heads], "HN" is the contiguous transpose [num_qo_heads, qo_indptr[-1]]. Prepared cuDNN graphs write HN directly when supported; other paths transpose and rescale as needed. A caller-provided lse must have the requested layout.
enable_pdl (bool) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/programmatic-dependent-launch.html Only effective on backends and devices that support PDL.
window_left (Optional[int]) – Per-call override for the left (inclusive) sliding-window size. When None, the value supplied to plan() is used. Pass -1 to disable the sliding window for this call.
sinks (Optional[torch.Tensor]) – Per-head attention-sink logits. When provided, the kernel applies the attention-with-sink variant: an additional virtual token whose logit is sinks[head_idx] is appended to each row of the softmax denominator. Shape: [num_qo_heads], dtype float32.
kv_cache_sf (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]) –
Per-block scale factors for NVFP4 KV cache. Accepts the same formats as paged_kv_cache:
- a tuple (k_scales, v_scales) of 4-D tensors, each with shape: [num_pages, page_size, num_kv_heads, head_dim // 16] if kv_layout is NHD, and [num_pages, num_kv_heads, page_size, head_dim // 16] if kv_layout is HND.
- a single 5-D tensor with shape: [num_pages, 2, page_size, num_kv_heads, head_dim // 16] if kv_layout is NHD, and [num_pages, 2, num_kv_heads, page_size, head_dim // 16] if kv_layout is HND, where dim 1 holds k (index 0) and v (index 1) scales.
For the tuple form, head_dim is head_dim_qk in k_scales and head_dim_vo in v_scales. A single 5-D tensor requires equal key and value scale widths.

Both tensors have dtype torch.float8_e4m3fn. k_scales uses a linear (row-major) layout, while v_scales must use TRT-LLM’s 4-token interleaved layout within each [page_size, head_dim_vo // 16] tile if backend is trtllm-gen. Use flashinfer.fp4_quantization.nvfp4_quantize_paged_kv_cache() to produce correctly formatted scale factors.

For the trtllm-gen backend with NHD layout, scale tensors are transposed to HND internally (incurring a copy). Use HND for better performance. Currently, NVFP4 KV supports fa2 and trtllm-gen backend.

With use_token_head_sf=True (plan time), kv_cache_sf must be a (k_sf, v_sf) tuple of float32 tensors mirroring the KV layout minus head_dim: [num_pages, page_size, num_kv_heads] if kv_layout is NHD, and [num_pages, num_kv_heads, page_size] if HND. The tensors may be non-contiguous views (e.g. sliced out of an inline slot layout).
use_fp16_softmax (Optional[bool]) – trtllm-gen backend only. Select the …Fp16Softmax… cubin variant (FP16 softmax accumulator). Currently only shipped for BF16 Q/KV/O context kernels. Ignored by other backends.
uses_spcompress (Optional[bool]) – trtllm-gen backend only. Select the …Spcomp… cubin variant (sparse compression). Currently only shipped for FP8 Q context kernels. Ignored by other backends.
skip_softmax_threshold_scale_factor (Optional[float]) – Threshold scale factor for skipping softmax operations. Providing a value enables skip-softmax sparsity as described in https://arxiv.org/abs/2512.12087. Defaults to None (standard attention). Higher values yield faster kernels at the cost of accuracy; the effective threshold equals the supplied factor divided by the context length.

Returns:

If return_lse is False, the attention output, shape: [qo_indptr[-1], num_qo_heads, head_dim_vo]. If return_lse is True, a tuple of two tensors:

The attention output, shape: [qo_indptr[-1], num_qo_heads, head_dim_vo].
The logsumexp of attention output, shape: [qo_indptr[-1], num_qo_heads] (lse_layout="NH", default) or [num_qo_heads, qo_indptr[-1]] ("HN"), in the base selected by lse_base.

Return type:

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

workspace_size(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, page_size: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, sm_scale: float | None = None, window_left: int = -1, logits_soft_cap: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, seq_lens: Tensor | None = None, seq_lens_q: Tensor | None = None, block_tables: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False, uniform_q_len: int | None = None) → Tuple[int, int]¶

Return the caller-owned workspace size required by plan().

The returned tuple is (float_workspace_size, int_workspace_size) in bytes. The inputs follow plan(); host-side planning tensors such as qo_indptr and paged_kv_indptr are copied to CPU in the same way as plan(). The wrapper’s float workspace buffer is only used as the device/stream selector for the underlying module query. This method does not allocate buffers and does not mutate cached plan state.

Parameters:

qo_indptr (torch.Tensor) – The indptr of the query/output tensor, shape: [batch_size + 1].
paged_kv_indptr (torch.Tensor) – The indptr of the paged kv-cache, shape: [batch_size + 1].
paged_kv_indices (torch.Tensor) – The page indices of the paged kv-cache, shape: [paged_kv_indptr[-1]].
paged_kv_last_page_len (torch.Tensor) – The number of entries in the last page of each request in the paged kv-cache, shape: [batch_size].
num_qo_heads (int) – The number of query/output heads.
num_kv_heads (int) – The number of key/value heads.
head_dim_qk (int) – The dimension of the query/key heads.
page_size (int) – The size of each page in the paged kv-cache.
head_dim_vo (Optional[int]) – The dimension of the value/output heads. If not provided, defaults to head_dim_qk.
custom_mask (Optional[torch.Tensor]) – The flattened boolean mask tensor; when provided the packed mask is generated internally. See plan() for the full description.
packed_custom_mask (Optional[torch.Tensor]) – The 1D packed uint8 mask tensor. If provided, custom_mask is ignored.
causal (bool) – Whether to apply a causal mask. Defaults to False.
pos_encoding_mode (str) – The position encoding applied inside attention kernels, could be NONE/ROPE_LLAMA (LLAMA style rotary embedding) /ALIBI. Defaults to NONE.
use_fp16_qk_reduction (bool) – Whether to use fp16 for qk reduction. Defaults to False.
sm_scale (Optional[float]) – Softmax scale. If None, defaults to 1.0 / sqrt(head_dim_qk).
window_left (int) – The left (inclusive) window size for the attention window, when set to -1, the window size will be set to the full length of the sequence. Defaults to -1.
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.
rope_scale (Optional[float]) – Scale factor applied during RoPE interpolation. Defaults to 1.0 when None.
rope_theta (Optional[float]) – Base value for the RoPE frequencies. Defaults to 1e4 when None.
q_data_type (Union[str, torch.dtype]) – The data type of the query tensor. Defaults to torch.float16.
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 (Optional[Union[str, torch.dtype]]) – The data type of the output tensor. If None, will be set to q_data_type.
prefix_len_ptr (Optional[torch.Tensor]) – A uint32 1D tensor indicating the prefix length of each prompt, shape: [batch_size].
token_pos_in_items_ptr (Optional[torch.Tensor]) – A uint16 1D tensor indicating the token position of each item within its item.
token_pos_in_items_len (int) – Zero-padding length for token_pos_in_items_ptr. Defaults to 0.
max_item_len_ptr (Optional[torch.Tensor]) – A uint16 vector containing the max token length of all items for each prompt.
seq_lens (Optional[torch.Tensor]) – A uint32 1D tensor indicating the kv sequence length of each prompt, shape: [batch_size].
seq_lens_q (Optional[torch.Tensor]) – A uint32 1D tensor indicating the q sequence length of each prompt, shape: [batch_size].
block_tables (Optional[torch.Tensor]) – A uint32 2D tensor indicating the block table of each prompt, shape: [batch_size, max_num_blocks_per_seq].
max_token_per_sequence (Optional[int]) – Required for cuDNN backend. The scalar max token length of each sequence.
max_sequence_kv (Optional[int]) – Maximum number of KV tokens per sequence for cuDNN backend.
fixed_split_size (Optional[int]) – The fixed split size for split-kv prefill, in pages.
disable_split_kv (bool) – Whether to disable the split-kv. Defaults to False.

Returns:

(float_workspace_size, int_workspace_size) in bytes.

Return type:

Tuple[int, int]

Example

>>> float_bytes, int_bytes = wrapper.workspace_size(...)
>>> wrapper.reset_workspace_buffer(
...     torch.empty(float_bytes, dtype=torch.uint8, device="cuda"),
...     torch.empty(int_bytes, dtype=torch.uint8, device="cuda"),
... )
>>> wrapper.plan(...)

class flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, kv_indptr_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None, variant_owns_mask: bool = False)¶

Wrapper class for prefill/append attention with a ragged (tensor) KV cache for a batch of requests.

Check our tutorial for ragged kv-cache layout.

Example

>>> import torch
>>> import flashinfer
>>> num_layers = 32
>>> num_qo_heads = 64
>>> num_kv_heads = 16
>>> head_dim = 128
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> prefill_wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
...     workspace_buffer, "NHD"
... )
>>> batch_size = 7
>>> nnz_kv = 100
>>> nnz_qo = 100
>>> qo_indptr = torch.tensor(
...     [0, 33, 44, 55, 66, 77, 88, nnz_qo], dtype=torch.int32, device="cuda:0"
... )
>>> kv_indptr = qo_indptr.clone()
>>> q_at_layer = torch.randn(num_layers, nnz_qo, num_qo_heads, head_dim).half().to("cuda:0")
>>> k_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")
>>> v_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")
>>> # create auxiliary data structures for batch prefill attention
>>> prefill_wrapper.plan(
...     qo_indptr,
...     kv_indptr,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
...     causal=True,
... )
>>> outputs = []
>>> for i in range(num_layers):
...     q = q_at_layer[i]
...     k = k_at_layer[i]
...     v = v_at_layer[i]
...     # compute batch prefill attention, reuse auxiliary data structures
...     o = prefill_wrapper.run(q, k, v)
...     outputs.append(o)
...
>>> outputs[0].shape
torch.Size([100, 64, 128])
>>>
>>> # below is another example of creating custom mask for batch prefill attention
>>> mask_arr = []
>>> qo_len = (qo_indptr[1:] - qo_indptr[:-1]).cpu().tolist()
>>> kv_len = (kv_indptr[1:] - kv_indptr[:-1]).cpu().tolist()
>>> for i in range(batch_size):
...     mask_i = torch.tril(
...         torch.full((qo_len[i], kv_len[i]), True, device="cuda:0"),
...         diagonal=(kv_len[i] - qo_len[i]),
...     )
...     mask_arr.append(mask_i.flatten())
...
>>> mask = torch.cat(mask_arr, dim=0)
>>> prefill_wrapper.plan(
...     qo_indptr,
...     kv_indptr,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
...     custom_mask=mask
... )
>>> outputs_custom_mask = []
>>> for i in range(num_layers):
...     q = q_at_layer[i]
...     k = k_at_layer[i]
...     v = v_at_layer[i]
...     # compute batch prefill attention, reuse auxiliary data structures
...     o_custom = prefill_wrapper.run(q, k, v)
...     assert torch.allclose(o_custom, outputs[i], rtol=1e-3, atol=1e-3)
...
>>> outputs_custom_mask[0].shape
torch.Size([100, 64, 128])

Note

__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, kv_indptr_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None, variant_owns_mask: bool = False) → None¶

Constructor of BatchPrefillWithRaggedKVCacheWrapper.

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. The buffer must be 16-byte aligned; tensors created by torch.empty satisfy this on supported devices.
kv_layout (str) – The layout of the input k/v tensors, could be either NHD or HND.
use_cuda_graph (bool) – Whether to enable CUDA graph capture for the prefill kernels, if enabled, the auxiliary data structures will be stored as the provided buffers.
qo_indptr_buf (Optional[torch.Tensor]) – The user reserved GPU buffer to store the qo_indptr array, the size of the buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True.
kv_indptr_buf (Optional[torch.Tensor]) – The user reserved GPU buffer to store the kv_indptr array, the size of the buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True.
custom_mask_buf (Optional[torch.Tensor]) – The user reserved GPU buffer to store the custom mask tensor, should be large enough to store the maximum possible size of the packed custom mask tensor during the lifetime of the wrapper. This argument is only effective when use_cuda_graph is True and custom mask will be used in attention computation.
mask_indptr_buf (Optional[torch.Tensor]) – The user reserved GPU buffer to store the mask_indptr array, the size of the buffer should be [batch_size]. This argument is only effective when use_cuda_graph is True and custom mask will be used in attention computation.
backend (str) – The implementation backend, could be auto/fa2/fa3/cudnn/cutlass or cute-dsl/cute-dsl-prims/cutile. Defaults to auto. If set to auto, the wrapper will automatically choose the backend based on the device architecture and kernel availability: fa3 on Hopper where it applies, fmha_v2 for large-sequence MHA on SM120, and on SM100a cudnn (when the installed cuDNN can consume token-unit indptrs directly, see _cudnn_supports_direct_seqlens()) or else cutlass on SM100a/SM110a, for the head dims those kernels serve. Sliding windows and attention sinks are supported by the eligible cudnn path, but exclude cutlass. RoPE, logits soft caps, custom masks and multi-item-scoring pointers exclude both upgrades, so auto stays on fa2 for them. Otherwise fa2. The cudnn backend takes qo_indptr/kv_indptr in the same token units as every other backend, whether selected explicitly or via auto. The cute-dsl backend uses the CuTe DSL attention kernel for Blackwell (SM100+). cute-dsl-prims is an explicit SM120-only packed FP8 prefill backend. The cutile backend uses the pure cuda.tile Python prefill kernel (Blackwell, opt-in); it requires qo_indptr == kv_indptr (equal-length prefill).
jit_args (Optional[List[Any]]) – If provided, the wrapper will use the provided arguments to create the JIT module, otherwise, the wrapper will use default attention implementation.
jit_kwargs (Optional[Dict[str, Any]]) – The keyword arguments to create the JIT module, defaults to None.
variant_owns_mask (bool) – If True, the attention variant supplied through jit_args computes the complete attention mask in its LogitsMask hook, and MaskMode.CUSTOM is selected without a mask tensor: the kernel evaluates LogitsMask on every KV tile instead of only on the causal/window boundary tiles, and no custom_mask/packed_custom_mask needs to be passed to plan(). Requires jit_args, because the default attention variant dereferences the custom mask buffer under MaskMode.CUSTOM. Only supported with backend="fa2" (the SM90 batch prefill kernels reject MaskMode.CUSTOM), and incompatible with prefix_len_ptr (multi-item scoring), which selects a different mask mode. Defaults to False.

plan(qo_indptr: Tensor, kv_indptr: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, window_left: int = -1, logits_soft_cap: float | None = None, sm_scale: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False, seq_lens: Tensor | None = None, seq_lens_q: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, v_indptr: Tensor | None = None, o_indptr: Tensor | None = None, qo_indptr_cpu: Tensor | None = None, kv_indptr_cpu: Tensor | None = None, uniform_q_len: int | None = None, use_token_head_sf: bool = False) → None¶

Plan batch prefill/append attention on Ragged KV-Cache for given problem specification.

Parameters:

qo_indptr (torch.Tensor) – The indptr of the query/output tensor, shape: [batch_size + 1]. Token units for every backend, including cudnn (it previously took element-unit offsets; that contract has been normalized). The cudnn backend also requires kv_layout="NHD".
kv_indptr (torch.Tensor) – The indptr of the key/value tensor, shape: [batch_size + 1].
qo_indptr_cpu (Optional[torch.Tensor]) – CPU mirrors of the corresponding indptrs, with the same shape, integer dtype, and values. When supplied, planning reads these mirrors without copying device indptrs back to the CPU. Device bindings still use qo_indptr and kv_indptr. The caller must keep each mirror consistent with its device tensor at every plan; values are not compared because that would require synchronization.
kv_indptr_cpu (Optional[torch.Tensor]) – CPU mirrors of the corresponding indptrs, with the same shape, integer dtype, and values. When supplied, planning reads these mirrors without copying device indptrs back to the CPU. Device bindings still use qo_indptr and kv_indptr. The caller must keep each mirror consistent with its device tensor at every plan; values are not compared because that would require synchronization.
num_qo_heads (int) – The number of query/output heads.
num_kv_heads (int) – The number of key/value heads.
head_dim_qk (int) – The dimension of the heads on query/key tensor.
head_dim_vo (Optional[int]) – The dimension of the heads on value/output tensor. If not provided, will be set to head_dim_qk.
custom_mask (Optional[torch.Tensor]) –
The flattened boolean mask tensor, shape: (sum(q_len[i] * k_len[i] for i in range(batch_size)). The elements in the mask tensor should be either True or False, where False means the corresponding element in the attention matrix will be masked out.

Please refer to the mask layout for more details about flattened layout of mask tensor.

When custom_mask is provided, and packed_custom_mask is not, the function will pack the custom mask tensor into a 1D packed mask tensor, which introduces additional overhead.
packed_custom_mask (Optional[torch.Tensor]) –
The 1D packed uint8 mask tensor, if provided, the custom_mask will be ignored. The packed mask tensor is generated by flashinfer.quantization.packbits().

If provided, the custom mask will be added to the attention matrix before softmax and after scaling. The mask tensor should be in the same device as the input tensors.
causal (bool) – Whether to apply causal mask to the attention matrix. This argument is ignored if mask is provided in plan().
pos_encoding_mode (str) – 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).
window_left (int) – The left (inclusive) window size for the attention window, when set to -1, the window size will be set to the full length of the sequence. Defaults to -1.
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_qk).
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 (Union[str, torch.dtype]) – The data type of the query tensor, defaults to torch.float16.
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 (Optional[Union[str, torch.dtype]]) – The data type of the output tensor. If None, will be set to q_data_type. For FP8 inputs, this should typically be set to torch.float16 or torch.bfloat16.
non_blocking (bool) – Whether to copy the input tensors to the device asynchronously, defaults to True.
prefix_len_ptr (Optional[torch.Tensor]) – prefix length. A uint32 1D tensor indicating the prefix length of each prompt. The tensor size is equal to the batch size.
token_pos_in_items_ptr (Optional[torch.Tensor]) – A uint16 1D tensor (it will be converted to uint16 in flashinfer) indicating the token position of each item and started from 0 (delimiter) for each item. E.g., if we have 3 items of length 3, 2, 4 respectively for this member. This vector will be looking like [0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0] with 4 delimiters indexed as 0. For batch size > 1, we will concat them as 1D with zero paddings to make sure each has the same length, the padding length is defined by token_pos_in_items_len - length of the raw token_pos_in_items_ptr for each prompt.
token_pos_in_items_len (int) – zero padding length for token_pos_in_items_ptr to better handle the bsz > 1 case. Still using the above 3,2,4 example. If we set token_pos_in_items_len to be 20, it will be [0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0] with 7 padded zeros. (note there’re 8 zeros in the end where the first one is the delimiter token 0 in the end of the prompt)
max_item_len_ptr (Optional[torch.Tensor]) – a uint16 vector contains the max token length of all items for each prompt
fixed_split_size (Optional[int],) – The fixed split size for split-kv FA2 prefill/decode, in pages. Recommend setting to the average sequence length of your workload. When enabled, will lead to deterministic softmax score reduction in the merge_states kernel, and therefore batch-size invariant outputs. See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ Note that compatibility with CUDA graph is NOT guaranteed, as even when bs is fixed, kv seq len can change and lead to a varied number of launched CTAs.
disable_split_kv (bool,) – Whether to disable the split-kv for determinism in CUDA Graph, defaults to False.
seq_lens (Optional[torch.Tensor]) – Only used by the cudnn backend. A 1D int32 tensor with the kv sequence length of each prompt, shape: [batch_size]. Optional: cuDNN derives it from kv_indptr.
seq_lens_q (Optional[torch.Tensor]) – Only used by the cudnn backend. A 1D int32 tensor with the q sequence length of each prompt, shape: [batch_size]. Optional: cuDNN derives it from qo_indptr.
max_token_per_sequence (Optional[int],) – Only used by the cudnn backend. The max q length of any sequence; defaults to max(qo_indptr[1:] - qo_indptr[:-1]).
max_sequence_kv (Optional[int],) – Only used by the cudnn backend. The max kv length of any sequence; defaults to max(kv_indptr[1:] - kv_indptr[:-1]).
v_indptr (Optional[torch.Tensor]) – Only used by the cudnn backend. Token-unit indptr of the value tensor; defaults to kv_indptr.
o_indptr (Optional[torch.Tensor]) – Only used by the cudnn backend. Token-unit indptr of the output tensor; defaults to qo_indptr.
uniform_q_len (Optional[int]) – Number of query rows of every request in the batch, when the caller can guarantee it for every plan replayed through this wrapper. Only the fa2 backend in CUDA-graph mode uses it (query tile sized for uniform_q_len * gqa_group_size packed rows instead of the ragged worst case); see BatchPrefillWithPagedKVCacheWrapper.plan(). Checked against qo_indptr on the host and dropped when the batch is not uniform. Defaults to None.
use_token_head_sf (bool) – Whether the KV cache uses the per-(token, head) FP8 scale layout: the KV cache is an FP8 (e4m3/e5m2) tensor with the standard layout (last dim head_dim), and a separate float32 scale tensor (mirroring the KV layout minus head_dim) is passed via kv_cache_sf at run() time. When enabled, the query must be fp16/bf16 and the backend is forced to fa2. Defaults to False.

Note

The plan() method should be called before any run() or run_return_lse() calls, auxiliary data structures will be created during this plan 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.

The plan() method cannot be used in Cuda Graph or in torch.compile.

reset_workspace_buffer(float_workspace_buffer: Tensor, int_workspace_buffer) → 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, *args, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, lse_base: str = 'log2', lse_layout: str = 'NH', enable_pdl: bool | None = None, kv_cache_sf: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None) → Tensor¶

run(q: Tensor, k: Tensor, v: Tensor, *args, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[True] = True, lse_base: str = 'log2', lse_layout: str = 'NH', enable_pdl: bool | None = None, kv_cache_sf: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None) → Tuple[Tensor, Tensor]

Compute batch prefill/append attention between query and kv-cache stored as ragged tensor.

Parameters:

q (torch.Tensor) – The query tensor, shape: [qo_indptr[-1], num_qo_heads, head_dim_qk]
k (torch.Tensor) – The key tensor, shape: [kv_indptr[-1], num_kv_heads, head_dim_qk]
v (torch.Tensor) – The value tensor, shape: [kv_indptr[-1], num_kv_heads, head_dim_vo]
*args – Additional arguments for the custom kernel.
q_scale (Optional[float]) – The calibration scale of fp8 query, if not provided, will be set to 1.0.
k_scale (Optional[float]) – The calibration scale of fp8 or nvfp4 key, if not provided, will be set to 1.0.
v_scale (Optional[float]) – The calibration scale of fp8 or nvfp4 value, if not provided, will be set to 1.0.
o_scale (Optional[float]) – The calibration scale of output, 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 logsumexp of attention output
lse_base (str) – Base of the returned lse when return_lse is True: "log2" (default, FlashInfer’s convention, max + log2(sum(exp2(...)))) or "ln" (natural log, the torch.logsumexp convention that merge kernels written with expf consume). "ln" lets the cuDNN backend hand out its native stats without a conversion kernel; the other backends rescale.
lse_layout (str) – Layout of the returned lse: "NH" (default) is [qo_indptr[-1], num_qo_heads], "HN" is the contiguous transpose [num_qo_heads, qo_indptr[-1]]. Prepared cuDNN graphs write HN directly when supported; other paths transpose and rescale as needed. A caller-provided lse must have the requested layout.
enable_pdl (bool) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/programmatic-dependent-launch.html Only effective on backends and devices that support PDL.
kv_cache_sf (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]) –
Per-block scale factors for NVFP4 KV input. Accepts either a single packed scale tensor or a (k_scales, v_scales) tuple matching the structure expected by the chosen backend. When None (default), the kernel runs without NVFP4 KV scaling. See flashinfer.fp4_quantization.nvfp4_quantize() for layout details.

With use_token_head_sf=True (plan time), kv_cache_sf must be a (k_sf, v_sf) tuple of float32 tensors mirroring the KV layout minus head_dim: [total_tokens, num_kv_heads] if kv_layout is NHD, and [num_kv_heads, total_tokens] if HND. The scale strides must mirror the KV strides divided by head_dim (the kernel derives the scale strides from the KV strides, as the NVFP4 path does); inline-slot views are supported by the paged path only.

Returns:

If return_lse is False, the attention output, shape: [qo_indptr[-1], num_qo_heads, head_dim_vo]. If return_lse is True, a tuple of two tensors:

The attention output, shape: [qo_indptr[-1], num_qo_heads, head_dim_vo].
The logsumexp of attention output, shape: [qo_indptr[-1], num_qo_heads] (lse_layout="NH", default) or [num_qo_heads, qo_indptr[-1]] ("HN"), in the base selected by lse_base.

Return type:

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

Causal + Bidirectional Ranges Prefill¶

A batch-prefill wrapper whose fa2 attention variant owns the whole mask: causal, plus an inclusive per-query key span attended in both directions. The spans are handed to BatchPrefillWithCausalBidirectionalRangesWrapper.run() as a compact int32 [total_q, 2] tensor and no mask is materialized, so nothing scales with qo_len * kv_len. The JIT module is specialized in the constructor, and the inherited options the variant makes meaningless are rejected rather than ignored.

class flashinfer.attention.BatchPrefillWithCausalBidirectionalRangesWrapper(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, pos_encoding_mode: str = 'NONE', q_data_type: dtype = torch.bfloat16, kv_data_type: dtype = torch.bfloat16, o_data_type: dtype | None = None, head_dim_qk: int = 128, head_dim_vo: int = 128, v_sf_swizzled: bool = False)¶

Paged batch prefill with a causal mask plus per-query bidirectional ranges.

The mask this wrapper applies is

(causal AND causal_window) OR (in_range AND range_window)

where in_range is decided per query token from an inclusive [start, end] key span supplied by the caller. Tokens inside their own span see each other in both directions; everything else stays causal. This is the shape a prefix-style batch produces, but the wrapper knows nothing about what the spans mean.

The mask is never materialized. A custom fa2 attention variant owns the whole expression and evaluates it on every KV tile from the compact range tensor, so nothing here grows with qo_len * kv_len.

Ranges are passed to run() as a contiguous int32 tensor of shape [total_q, 2] on the query device, one row per scheduled query token in the same order as the query tensor. Row i is the inclusive [start, end] span of query i in absolute key positions within its request; (-1, -1) marks a query with no span, which then keeps the plain causal mask.

The wrapper is fa2 only, because fa2 is the only backend whose kernels evaluate a custom LogitsMask per KV tile. It is not related to, and cannot be combined with, prefix_len_ptr (multi-item scoring), which selects a different mask mode and owns the mask itself.

The JIT module is specialized in the constructor on the query, KV and output dtypes, both head dimensions and the positional-encoding mode. plan() and workspace_size() reject any value that differs from what was compiled, and reject the inherited options that would otherwise be accepted and silently ignored. v_sf_swizzled reads the V scale factors of a packed NVFP4 cache in the trtllm-gen per-page swizzle.

Example

>>> import torch, flashinfer
>>> workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> wrapper = flashinfer.BatchPrefillWithCausalBidirectionalRangesWrapper(
...     workspace, q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16,
...     head_dim_qk=128, head_dim_vo=128,
... )  

__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, pos_encoding_mode: str = 'NONE', q_data_type: dtype = torch.bfloat16, kv_data_type: dtype = torch.bfloat16, o_data_type: dtype | None = None, head_dim_qk: int = 128, head_dim_vo: int = 128, v_sf_swizzled: bool = False) → None¶

Constructor of BatchPrefillWithPagedKVCacheWrapper.

Parameters:

float_workspace_buffer (torch.Tensor) – The user reserved workspace buffer used to store intermediate attention results in 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. The buffer must be 16-byte aligned; tensors created by torch.empty satisfy this on supported devices.
kv_layout (str) – The layout of the input k/v tensors, could be either NHD or HND.
use_cuda_graph (bool) – Whether to enable CUDA graph capture for the prefill kernels, if enabled, the auxiliary data structures will be stored in provided buffers. The batch_size cannot change during the lifecycle of this wrapper when CUDAGraph is enabled. With backend="auto", the first successful FA2/FA3 plan fixes the backend for subsequent plans so captured scheduling metadata stays compatible. Inputs requiring a different backend need a new wrapper.
qo_indptr_buf (Optional[torch.Tensor]) – The user reserved buffer to store the qo_indptr array, the size of the buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True.
paged_kv_indptr_buf (Optional[torch.Tensor]) – The user reserved buffer to store the paged_kv_indptr array, the size of this buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True.
paged_kv_indices_buf (Optional[torch.Tensor]) – The user reserved buffer to store the paged_kv_indices array, should be large enough to store the maximum possible size of the paged_kv_indices array during the lifetime of the wrapper. This argument is only effective when use_cuda_graph is True.
paged_kv_last_page_len_buf (Optional[torch.Tensor]) – The user reserved buffer to store the paged_kv_last_page_len array, the size of the buffer should be [batch_size]. This argument is only effective when use_cuda_graph is True.
custom_mask_buf (Optional[torch.Tensor]) – The user reserved buffer to store the custom mask tensor, should be large enough to store the maximum possible size of the packed custom mask tensor during the lifetime of the wrapper. This argument is only effective when use_cuda_graph is set to True and the custom mask will be used in attention computation.
mask_indptr_buf (Optional[torch.Tensor]) – The user reserved buffer to store the mask_indptr array, the size of the buffer should be [batch_size + 1]. This argument is only effective when use_cuda_graph is True and the custom mask will be used in attention computation.
backend (str) – The implementation backend, could be auto/fa2/fa3/cudnn/trtllm-gen cute-dsl/cute-dsl-prims. Defaults to auto. If set to auto, the wrapper will automatically choose the backend based on the device architecture and kernel availability. In CUDA-graph mode, once FA2 has been planned, later plans keep FA2 so its captured scheduling metadata is refreshed. Use a separate wrapper to select cuDNN for a new graph after planning FA2. The cute-dsl backend uses the CuTe DSL attention kernel for Blackwell (SM100+). cute-dsl-prims backend is a SM120-only FP8 attention backend implemented using CUTLASS primitives.
jit_args (Optional[List[Any]]) – If provided, the wrapper will use the provided arguments to create the JIT module, otherwise, the wrapper will use default attention implementation.
jit_kwargs (Optional[Dict[str, Any]]) – The keyword arguments to create the JIT module, defaults to None.
variant_owns_mask (bool) – If True, the attention variant supplied through jit_args computes the complete attention mask in its LogitsMask hook, and MaskMode.CUSTOM is selected without a mask tensor: the kernel evaluates LogitsMask on every KV tile instead of only on the causal/window boundary tiles, and no custom_mask/packed_custom_mask needs to be passed to plan(). Requires jit_args, because the default attention variant dereferences the custom mask buffer under MaskMode.CUSTOM. Only supported with backend="fa2" (the SM90 batch prefill kernels reject MaskMode.CUSTOM), and incompatible with prefix_len_ptr (multi-item scoring), which selects a different mask mode. Defaults to False.
max_total_num_rows (Optional[int]) – Maximum total number of query tokens across requests in CUDA graph mode. Declaring this capacity lets a later plan() accept more query tokens than the first call while retaining the same metadata buffers. If omitted, the first plan() fixes the capacity. Only valid with use_cuda_graph=True. Every plan must stay within this capacity; each captured graph still requires compatible query/output buffers.

begin_forward(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, page_size: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, sm_scale: float | None = None, window_left: int = -1, logits_soft_cap: float | None = None, q_data_type: str | dtype | None = None, kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False) → None¶

Plan a batch whose mask this variant owns.

The signature mirrors the parent wrapper so that a caller can swap the two, but every argument that the variant makes meaningless is rejected instead of ignored, and the dtypes and head dimensions must match what the constructor compiled. Backend-specific planning arguments of the parent (seq_lens, block_tables, rope_scale, …) are absent here: this wrapper is fa2 only and applies no positional encoding.

forward(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], bidirectional_ranges: Tensor, causal_window_left: int = -1, range_window_left: int = -1, q_scale: float | None = None, k_scale: float | None = None, v_scale: float | None = None, enable_pdl: bool | None = None, kv_cache_sf: Tensor | Tuple[Tensor, Tensor] | None = None) → Tensor¶

Warning: deprecated, use run() instead.

The parent’s forward sets plan state from keyword arguments this wrapper refuses and has no parameter for the ranges, so it is replaced rather than inherited: reaching it would assign _causal, _window_left, _logits_soft_cap and _sm_scale and only then fail inside prepare_jit_additional_args.

forward_return_lse(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], bidirectional_ranges: Tensor, causal_window_left: int = -1, range_window_left: int = -1, q_scale: float | None = None, k_scale: float | None = None, v_scale: float | None = None, enable_pdl: bool | None = None, kv_cache_sf: Tensor | Tuple[Tensor, Tensor] | None = None) → Tuple[Tensor, Tensor]¶

Warning: deprecated, use run_return_lse() instead.

Replaced for the same reason as forward().

plan(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, page_size: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, sm_scale: float | None = None, window_left: int = -1, logits_soft_cap: float | None = None, q_data_type: str | dtype | None = None, kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False) → None¶

Plan a batch whose mask this variant owns.

run(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], bidirectional_ranges: Tensor, causal_window_left: int = -1, range_window_left: int = -1, q_scale: float | None = None, k_scale: float | None = None, v_scale: float | None = None, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, enable_pdl: bool | None = None, kv_cache_sf: Tensor | Tuple[Tensor, Tensor] | None = None) → Tensor¶

run(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], bidirectional_ranges: Tensor, causal_window_left: int = -1, range_window_left: int = -1, q_scale: float | None = None, k_scale: float | None = None, v_scale: float | None = None, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[True] = True, enable_pdl: bool | None = None, kv_cache_sf: Tensor | Tuple[Tensor, Tensor] | None = None) → Tuple[Tensor, Tensor]

Run causal + bidirectional-ranges attention.

Parameters:

q (torch.Tensor) – Query tensor, shape [total_q, num_qo_heads, head_dim_qk].
paged_kv_cache (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) – The paged KV cache, as accepted by the parent wrapper.
bidirectional_ranges (torch.Tensor) – Contiguous int32 tensor of shape [total_q, 2] on the query device. Row i is the inclusive [start, end] key span that query i attends bidirectionally, or (-1, -1) for none. A contiguous row slice of a larger [N, 2] buffer is accepted; a non-contiguous tensor is rejected rather than copied, so the kernel argument stays a stable pointer under CUDA graph capture.
causal_window_left (int) – Sliding window on the causal part, with the same meaning as window_left everywhere else in FlashInfer: a key is kept when q_abs - kv <= causal_window_left, so N keeps N + 1 keys and 0 keeps the diagonal alone. A negative value disables the window and leaves the causal part unbounded. Pass the value you would pass to any other prefill wrapper; no conversion is needed.
range_window_left (int) –
Sliding window on the bidirectional part, as a count of keys rather than as a left bound: a key inside a span is kept when q_abs - kv < range_window_left, so N keeps the N nearest backward distances, 0..N-1. 0 and any negative value leave the spans unclamped, which is what a span normally wants. Only the backward edge is clamped; a span’s forward side is untouched.

This is deliberately not causal_window_left’s contract. That one carries FlashInfer’s window_left, where N keeps 0..N and 0 keeps the diagonal alone. The same number therefore means one fewer key here, and 0 means the opposite of what it means there.
q_scale (Optional[float]) – Scalar calibration scales folded into sm_scale. Per-head scale tensors are rejected: this variant declares a single double sm_scale, so a tensor could not be folded into it.
k_scale (Optional[float]) – Scalar calibration scales folded into sm_scale. Per-head scale tensors are rejected: this variant declares a single double sm_scale, so a tensor could not be folded into it.
v_scale (Optional[float]) – Scalar calibration scale of the value cache. It never reaches the variant: the parent rescales the output once the kernel has returned. Only a scalar is exposed here, so the public contract of this wrapper stays a single output rescaling; per-head tensors are rejected.
out (Optional[torch.Tensor]) – Output tensor; allocated internally when omitted. Its dtype must be the o_data_type this wrapper was constructed with.
lse (Optional[torch.Tensor]) – Log-sum-exp tensor, allocated internally when omitted.
return_lse (bool) – Whether to return the log-sum-exp alongside the output.
enable_pdl (Optional[bool]) – Programmatic dependent launch, as in the parent wrapper.
kv_cache_sf (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]) – Scale factors for a packed NVFP4 KV cache. Required when the wrapper was constructed with a packed fp4 kv_data_type, and meaningless otherwise.

Notes

window_left, sinks and the parent’s remaining keyword arguments are deliberately absent. Each of them either contradicts the mask this variant owns or falls outside the contract this specialized wrapper offers, so it is not exposed rather than accepted and then ignored.

workspace_size(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, page_size: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, sm_scale: float | None = None, window_left: int = -1, logits_soft_cap: float | None = None, q_data_type: str | dtype | None = None, kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False) → Tuple[int, int]¶

Caller-owned workspace size for a plan() with these arguments.

Validates exactly what plan() validates, so a caller cannot size a workspace for a configuration that plan() would then reject.

Unified BatchAttention¶

The BatchAttention class provides a holistic attention wrapper that automatically dispatches between paged-prefill and paged-decode based on per-request sequence lengths. It is the recommended entry point for serving stacks that batch mixed prefill/decode requests in a single kernel launch.

class flashinfer.attention.BatchAttention(kv_layout: str = 'NHD', device: str = 'cuda')¶

Holistic batched attention wrapper that fuses paged-prefill and paged-decode requests into a single kernel launch.

BatchAttention dispatches between prefill-style and decode-style execution per request based on the qo_indptr / kv_indptr ranges supplied to plan(), so a serving stack can submit a mixed batch (e.g. some prompts in prefill, others in decode) without splitting it into two separate wrappers. Workspace buffers are owned by the instance and reused across plan() / run() calls.

Parameters:

kv_layout (str) – Layout of the paged KV-cache tensors, either "NHD" (token-major) or "HND" (head-major). Defaults to "NHD".
device (str) – CUDA device that owns the internal workspace buffers, e.g. "cuda" or "cuda:0". Defaults to "cuda".

__init__(kv_layout: str = 'NHD', device: str = 'cuda')¶

Allocate workspace buffers and bind the wrapper to a CUDA device.

See BatchAttention for the meaning of each parameter.

plan(qo_indptr: Tensor, kv_indptr: Tensor, kv_indices: Tensor, kv_len_arr: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, head_dim_vo: int, page_size: int, causal: bool = False, sm_scale: float | None = None, logits_soft_cap: float | None = None, q_data_type: dtype = torch.bfloat16, kv_data_type: dtype = torch.bfloat16, use_profiler: bool = False) → None¶

Plan the holistic attention kernel for a specific batch shape.

Should be called before any run() call. The plan is cached on the instance and reused across subsequent run() invocations with the same layout.

Parameters:

qo_indptr (torch.Tensor) – CSR-style query offsets, shape [batch_size + 1], dtype int32.
kv_indptr (torch.Tensor) – CSR-style page offsets into kv_indices, shape [batch_size + 1], dtype int32.
kv_indices (torch.Tensor) – Page indices into the paged KV-cache, shape [kv_indptr[-1]], dtype int32.
kv_len_arr (torch.Tensor) – Per-request KV-cache lengths in tokens, shape [batch_size], dtype int32.
num_qo_heads (int) – Number of query / output heads.
num_kv_heads (int) – Number of key / value heads. Must divide num_qo_heads.
head_dim_qk (int) – Per-head dimension of the query / key tensors.
head_dim_vo (int) – Per-head dimension of the value / output tensors.
page_size (int) – Page size of the paged KV-cache.
causal (bool) – Whether to apply a causal mask. Defaults to False.
sm_scale (Optional[float]) – Softmax scale. If None, defaults to 1/sqrt(head_dim_qk).
logits_soft_cap (Optional[float]) – Logits soft-cap value. None or 0 disables capping.
q_data_type (torch.dtype) – Dtype of the query tensor. Defaults to torch.bfloat16.
kv_data_type (torch.dtype) – Dtype of the key / value tensors. Defaults to torch.bfloat16.
use_profiler (bool) – Whether to compile the profiler-enabled variant of the kernel. Defaults to False.

prewarm_paged_kv_stride_variant(variant: str = 'independent') → None¶

Load the independent data-stride variant after plan().

Call this before CUDA graph capture when K and V will have different page, token, or head strides. The only supported variant is "independent". The equal-stride primary is loaded by plan().

The independent module is compiled on demand and is not included in the default precompiled cache. With JIT disabled, a compatible cache provider must supply it, or it must already be loaded in this process.

Parameters:: variant (str, optional) – The variant to prewarm. The only supported value is "independent", which is also the default.

Run the planned holistic attention kernel.

Parameters:

q (torch.Tensor) – Query tensor, shape [total_qo_tokens, num_qo_heads, head_dim_qk].
kv_cache (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) – Either a single packed paged KV-cache tensor (when K and V share storage) or a (k_cache, v_cache) pair. Layout must match the kv_layout passed to __init__(). Unequal K/V data strides load a separate module on first use. Call prewarm_paged_kv_stride_variant() before CUDA graph capture when these strides differ.
out (Optional[torch.Tensor]) – Optional output buffer. If None, a new tensor is allocated with the same shape as q.
lse (Optional[torch.Tensor]) – Optional log-sum-exp buffer, shape [total_qo_tokens, num_qo_heads], dtype float32. Allocated if None.
k_scale (Optional[torch.Tensor]) – FP8 dequantization scale for k. Pre-multiplied into sm_scale.
v_scale (Optional[torch.Tensor]) – FP8 dequantization scale for v. Applied to the output.
logits_soft_cap (float) – Logits soft-cap value. Must be consistent with the logits_soft_cap passed to plan() (a non-zero value here requires a non-zero plan-time value too).
profiler_buffer (Optional[torch.Tensor]) – Profiler buffer. Required if the wrapper was planned with use_profiler=True.
kv_cache_sf (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]) – Optional scale tensors for NVFP4 KV-cache (one tensor or a (k_sf, v_sf) pair, mirroring the structure of kv_cache).

Returns:

(out, lse) — the attention output and its log-sum-exp.

Return type:

Tuple[torch.Tensor, torch.Tensor]

class flashinfer.attention.BatchAttentionWithAttentionSinkWrapper(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, q_data_type: dtype = torch.bfloat16, kv_data_type: dtype = torch.bfloat16, head_dim_qk: int = 128, head_dim_vo: int = 128, window_left: int = -1)¶

Wrapper for prefill and decode attention with paged KV-cache that adds support for attention sinks. This class extends BatchPrefillWithPagedKVCacheWrapper, providing a convenient interface for using attention sinks during prefill or decode attention.

__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, q_data_type: dtype = torch.bfloat16, kv_data_type: dtype = torch.bfloat16, head_dim_qk: int = 128, head_dim_vo: int = 128, window_left: int = -1) → None¶

Create an attention-sink wrapper with a sink-specific JIT variant.

Parameters:

float_workspace_buffer (torch.Tensor) – Workspace buffer for split-K intermediate results.
kv_layout (str) – KV-cache layout, either "NHD" or "HND".
use_cuda_graph (bool) – Whether the wrapper will be used with CUDA Graph capture.
qo_indptr_buf (torch.Tensor, optional) – CUDA Graph buffer for the query/output indptr array.
paged_kv_indptr_buf (torch.Tensor, optional) – CUDA Graph buffer for the paged-KV indptr array.
paged_kv_indices_buf (torch.Tensor, optional) – CUDA Graph buffer for the paged-KV indices array.
paged_kv_last_page_len_buf (torch.Tensor, optional) – CUDA Graph buffer for last-page lengths.
custom_mask_buf (torch.Tensor, optional) – CUDA Graph buffer for a packed custom mask.
mask_indptr_buf (torch.Tensor, optional) – CUDA Graph buffer for the custom-mask indptr array.
backend (str) – Attention backend. "auto" selects a supported backend.
pos_encoding_mode (str) – Positional-encoding mode for the sink attention variant.
use_fp16_qk_reduction (bool) – Whether to use FP16 QK accumulation.
q_data_type (torch.dtype) – Query data type used to specialize the sink attention variant.
kv_data_type (torch.dtype) – KV-cache data type used to specialize the sink attention variant.
head_dim_qk (int) – Query/key head dimension.
head_dim_vo (int) – Value/output head dimension.
window_left (int) – Left sliding-window size. -1 disables sliding-window masking.

Compute batch prefill/append attention between query and paged kv-cache.

Parameters:

q (torch.Tensor) – The query tensor, shape: [qo_indptr[-1], num_qo_heads, head_dim_qk]
paged_kv_cache (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) –
The paged KV-Cache stored as a tuple of tensors or a single tensor:
- a tuple (k_cache, v_cache) of 4-D tensors, each with shape: [max_num_pages, page_size, num_kv_heads, head_dim] if kv_layout is NHD, and [max_num_pages, num_kv_heads, page_size, head_dim] if kv_layout is HND.
- a single 5-D tensor with shape: [max_num_pages, 2, page_size, num_kv_heads, head_dim] if kv_layout is NHD, and [max_num_pages, 2, num_kv_heads, page_size, head_dim] if kv_layout is HND. Where paged_kv_cache[:, 0] is the key-cache and paged_kv_cache[:, 1] is the value-cache.
In these shapes, head_dim is head_dim_qk for the key cache and head_dim_vo for the value cache. For NVFP4 torch.uint8 caches, the stored last dimensions are head_dim_qk // 2 and head_dim_vo // 2, respectively. A single 5-D tensor requires equal key and value storage widths.
*args – Additional arguments for custom kernels.
q_scale (Optional[Union[float, torch.Tensor]]) – The calibration scale of query for fp8 input, if not provided, will be set to 1.0.
k_scale (Optional[Union[float, torch.Tensor]]) – The calibration scale of key for fp8 or nvfp4 input, if not provided, will be set to 1.0.
v_scale (Optional[Union[float, torch.Tensor]]) – The calibration scale of value for fp8 or nvfp4 input, 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 logsumexp of attention output
lse_base (str) – Base of the returned lse when return_lse is True: "log2" (default, FlashInfer’s convention, max + log2(sum(exp2(...)))) or "ln" (natural log, the torch.logsumexp convention that merge kernels written with expf consume). "ln" lets the cuDNN backend hand out its native stats without a conversion kernel; the other backends rescale.
lse_layout (str) – Layout of the returned lse: "NH" (default) is [qo_indptr[-1], num_qo_heads], "HN" is the contiguous transpose [num_qo_heads, qo_indptr[-1]]. Prepared cuDNN graphs write HN directly when supported; other paths transpose and rescale as needed. A caller-provided lse must have the requested layout.
enable_pdl (bool) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/programmatic-dependent-launch.html Only effective on backends and devices that support PDL.
window_left (Optional[int]) – Per-call override for the left (inclusive) sliding-window size. When None, the value supplied to plan() is used. Pass -1 to disable the sliding window for this call.
sinks (Optional[torch.Tensor]) – Per-head attention-sink logits. When provided, the kernel applies the attention-with-sink variant: an additional virtual token whose logit is sinks[head_idx] is appended to each row of the softmax denominator. Shape: [num_qo_heads], dtype float32.
kv_cache_sf (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]) –
Per-block scale factors for NVFP4 KV cache. Accepts the same formats as paged_kv_cache:
- a tuple (k_scales, v_scales) of 4-D tensors, each with shape: [num_pages, page_size, num_kv_heads, head_dim // 16] if kv_layout is NHD, and [num_pages, num_kv_heads, page_size, head_dim // 16] if kv_layout is HND.
- a single 5-D tensor with shape: [num_pages, 2, page_size, num_kv_heads, head_dim // 16] if kv_layout is NHD, and [num_pages, 2, num_kv_heads, page_size, head_dim // 16] if kv_layout is HND, where dim 1 holds k (index 0) and v (index 1) scales.
For the tuple form, head_dim is head_dim_qk in k_scales and head_dim_vo in v_scales. A single 5-D tensor requires equal key and value scale widths.

Both tensors have dtype torch.float8_e4m3fn. k_scales uses a linear (row-major) layout, while v_scales must use TRT-LLM’s 4-token interleaved layout within each [page_size, head_dim_vo // 16] tile if backend is trtllm-gen. Use flashinfer.fp4_quantization.nvfp4_quantize_paged_kv_cache() to produce correctly formatted scale factors.

For the trtllm-gen backend with NHD layout, scale tensors are transposed to HND internally (incurring a copy). Use HND for better performance. Currently, NVFP4 KV supports fa2 and trtllm-gen backend.

With use_token_head_sf=True (plan time), kv_cache_sf must be a (k_sf, v_sf) tuple of float32 tensors mirroring the KV layout minus head_dim: [num_pages, page_size, num_kv_heads] if kv_layout is NHD, and [num_pages, num_kv_heads, page_size] if HND. The tensors may be non-contiguous views (e.g. sliced out of an inline slot layout).
use_fp16_softmax (Optional[bool]) – trtllm-gen backend only. Select the …Fp16Softmax… cubin variant (FP16 softmax accumulator). Currently only shipped for BF16 Q/KV/O context kernels. Ignored by other backends.
uses_spcompress (Optional[bool]) – trtllm-gen backend only. Select the …Spcomp… cubin variant (sparse compression). Currently only shipped for FP8 Q context kernels. Ignored by other backends.
skip_softmax_threshold_scale_factor (Optional[float]) – Threshold scale factor for skipping softmax operations. Providing a value enables skip-softmax sparsity as described in https://arxiv.org/abs/2512.12087. Defaults to None (standard attention). Higher values yield faster kernels at the cost of accuracy; the effective threshold equals the supplied factor divided by the context length.

Returns:

If return_lse is False, the attention output, shape: [qo_indptr[-1], num_qo_heads, head_dim_vo]. If return_lse is True, a tuple of two tensors:

The attention output, shape: [qo_indptr[-1], num_qo_heads, head_dim_vo].
The logsumexp of attention output, shape: [qo_indptr[-1], num_qo_heads] (lse_layout="NH", default) or [num_qo_heads, qo_indptr[-1]] ("HN"), in the base selected by lse_base.

Return type:

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

class flashinfer.attention.BatchDecodeWithAttentionSinkWrapper(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', backend: str = 'auto', q_data_type: dtype = torch.bfloat16, kv_data_type: dtype = torch.bfloat16, head_dim_qk: int = 128, head_dim_vo: int = 128, window_left: int = -1, **kwargs: Any)¶

Decode counterpart of BatchAttentionWithAttentionSinkWrapper.

It runs the same attention-sink kernel through the tensor-core decode path, so it is planned and captured in CUDA graphs like BatchDecodeWithPagedKVCacheWrapper, which receives the remaining keyword arguments. Call run(q, paged_kv_cache, sinks, sm_scale, ...).

__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', backend: str = 'auto', q_data_type: dtype = torch.bfloat16, kv_data_type: dtype = torch.bfloat16, head_dim_qk: int = 128, head_dim_vo: int = 128, window_left: int = -1, **kwargs: Any) → None¶

Constructor of BatchDecodeWithPagedKVCacheWrapper.

Parameters:

float_workspace_buffer (torch.Tensor. Must be initialized to 0 for its first use.) – 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. The buffer must be 16-byte aligned; tensors created by torch.empty satisfy this on supported devices.
kv_layout (str) – The layout of the input k/v tensors, could be either NHD or HND.
use_cuda_graph (bool) – Whether to enable CUDAGraph for batch decode attention, if enabled, the auxiliary data structures will be stored as the provided buffers. The batch_size cannot change during the lifecycle of this wrapper when CUDAGraph is enabled.
use_tensor_cores (bool) – Whether to use tensor cores for the computation. Will be faster for large group size in grouped query attention. Defaults to False.
paged_kv_indptr_buffer (Optional[torch.Tensor]) – The user reserved buffer on GPU to store the indptr of the paged kv cache, the size of the buffer should be [batch_size + 1]. Only needed when use_cuda_graph is True.
paged_kv_indices_buffer (Optional[torch.Tensor]) – The user reserved buffer on GPU to store the page indices of the paged kv cache, should be large enough to store the maximum number of page indices (max_num_pages) during the lifecycle of this wrapper. Only needed when use_cuda_graph is True.
paged_kv_last_page_len_buffer (Optional[torch.Tensor]) – The user reserved buffer on GPU to store the number of entries in the last page, the size of the buffer should be [batch_size]. Only needed when use_cuda_graph is True.
backend (str) – The implementation backend, could be auto/fa2/fa3/trtllm-gen/ cute-dsl/cudnn. Defaults to auto. If set to auto, the wrapper will automatically choose the backend based on the device architecture and kernel availability. On SM100 / SM103 with cudnn-frontend 1.30+ installed, auto resolves to cudnn for multi-token decode rows (2 <= q_len_per_req <= 4, speculative / MTP verification) of fp16 / bf16 head_dim-128 GQA models whose group divides 128, without RoPE / soft-cap / window, when cuDNN’s default SM100 engine serves them with its FROST decode tile and that tile has enough work: q_len_per_req * group in [32, 128] rows per CTA and batch_size * num_kv_heads >= 64 CTAs, where it measures at 0.35-0.95x fa2 (see _auto_decode_prefers_cudnn); single-token decode stays on fa2, except head_dim-256 GQA models (groups 4, 8, 16) with batch_size * num_kv_heads in [96, 256], where cuDNN’s d256 decode tile measures at 0.5-0.95x fa2 (from 64 under use_cuda_graph=True on cudnn-frontend 1.31+, which is told the graph is replayed and splits KV). Multi-token rows without tensor cores have no fa2 kernel, so there auto takes cudnn for every plan its decode path can run (such a plan used to raise). A caller-owned block_tables narrower than 2048 tokens keeps fa2 on cudnn-frontend 1.30 (its placement puts the backend’s prefill-class engine first below that; 1.31+ lifts the limit). auto takes cudnn only where the frontend’s FROST engines can run (its own CuTe-DSL check: nvidia-cutlass-dsl 4.7.0 or newer, above the 4.6.2 the packages’ dependency metadata admit), and callers that pin fa2’s KV split (fixed_split_size / disable_split_kv, the batch-invariance controls) keep fa2 wherever it can serve the rows. resolved_backend reports the choice after plan(). FLASHINFER_DECODE_AUTO_CUDNN=0 keeps fa2, =1 applies the rule on an older frontend too. Under CUDA graphs (use_cuda_graph=True) auto takes cudnn only when plan() receives a caller-owned block_tables (the auto-built table is captured at its first width and cannot grow) or when fa2 cannot serve the rows (q_len_per_req > 1 without tensor cores); the resolution then becomes part of the frozen graph shape and a later plan() may not change it. The cute-dsl backend uses the CuTe DSL GQA decode kernel for Blackwell (SM100+) and only supports a subset of features (equal head_dim_qk/vo, no RoPE/ALiBi/soft-cap). The cudnn backend runs cuDNN’s paged SDPA decode graph (flashinfer.decode.cudnn_batch_decode_with_kv_cache()). It supports fp16/bf16 q/kv/o of one dtype, GQA, return_lse, CUDA graphs, multi-token decode (q_len_per_req > 1, bottom-right causal), a left sliding window (window_left) and attention sinks, but not RoPE, soft-cap, fp8 or NVFP4 KV. A sink at q_len_per_req == 1 is served when the cuDNN stack’s SDPA engines accept it (cudnn-frontend 1.30+ with the FROST engines enabled); the backend engine raises a not-supported error at the first run. Pass a fixed block_tables to plan() when capturing CUDA graphs so the captured block-table buffer is stable across re-plans. On SM90, tensor-core auto graph-mode wrappers may select XQA MTP for FP16 NHD, page16, head128, 32Q/8KV-head plans with uniform KV lengths: q8/KV8192 at batch4/16/64, or q4/q8/KV1024 at batch16. This provisional rule excludes ordinary wrappers, explicit backends, q2 and unmeasured shapes. It is based on selected historical H100 graph results, not a guarantee of gains on every SM90 GPU. The preferred route is frozen for graph-mode wrappers; changes in eligibility require a separate wrapper and graph capture. Eager warmups of such wrappers use their selected graph-mode route. workspace_size is unavailable for XQA-selected shapes.
jit_args (Optional[List[Any]]) – If provided, the wrapper will use the provided arguments to create the JIT module, otherwise, the wrapper will use default attention implementation.
jit_kwargs (Optional[Dict[str, Any]]) – The keyword arguments to create the JIT module, defaults to None.

SM120 NVFP4 Attention¶

`nvfp4_attention_sm120_quantize_qkv`(q, k, v)	Preprocess and quantize dense Q/K/V tensors for SM120 NVFP4 attention.
`nvfp4_attention_sm120_fwd`(q_fp4, k_fp4, ...)	Run SM120 NVFP4 attention on pre-quantized Q/K/V tensors.

flashinfer.mla¶

MLA (Multi-head Latent Attention) is an attention mechanism proposed in DeepSeek series of models ( DeepSeek-V2, DeepSeek-V3, and DeepSeek-R1).

Sliding-window MLA (SM90)¶

flashinfer.mla.batch_decode_with_kv_cache_mla_swa(query: Tensor, kv_cache: Tensor, workspace_buffer: Tensor, block_tables: Tensor, seq_lens: Tensor, *, window_left: int, sm_scale: float, out: Tensor | None = None, backend: str = 'auto') → Tensor¶

Causal sliding-window MLA decode over a shared latent Page64 KV cache.

Parameters:

query (torch.Tensor) – BF16 [B, q_len_per_req, 64, 1088]; B1/2/4/8/16 and q_len1/4.
kv_cache (torch.Tensor) – BF16 contiguous [num_pages, 64, 1088]. The first 1024 features are both latent K and V; the last 64 features are positional K. Attended K/V values must be finite. Masked page slots need not be initialized; they do not contribute even if they contain NaN or Inf.
workspace_buffer (torch.Tensor) – Caller-owned CUDA uint8 buffer sized by get_mla_swa_workspace_size_bytes(). Do not share writable workspace across concurrent calls/streams.
block_tables (torch.Tensor) – CUDA int32 [B, max_pages], mapping logical pages to physical pages.
seq_lens (torch.Tensor) – CUDA int32 [B], total KV lengths including the query tokens. Query i has position seq_lens[b] - q_len_per_req + i. Valid page IDs must be in range and each nonempty length must fit its table; invalid metadata values are not checked synchronously. Zero-length dummy rows produce zero output.
window_left (int) – Number of past tokens; currently 512 (513 including the current token).
sm_scale (float) – Positive finite attention scale, supplied by the model.
out (torch.Tensor, optional) – Contiguous BF16 [B, q_len_per_req, 64, 1024]. If omitted, allocated. Pass a persistent output and warm up eagerly before Graph capture.
backend (str) – "auto" or "sm90". Both use the SM90 implementation; there is no hidden fallback to a different operator/precision on other devices.

Returns:

The supplied or newly allocated output. All split reductions retain FP32 until the final BF16 write. Use flashinfer.autotune(True) before capture to search the six tactics; otherwise a frozen LUT is used.

Return type:

torch.Tensor

flashinfer.mla.get_mla_swa_workspace_size_bytes(batch_size: int, q_len_per_req: int, num_q_heads: int = 64) → int¶

Return byte workspace capacity for every SM90 MLA+SWA tuning candidate.

Parameters:

batch_size (int) – Local batch size, one of 1, 2, 4, 8, 16.
q_len_per_req (int) – Query length per request, 1 or 4.
num_q_heads (int) – Local query head count; currently 64.

Returns:

Capacity for FP32 partial output and LSE with up to nine splits, including 256-byte padding. The final output is not part of workspace. Even DSM/direct calls require this capacity so autotune can change tactics.

Return type:

int

PageAttention for MLA¶

`trtllm_batch_decode_with_kv_cache_mla`(query, ...)	Execute MLA decode or prefill through the selected backend: TRTLLM-GEN, CuteDSL, XQA, Cake, or SM120/SM121 sparse kernels.
`trtllm_prefill_with_kv_cache_mla`(query, ...)	Run MLA with decode-compatible semantics for prefill callers.
`trtllm_batch_decode_sparse_mla_dsv4`(query, ...)	Decode DeepSeek V4 sparse MLA.
`nvfp4_quantize_pack_sparse_mla_cache`(...[, ...])	Quantize complete DeepSeek-V4 latent-KV pages to the NVFP4 cache ABI.
`nvfp4_quantize_append_sparse_mla_cache`(...)	Quantize and append DeepSeek-V4 latent KV by physical cache slot.
`dsv41_fp4_quantize_pack_sparse_mla_cache`(...)	Quantize complete DeepSeek-V4.1 latent-KV pages to the V41_FP4 ABI.
`dsv41_fp4_quantize_append_sparse_mla_cache`(...)	Quantize and append DeepSeek-V4.1 latent KV by physical cache slot.
`dsv41_fp8_quantize_pack_sparse_mla_cache`(...)	Quantize complete DeepSeek-V4.1 latent-KV pages to the 528-byte DSV4_1 FP8 ABI.
`dsv41_fp8_quantize_append_sparse_mla_cache`(...)	Quantize and append DeepSeek-V4.1 latent KV by physical slot into a 528-byte FP8 cache.
`convert_compressed_page_aligned_sparse_indices_to_hca_metadata`(...)	Convert combined sparse indices into reusable HCA metadata.
`DSV4HCAMetadata`(hca_swa_indices, ...[, ...])	Reusable metadata for the CuTe DSL causal HCA backend.
`xqa_batch_decode_with_kv_cache_mla`(query, ...)	XQA-backend batched MLA decode.
`supported_sparse_mla_sm120_configs`(*[, ...])	Enumerate the instantiated SM120 sparse-MLA decode kernel configurations.
`SparseMLASm120DecodeConfig`(d_qk, ...[, ...])	Instantiated decode-kernel set for one SM120 sparse-MLA kernel family.
`SparseMLASm120Wrapper`	alias of `_SparseMLAPagedAttentionRunner`
`cake_sparse_mla_sm120_dsv4_nvfp4_decode`(q, ...)	Run the allocation-free Cake SM120 NVFP4 sparse-MLA decode.
`cake_sparse_mla_sm120_dsv4_nvfp4_prefill`(q, ...)	Run the allocation-free Cake SM120 NVFP4 sparse-MLA prefill (one launch, no scratch).
`cake_sparse_mla_sm120_dsv4_nvfp4_select_kernel`(*, ...)	`"decode"` or `"prefill"` for one `backend="cake"` call (measured crossover).
`cake_sparse_mla_sm120_dsv4_nvfp4_plan_prefill`(*, ...)	Return `head_tiles` for one prefill call (measured crossover, see `_PREFILL_POLICY`).
`cake_sparse_mla_sm120_dsv4_nvfp4_plan_head_tiles`(*, ...)	Return the number of 16-head tiles per decode CTA (1 or 2).
`cake_sparse_mla_sm120_dsv4_nvfp4_plan_splits`(*, ...)	Return `(num_splits, chunks_per_block)` for one decode call.
`cake_sparse_mla_sm120_dsv4_nvfp4_scratch_bytes`(...)	Workspace bytes that cover every split plan of this shape (partials + LSE + alignment slack).
`cake_dsv4_nvfp4_rope_quantize_insert`(q, kv, ...)	Rotate `q`, then rotate, quantize and insert `kv` into the DeepSeek-V4 NVFP4 cache.
`cake_dsv4_nvfp4_kv_rope_quantize_insert`(kv, ...)	Rotate, quantize and insert latent KV rows into the DeepSeek-V4 NVFP4 cache (KV only).

Note

With backend="cake" on SM120/SM121 (kv_cache_format="nvfp4"), every call runs the Cake split decode kernel or the single-launch prefill kernel as chosen by cake_sparse_mla_sm120_dsv4_nvfp4_select_kernel from the token, head and candidate counts and the SM count (DeepSeek-V4 sparse MLA SM120 tracker: flashinfer#4254).

Note

The NVFP4 cache that route reads is written per layer by the fused SM120/SM121 writers cake_dsv4_nvfp4_rope_quantize_insert (sliding-window pool: GPT-J RoPE of the query and latent KV, head-padded q_out or, with q_inplace=True and no head padding, the query rotated in place, NVFP4 quantization and paged insert in one launch) and cake_dsv4_nvfp4_kv_rope_quantize_insert (compressed pool with compress_ratio 1 or 2, speculative context with ratio 1); both produce the bytes of nvfp4_quantize_append_sparse_mla_cache applied to the BF16-rounded roped rows.

Note

With backend="cute-dsl", pass hca_swa_indices as absolute rows into the flattened SWA cache and hca_compressed_block_tables as physical compressed-cache page IDs. The SWA table has shape [B * Q, 128] and may express ring rotation or wraparound. Combined tables whose compressed segment is a canonical page expansion can opt into compatibility conversion with hca_sparse_indices_format="compressed-page-aligned". SWA entries remain arbitrary absolute rows. Precompute that conversion before a CUDA Graph or a latency-sensitive loop.

Note

kv_cache_format="nvfp4" (the 384-byte-per-token DeepSeek-V4 NVFP4 sparse cache: 448 NoPE values as E2M1 with one E4M3 scale per 16 values, 64 BF16 RoPE values, page_size * 352 data bytes followed by page_size * 32 scale bytes per page) is consumed by backend="sparse" on SM120 / SM121 and by backend="cake" on SM100 / SM103 (B200 / GB300). The CAKE route takes a BF16 query, two independent tables (sparse_indices over swa_kv_cache and extra_sparse_indices over compressed_kv_cache with their own *_topk_lens; -1 entries are masked), sinks, a caller-owned workspace_buffer sized by flashinfer.mla.cake_dsv4.get_cake_dsv4_workspace_bytes() and is CUDA-Graph safe. Build the cache with nvfp4_quantize_pack_sparse_mla_cache() / nvfp4_quantize_append_sparse_mla_cache() (one implementation for all four architectures). backend="auto" keeps selecting TRTLLM-GEN on SM100 / SM103; pass backend="cake" explicitly.

class flashinfer.mla.BatchMLAPagedAttentionWrapper(float_workspace_buffer: Tensor, use_cuda_graph: bool = False, qo_indptr: Tensor | None = None, kv_indptr: Tensor | None = None, kv_indices: Tensor | None = None, kv_len_arr: Tensor | None = None, backend: str = 'auto')¶

Wrapper for MLA PagedAttention on DeepSeek models.

This wrapper is intended for decode and incremental prefill with the Matrix Absorption formulation of MLA, where the query/key and value/output projections are absorbed before attention. For the non-absorbed MLA prefill path, use the appropriate prefill wrapper instead.

The planned-wrapper surface owns FA2, FA3, CUTLASS, cuTile, TRTLLM-GEN, XQA, and CuTe DSL planning and execution. Call plan() once with canonical metadata before invoking run(); the plan captures the supported input/output contract and the concrete backend’s metadata representation.

See MLA Page Layout for the paged KV-cache layout and the FlashInfer 0.2 release post for the computation and Matrix Absorption background.

__init__(float_workspace_buffer: Tensor, use_cuda_graph: bool = False, qo_indptr: Tensor | None = None, kv_indptr: Tensor | None = None, kv_indices: Tensor | None = None, kv_len_arr: Tensor | None = None, backend: str = 'auto') → None¶

Construct a planned MLA wrapper.

Parameters:

float_workspace_buffer (torch.Tensor) – Caller-owned workspace for intermediate attention results. A 128 MiB buffer is the usual starting point; it must be on the same device as the query and KV-cache tensors.
use_cuda_graph (bool, optional) – Enable CUDA-graph-compatible planning. When enabled, the optional metadata buffers below are copied at plan() time to preserve capture-time pointers. The captured batch shape cannot change.
qo_indptr (Optional[torch.Tensor]) – Caller-reserved int32 buffers of shape [batch_size + 1] for CSR metadata. Used only with use_cuda_graph=True.
kv_indptr (Optional[torch.Tensor]) – Caller-reserved int32 buffers of shape [batch_size + 1] for CSR metadata. Used only with use_cuda_graph=True.
kv_indices (Optional[torch.Tensor]) – Caller-reserved int32 CSR page-index buffer, sized for the maximum planned number of pages. Used only with CUDA graphs.
kv_len_arr (Optional[torch.Tensor]) – Caller-reserved int32 buffer of shape [batch_size] for CSR KV lengths. Used only with CUDA graphs.
backend ({"auto", "autotune", "fa2", "fa3", "cutlass", "cutile", "trtllm-gen", "xqa", "cute-dsl", "cute-dsl-monolithic", "cute-dsl-modular", "cute-dsl-rubin-mtp"}) – Requested policy or concrete backend. "auto" selects a backend in plan() based on architecture and request facts. See _auto_policy.py. Explicit requests remain strict; "cute-dsl" is a family alias. "cute-dsl-rubin-mtp" supports SM107 causal absorbed MLA with uniform Q=2/4, 128 heads, latent/RoPE widths 512/64, page size 64/128, and FP8 E4M3 query, KV, and output tensors. It requires a CuTe DSL compiler with native SM107 and mixed-CGA support. "autotune" benchmarks and selects among eligible backends during run() inside with flashinfer.autotune(True):. See _auto_policy.py. Canonical dense or CSR metadata is accepted; independent split query or KV storage requires a backend with native split-input support. Existing non-FA graph plans allow one plan per wrapper. Use a new wrapper when preparing a different backend for a new capture.

plan(*, metadata: MLAPlanMetadata, num_heads: int, head_dim_ckv: int, head_dim_kpe: int, page_size: int, causal: bool, sm_scale: float, q_data_type: dtype, kv_data_type: dtype, use_profiler: bool = False, query_layout: Literal['packed', 'split'] = 'packed', kv_cache_layout: Literal['packed', 'split'] = 'packed', lse_mode: Literal['none', 'base2', 'basee'] = 'none', output_dtype: dtype | None = None, output_scale: Literal['none', 'per-tensor'] = 'none', scale_mode: Literal['default', 'kv-per-tensor', 'bmm-scalar', 'bmm-tensor'] = 'default', skip_softmax: bool = False, enable_pdl: bool | None = None, use_sinks: bool = False) → None¶

plan(qo_indptr: Tensor, kv_indptr: Tensor, kv_indices: Tensor, kv_len_arr: Tensor, num_heads: int, head_dim_ckv: int, head_dim_kpe: int, page_size: int, causal: bool, sm_scale: float, q_data_type: dtype, kv_data_type: dtype, use_profiler: bool = False, *, query_layout: Literal['packed', 'split'] = 'split', kv_cache_layout: Literal['packed', 'split'] = 'split', lse_mode: Literal['none', 'base2', 'basee'] = 'none', output_dtype: dtype | None = None, output_scale: Literal['none', 'per-tensor'] = 'none', scale_mode: Literal['default', 'kv-per-tensor', 'bmm-scalar', 'bmm-tensor'] = 'default', skip_softmax: bool = False, enable_pdl: bool | None = None, use_sinks: bool = False) → None

plan(*, cum_seq_lens_q: Tensor, block_tables: Tensor, seq_lens: Tensor, num_heads: int, head_dim_ckv: int, head_dim_kpe: int, page_size: int, causal: bool, sm_scale: float, q_data_type: dtype, kv_data_type: dtype, max_q_len: int | None = None, use_profiler: bool = False, query_layout: Literal['packed', 'split'] = 'split', kv_cache_layout: Literal['packed', 'split'] = 'split', lse_mode: Literal['none', 'base2', 'basee'] = 'none', output_dtype: dtype | None = None, output_scale: Literal['none', 'per-tensor'] = 'none', scale_mode: Literal['default', 'kv-per-tensor', 'bmm-scalar', 'bmm-tensor'] = 'default', skip_softmax: bool = False, enable_pdl: bool | None = None, use_sinks: bool = False) → None

Plan a concrete MLA backend from canonical metadata.

Prefer metadata=MLAPlanMetadata.csr(...) for the CSR form or metadata=MLAPlanMetadata.dense(...) for the dense page-table form. MLAPlanMetadata.dual(...) may be used when both representations already exist; the planner verifies that they describe the same requests and page mapping before publishing the plan. FA2 and FA3 consume CSR metadata natively, while CUTLASS and cuTile consume dense metadata.

Metadata tensors may be on CPU or the wrapper device. They are normalized to the device required by the selected backend; tensors on another accelerator device are rejected. Passing flat CSR or dense metadata fields remains supported for compatibility, but is deprecated in favor of the metadata= object form. TRTLLM-GEN, XQA, and CuTe DSL CUDA graph plans require supplied contiguous int32 dense metadata on the wrapper device. Retain and update those device tensors in place for replay; CPU and CSR-only graph metadata are rejected. Query offsets and lengths must remain fixed after graph planning. backend="autotune" instead snapshots metadata during planning; modifying caller metadata does not update that plan. Replan eagerly, or create a new wrapper for a different CUDA graph workload.

The plan also declares the later run() contract. In particular, query_layout, kv_cache_layout, lse_mode, output_dtype, output_scale, scale_mode, skip_softmax, and use_sinks must agree with the subsequent call. Backends may also consume enable_pdl at plan time. Sequence variability is inferred from canonical metadata by backends that require a specialized plan. Canonical metadata defaults to packed inputs; the legacy flat forms retain their historical split-input defaults. Deprecated flat CSR plans also temporarily retain dynamic LSE behavior on FA2/FA3 and emit a warning when it is used.

Parameters:

metadata (Optional[MLAPlanMetadata]) – Preferred canonical CSR, dense, or dual metadata representation.
qo_indptr (Optional[torch.Tensor]) – Deprecated flat CSR metadata fields.
kv_indptr (Optional[torch.Tensor]) – Deprecated flat CSR metadata fields.
kv_indices (Optional[torch.Tensor]) – Deprecated flat CSR metadata fields.
kv_len_arr (Optional[torch.Tensor]) – Deprecated flat CSR metadata fields.
cum_seq_lens_q (Optional[torch.Tensor]) – Deprecated flat dense page-table metadata fields.
block_tables (Optional[torch.Tensor]) – Deprecated flat dense page-table metadata fields.
seq_lens (Optional[torch.Tensor]) – Deprecated flat dense page-table metadata fields.
max_q_len (Optional[int]) – Maximum dense query length; inferred from query metadata when omitted. For CuTe DSL monolithic plans created with nonuniform query lengths, this also sets the per-request CUDA graph launch capacity. Those plans may update query offsets in place while preserving batch size, total tokens, tensor addresses/shapes, and positive request lengths within this capacity. Initially uniform plans keep query lengths fixed. Other selected backends retain their own graph contracts; automatic selection does not guarantee support for query redistribution. Direct graph replay does not validate updated offsets in Python.
num_heads (Optional[int]) – Number of query heads.
head_dim_ckv (Optional[int]) – Compressed-KV and RoPE feature widths.
head_dim_kpe (Optional[int]) – Compressed-KV and RoPE feature widths.
page_size (Optional[int]) – Number of KV tokens in a cache page.
causal (Optional[bool]) – Whether the planned attention is causal.
sm_scale (Optional[float]) – Softmax scale captured by the plan.
q_data_type (Optional[torch.dtype]) – Query and KV-cache dtypes accepted by the selected backend.
kv_data_type (Optional[torch.dtype]) – Query and KV-cache dtypes accepted by the selected backend.
use_profiler (bool) – Whether to enable backend profiler support.
query_layout (Optional[{"packed", "split"}]) – Tensor representations accepted by subsequent run() calls.
kv_cache_layout (Optional[{"packed", "split"}]) – Tensor representations accepted by subsequent run() calls.
lse_mode ({"none", "base2", "basee"}) – Required log-sum-exp output mode.
output_dtype (Optional[torch.dtype]) – Required output dtype; defaults to q_data_type. Set torch.float8_e4m3fn for Rubin MTP; monolithic CuTe also supports this output dtype, while modular CuTe requires FP16/BF16.
output_scale ({"none", "per-tensor"}) – Required output scaling mode.
scale_mode ({"default", "kv-per-tensor", "bmm-scalar", "bmm-tensor"}) – Required KV or BMM scale mode for subsequent run() calls. TRTLLM-GEN supports "bmm-tensor" for FP8 query/KV inputs.
skip_softmax (bool) – Whether the plan must support the skip-softmax threshold feature.
enable_pdl (optional) – Backend-specific planning controls. A backend that does not support a non-default value rejects it during planning.
use_sinks (optional) – Backend-specific planning controls. A backend that does not support a non-default value rejects it during planning.

Notes

Positional arguments are deprecated; use keyword arguments for all plan() parameters. Flat metadata arguments are also deprecated; use an MLAPlanMetadata object instead.

Run one planned MLA attention request.

Preferred structural input form

Pass query= and kv_cache=. Each may be a packed tensor whose final dimension joins its no-PE and PE features, or a split pair of tensors. The representation must satisfy the layouts declared by plan(). Adjacent split views can be reinterpreted as packed tensors without a copy; FA2 and FA3 also accept independent split inputs natively. Packed-native CUTLASS plans reject independent split KV caches instead of silently copying them.

Output forms

With return_lse=False (the default), the method returns the output tensor. With return_lse=True, it returns (output, lse) when the canonical plan was created with a matching LSE mode. return_lse_base_on_e selects natural-log rather than base-2 LSE output. Deprecated flat CSR FA2/FA3 plans temporarily retain dynamic LSE behavior and emit a warning. Caller-provided out and lse buffers are used directly and returned by identity.

Legacy separate-parameter form

The separate q_nope / q_pe and ckv_cache / kpe_cache parameters remain a deprecated compatibility form. Pass the same split tensors through query=(q_nope, q_pe) and kv_cache=(ckv_cache, kpe_cache) instead. Structural split values are not deprecated. Positional arguments are deprecated independently; each warning is emitted once per wrapper instance.

Parameters:

query (object) – Preferred packed or split structural query value.
kv_cache (object) – Preferred packed or split structural KV-cache value.
q_nope (Optional[torch.Tensor]) – Legacy split query tensors. Supply both together.
q_pe (Optional[torch.Tensor]) – Legacy split query tensors. Supply both together.
ckv_cache (Optional[torch.Tensor]) – Legacy split compressed-KV and PE cache tensors. Supply both together.
kpe_cache (Optional[torch.Tensor]) – Legacy split compressed-KV and PE cache tensors. Supply both together.
out (Optional[torch.Tensor]) – Caller-owned output buffer. When o_scale is provided, it must be an FP8 tensor for CUTLASS output.
lse (Optional[torch.Tensor]) – Caller-owned log-sum-exp output buffer for an LSE-enabled plan.
return_lse (bool) – Return the LSE tensor in addition to the output.
profiler_buffer (Optional[torch.Tensor]) – Backend profiler output buffer.
kv_len (Optional[torch.Tensor]) – CUTLASS/cuTile metadata aliases. A planned request may omit them; the deprecated unplanned CUTLASS path requires both. Runtime metadata is a trusted hot-path input: every length must be nonnegative and fit within its page-table row, every live page ID must index kv_cache, and callers must not mutate either tensor while a launch is in flight. These values are not synchronized to the host for validation so that CUDA-graph capture remains valid.
page_table (Optional[torch.Tensor]) – CUTLASS/cuTile metadata aliases. A planned request may omit them; the deprecated unplanned CUTLASS path requires both. Runtime metadata is a trusted hot-path input: every length must be nonnegative and fit within its page-table row, every live page ID must index kv_cache, and callers must not mutate either tensor while a launch is in flight. These values are not synchronized to the host for validation so that CUDA-graph capture remains valid.
return_lse_base_on_e (bool) – Return natural-log rather than base-2 LSE values.
o_scale (Optional[float]) – Per-tensor FP8 output scale supported by CUTLASS.
ckv_scale (optional) – Per-tensor or per-token FP8 KV-cache scales required by plans that selected scale_mode="kv-per-tensor".
ckv_scale_arr (optional) – Per-tensor or per-token FP8 KV-cache scales required by plans that selected scale_mode="kv-per-tensor".
kpe_scale (optional) – Per-tensor or per-token FP8 KV-cache scales required by plans that selected scale_mode="kv-per-tensor".
sinks (Optional[torch.Tensor]) – Per-head sink values. This must be present exactly when the plan used use_sinks=True.
skip_softmax_threshold_scale_factor (Optional[float]) – Skip-softmax threshold scale. This must be present exactly when the plan used skip_softmax=True.
bmm1_scale (Optional[float or torch.Tensor]) – Backend-specific BMM scale pair. Both values must be provided together, use the same scalar/tensor family, and match the planned scale_mode. TRTLLM-GEN tensor scales must each contain one float32 value on the query device. Their values may change between runs and CUDA graph replays; graph replay reads the captured tensors.
bmm2_scale (Optional[float or torch.Tensor]) – Backend-specific BMM scale pair. Both values must be provided together, use the same scalar/tensor family, and match the planned scale_mode. TRTLLM-GEN tensor scales must each contain one float32 value on the query device. Their values may change between runs and CUDA graph replays; graph replay reads the captured tensors.

Notes

Positional arguments are deprecated; use keyword arguments for plan() and run(). The separate q_nope / q_pe and ckv_cache / kpe_cache parameters are also deprecated; pass structural query= and kv_cache= values instead. An explicitly requested CUTLASS backend may still run without plan() when both kv_len and page_table are supplied, but that compatibility path is also deprecated.