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 KV Cache for single request, return 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)¶

Wrapper class for decode attention with paged kv-cache (first proposed in vLLM) for 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) → 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. 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.
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.

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. q_len_per_req > 1 is currently supported on the fa2 tensor-core backend (and natively by trtllm-gen/cute-dsl). 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.

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().

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().
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.
*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-c-programming-guide/#programmatic-dependent-launch-and-synchronization 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 fa2 and trtllm-gen backend.

Returns:

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

attention output, shape: [batch_size, num_qo_heads, head_dim]
logsumexp of attention scores, shape: [batch_size, 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 paged kv-cache (first proposed in vLLM) for 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. :param q: Query tensor with shape [batch_size, beam_width, num_q_heads, head_dim] if not using speculative decoding, or [batch_size, beam_width, q_seq_len, num_q_heads, head_dim] if using speculative decoding. q_seq_len is the number of speculative decoding tokens. Data type should be torch.float16 or torch.bfloat16. Now only beam_width 1 is supported. :type q: torch.Tensor :param k_cache: Paged K cache tensor with shape [num_pages, page_size, num_kv_heads, head_dim] if kv_layout is NHD, or [num_pages, num_kv_heads, page_size, head_dim] if kv_layout is HND. Data type should match query tensor or be torch.float8_e4m3fn, in which case xqa will run fp8 calculation. Should be the same data type as v_cache. When using NVFP4 KV, the data type is torch.uint8, and the last dimension should be head_dim / 2. :type k_cache: torch.Tensor :param v_cache: Paged V cache tensor with shape [num_pages, page_size, num_kv_heads, head_dim] if kv_layout is NHD, or [num_pages, num_kv_heads, page_size, head_dim] if kv_layout is HND. Data type should match query tensor or be torch.float8_e4m3fn, in which case xqa will run fp8 calculation. Should be the same data type as k_cache. When using NVFP4 KV, the data type is torch.uint8, and the last dimension should be head_dim / 2. :type v_cache: torch.Tensor :param k_sf_cache: Optional scale factor cache tensor for the K cache. Use when NVFP4 KV is used. Expected shape is [num_pages, page_size, num_kv_heads, head_dim / 16] if kv_layout is NHD, or [num_pages, num_kv_heads, page_size, head_dim / 16] if kv_layout is HND. Should be the same data type as v_sf_cache. Data type should be torch.uint8. :type k_sf_cache: Optional[torch.Tensor] :param v_sf_cache: Optional scale factor cache tensor for the V cache. Use when NVFP4 KV is used. Expected shape is [num_pages, page_size, num_kv_heads, head_dim / 16] if kv_layout is NHD, or [num_pages, num_kv_heads, page_size, head_dim / 16] if kv_layout is HND. Should be the same data type as k_sf_cache. Data type should be torch.uint8. :type v_sf_cache: Optional[torch.Tensor] :param page_table: Page table tensor with shape batch_size, nb_pages_per_seq. Data type should be torch.int32. K and V share the same table. :type page_table: torch.Tensor :param seq_lens: Sequence lengths tensor with shape [batch_size, beam_width]. Data type should be torch.uint32. :type seq_lens: torch.Tensor :param output: Output tensor with shape that matches the query tensor. Data type should match query tensor or kv tensor. This tensor will be modified in-place. :type output: torch.Tensor :param workspace_buffer: Workspace buffer for temporary computations. Data type should be torch.uint8. :type workspace_buffer: torch.Tensor :param semaphores: Semaphore buffer for synchronization. Data type should be torch.uint32. :type semaphores: torch.Tensor :param num_kv_heads: Number of key-value heads in the attention mechanism. :type num_kv_heads: int :param page_size: Size of each page in the paged KV cache. Must be one of [16, 32, 64, 128]. :type page_size: int :param sinks: Attention sink values with shape [num_kv_heads, head_group_ratio]. Data type should be torch.float32. If None, no attention sinks are used. :type sinks: Optional[torch.Tensor], default=None :param q_scale: Scale factor for query tensor. :type q_scale: Union[float, torch.Tensor], default=1.0 :param kv_scale: Scale factor for KV cache. :type kv_scale: Union[float, torch.Tensor], default=1.0 :param sliding_win_size: Sliding window size for attention. If 0, no sliding window is used. With speculative decoding (q_seq_len > 1), the window is applied per draft row, assuming draft tokens occupy consecutive positions at the end of the sequence (linear chains, causal or full masks); tree-structured drafts are not supported with sliding window. :type sliding_win_size: int, default=0 :param kv_layout: The layout of the KV cache. Can be either NHD or HND. :type kv_layout: str, default="NHD" :param sm_count: Number of streaming multiprocessors to use. If None, will be inferred from the device. :type sm_count: Optional[int], default=None :param enable_pdl: Whether to enable PDL (Persistent Data Loader) optimization. If None, will be set to True if hardware supports it. :type enable_pdl: Optional[bool], default=None :param rcp_out_scale: Reciprocal of output scale factor. :type rcp_out_scale: float, default=1.0 :param q_seq_len: Query sequence length. When > 1, enables speculative decoding mode. With ragged Q (q_cu_seq_lens set), this is the maximum draft length across the batch. :type q_seq_len: int, default=1 :param mask: Draft-block attention mask for speculative decoding mode (when q_seq_len > 1). Shape: [batch_size, q_seq_len, mask_size_per_row] where mask_size_per_row = ((q_seq_len + 31) // 32) * 2, or [total_q_tokens, mask_size_per_row] with ragged Q, where each request's rows cover its own draft length. Data type should be torch.uint16 (bit-packed format, aligned to 32 bits). :type mask: Optional[torch.Tensor], default=None :param q_cu_seq_lens: Cumulative draft lengths [batch_size + 1] (torch.int32 or torch.uint32, on device) enabling ragged Q: requests may have different draft lengths, including 0 for a request with no draft tokens this step. When set, q and output are packed as [total_q_tokens, num_q_heads, head_dim], q_seq_len must be the maximum draft length, and mask rows are packed by the same cumulative offsets. :type q_cu_seq_lens: Optional[torch.Tensor], default=None.

xqa_mla(q, k_cache, v_cache, page_table, ...)

Apply attention with paged KV cache using XQA MLA (Multi-Head Latent Attention) kernel. :param q: Query tensor with shape [batch_size, beam_width, num_q_heads, head_dim]. Data type should be torch.float8_e4m3fn. Now only beam_width 1 is supported. :type q: torch.Tensor :param k_cache: Paged K cache tensor with shape [total_num_cache_heads, head_dim]. Data type should be torch.float8_e4m3fn :type k_cache: torch.Tensor :param v_cache: Paged V cache tensor with shape [total_num_cache_heads, head_dim]. Data type should be torch.float8_e4m3fn :type v_cache: torch.Tensor :param page_table: Page table tensor with shape batch_size, nb_pages_per_seq. Data type should be torch.int32. K and V share the same table. :type page_table: torch.Tensor :param seq_lens: Sequence lengths tensor with shape [batch_size, beam_width]. Data type should be torch.uint32. :type seq_lens: torch.Tensor :param output: Output tensor with shape [batch_size, beam_width, num_q_heads, head_dim]. Data type should be torch.bfloat16. This tensor will be modified in-place. :type output: torch.Tensor :param workspace_buffer: Workspace buffer for temporary computations. Data type should be torch.uint8. :type workspace_buffer: torch.Tensor :param semaphores: Semaphore buffer for synchronization. Data type should be torch.uint32. :type semaphores: torch.Tensor :param page_size: Size of each page in the paged KV cache. Must be one of [16, 32, 64, 128]. :type page_size: int :param q_scale: Scale factor for query tensor. :type q_scale: Union[float, torch.Tensor], default=1.0 :param kv_scale: Scale factor for KV cache. :type kv_scale: Union[float, torch.Tensor], default=1.0 :param sm_count: Number of streaming multiprocessors to use. If None, will be inferred from the device. :type sm_count: Optional[int], default=None :param enable_pdl: Whether to enable PDL (Persistent Data Loader) optimization. If None, will be set to True if hardware supports it. :type enable_pdl: Optional[bool], default=None.

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 KV cache for single request, 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, ...)
`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)¶

Wrapper class for prefill/append attention with paged kv-cache for 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) → 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.
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. 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.

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) → 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.
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).

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]
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.
*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]] (one fused transpose + rescale kernel). 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-c-programming-guide/#programmatic-dependent-launch-and-synchronization 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.
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 // 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.
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]. If return_lse is True, a tuple of two tensors:

The attention output, shape: [qo_indptr[-1], num_qo_heads, head_dim].
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 ragged (tensor) kv-cache for 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, when no sliding window, logits soft cap, custom mask or multi-item-scoring pointer is in play – neither path accepts those, 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, uniform_q_len: int | None = None) → 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].
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.

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]] (one fused transpose + rescale kernel). 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-c-programming-guide/#programmatic-dependent-launch-and-synchronization 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.

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)¶

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.

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) → 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.
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. 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.

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 (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¶

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

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).

PageAttention for MLA¶

`trtllm_batch_decode_with_kv_cache_mla`(query, ...)	Decode MLA with Blackwell, TRTLLM-GEN, CuteDSL, XQA, or 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"), one query token runs the Cake split decode kernel and several tokens run the decode or the single-launch prefill kernel as chosen by cake_sparse_mla_sm120_dsv4_nvfp4_select_kernel (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 MLA blog 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", "fa2", "fa3", "cutlass", "cutile", "trtllm-gen", "xqa", "cute-dsl", "cute-dsl-monolithic", "cute-dsl-modular"}) – 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. Compilation and selection finish before normal execution or capture. 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.

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.
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.
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. Tensor BMM scales are reserved for later backend support and are rejected by this wrapper version.
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. Tensor BMM scales are reserved for later backend support and are rejected by this wrapper version.

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.