API Reference

ffpa_attn.ffpa_attn_func

ffpa_attn_func(query: Tensor, key: Tensor, value: Tensor, attn_mask: Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs: object) -> torch.Tensor

FFPA: Faster Flash Prefill Attention for large headdims (D > 256).

Signature aligned with torch.nn.functional.scaled_dot_product_attention. Dispatches by query.dtype (fp16 / bf16) and acc through a single registered torch op (torch.ops.ffpa_attn.attn), keeping the Python layer minimal and fully compatible with torch.compile.

Supports cross-attention where query seqlen (Nq) differs from key/value seqlen (Nkv) and grouped-query attention where query has more heads than key/value (MQA is the Nh_kv == 1 special case). key and value must share the same Nh_kv and the same Nkv. Causal masking is supported via is_causal=True with queries aligned to the tail of the KV sequence (Nkv >= Nq required).

Backward pass is supported via :class:FFPAAttnFunc. The public API falls back to SDPA for cases FFPA does not currently support directly (small-D, D > 1024, and unsupported large-D dropout), and otherwise keeps the existing FFPA forward plus SDPA/FFPA backward routing. Large-D Triton forward and backward support explicit additive attn_mask gradients. forward_backend only affects the large-D path and is accepted as an FFPA-specific keyword inside **kwargs so the explicit signature remains aligned with :func:torch.nn.functional.scaled_dot_product_attention.

Parameters:

Name	Type	Description	Default
query	Tensor	Query tensor with layout [B, Nh_q, Nq, D]; dtype must be torch.float16 or torch.bfloat16 and match key / value.	required
key	Tensor	Key tensor with layout [B, Nh_kv, Nkv, D]; same dtype as query. Nh_q must be an integer multiple of Nh_kv (group_size = Nh_q / Nh_kv).	required
value	Tensor	Value tensor with layout [B, Nh_kv, Nkv, D]; same dtype as query. key and value must share the same Nh_kv and Nkv.	required
attn_mask	Tensor | None	Optional attention mask broadcastable to [B, Nh_q, Nq, Nkv]. Boolean masks follow SDPA semantics where True means the element participates in attention; floating masks are additive attention bias. Large-D Triton supports additive mask gradients. forward_backend='cutedsl' rejects any non-None attn_mask with NotImplementedError (no silent fallback). :param dropout_p: Dropout probability. Large-D CUDA and Triton implement SDPA-style attention dropout; unsupported backends route to SDPA, except that forward_backend='cutedsl' raises NotImplementedError for dropout_p > 0 instead of falling back.	None
is_causal	bool	When True, apply a causal attention mask so that query row r only attends to KV positions k <= r + (Nkv - Nq) (standard queries aligned to KV tail convention). Requires Nkv >= Nq. Non-causal tiles pay only one compare-and-branch per KV tile; diagonal tiles apply a per-fragment -inf mask.	False
scale	float | None	Pre-softmax scaling factor applied to QK^T. Defaults to 1 / sqrt(D) (standard attention scale) when None.	None
enable_gqa	bool	Grouped-query attention mode. Defaults to False to match SDPA exactly. When False, the large-D FFPA path requires query and key/value to have the same number of heads. Pass True to opt into GQA/MQA semantics explicitly.	False
kwargs	object	FFPA-specific extension kwargs. Supported keys are backend (str or Backend instance), forward_backend, and backward_backend. backend is a shorthand that auto-fills both forward_backend and backward_backend with the same config: backend="cutedsl" is equivalent to passing forward_backend=CuTeDSLBackend(), backward_backend=CuTeDSLBackend(). Explicit forward_backend / backward_backend take priority over backend. Any other kwarg raises TypeError.	{}

Returns:

Type	Description
Tensor	Output tensor O with layout [B, Nh_q, Nq, D], filled with the attention output softmax(scale * QK^T) V.

Raises:

Type	Description
TypeError	if query.dtype is neither fp16 nor bf16.
ValueError	if the selected backend config is invalid, if acc is not one of {'f16', 'f32'}, if bf16 activations are combined with acc='f16' (no bf16-acc mma PTX instruction exists on supported architectures), or if is_causal=True is combined with Nkv < Nq.
NotImplementedError	propagated from SDPA or FFPA backends for unsupported backend-specific combinations.

ffpa_attn.ffpa_attn_varlen_func

ffpa_attn_varlen_func(q: Tensor, k: Tensor, v: Tensor, cu_seqlens_q: Tensor, cu_seqlens_k: Tensor | None, max_seqlen_q: int, max_seqlen_k: int, *, dropout_p: float = 0.0, softmax_scale: float | None = None, causal: bool = False, enable_gqa: bool = False, return_lse: bool = False, **kwargs: object) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

FFPA variable-length attention (packed THD, FlashAttention-style).

Signature aligned with Dao-AILab flash_attn_varlen_func. Inputs are packed THD: q is [T_q, H_q, D] and k / v are [T_k, H_kv, D]. Sequence boundaries are described by cu_seqlens_q and cu_seqlens_k (int32 CUDA tensors of length B+1 starting at 0). When cu_seqlens_k is None it defaults to cu_seqlens_q (self-attention).

Only the CuTeDSL backend is supported: SM8x/SM90, large 64-aligned head dims, fp16 / bf16. Any unsupported case raises an actionable error immediately — there is no silent fallback to dense / per-sequence paths. Callers needing other shapes / backends should unpack the batch and call :func:ffpa_attn_func per sequence.

Parameters:

Name	Type	Description	Default
q	Tensor	Query tensor of shape [T_q, H_q, D].	required
k	Tensor	Key tensor of shape [T_k, H_kv, D].	required
v	Tensor	Value tensor of shape [T_k, H_kv, D].	required
cu_seqlens_q	Tensor	[B+1] int32 CUDA tensor; cu_seqlens_q[0] == 0 and cu_seqlens_q[-1] == T_q.	required
cu_seqlens_k	Tensor | None	Same convention for keys; defaults to cu_seqlens_q if None.	required
max_seqlen_q	int	Maximum per-sequence query length across the batch.	required
max_seqlen_k	int	Maximum per-sequence key length across the batch.	required
dropout_p	float	FlashAttention-compat; must be 0.0 (CuTeDSL has no dropout support).	0.0
softmax_scale	float | None	Pre-softmax scaling factor; defaults to 1 / sqrt(D).	None
causal	bool	Apply a lower-right (tail-aligned) causal mask.	False
enable_gqa	bool	Opt-in to GQA/MQA (H_q != H_kv). When False, H_q must equal H_kv.	False
return_lse	bool	When True, also return the log-sum-exp tensor of shape [H_q, T_q] in fp32 (CUDA convention).	False
kwargs	object	Most kwargs are recognized-and-rejected by :func:ffpa_attn.cute._check_supported_options — passing a non-default value for window_size, softcap, sink, attention_mask / attn_mask, block_mask, score_mod, aux_tensors, seqused_k, block_table, num_splits, or alibi_slopes raises NotImplementedError (see :raises:). Remaining kwargs are forwarded to the cutedsl wrapper for unsupported option checking.	{}

Returns:

Type	Description
Tensor | tuple[Tensor, Tensor]	out of shape [T_q, H_q, D] if return_lse=False, otherwise (out, lse).

Raises:

Type	Description
NotImplementedError	for unsupported CuTeDSL head dims or hardware, dropout_p > 0, or any non-default unsupported kwarg: window_size, softcap, sink, attention_mask / attn_mask, block_mask, score_mod, aux_tensors, seqused_k, block_table, num_splits, alibi_slopes. :raises TypeError: if cu_seqlens_* is not int32, or if dtype is not fp16/bf16.
ValueError	for shape mismatches between q/k/v, malformed cu_seqlens_*, or enable_gqa=False with H_q != H_kv.

Backend Configuration

ffpa_attn.Backend dataclass

Base backend configuration.

Attributes:

Name	Type	Description
name	str	Backend identifier (e.g. "triton", "cutedsl").
forward	bool | None	Whether this instance configures the forward pass. None (default) means "not explicitly set"; resolved by :meth:__post_init__.
backward	bool | None	Whether this instance configures the backward pass. Same None semantics as forward.

ffpa_attn.TritonBackend dataclass

Bases: Backend

Triton forward + backward backend (default).

Attributes:

Name	Type	Description
autotune	bool	Enable Triton autotuning for kernel parameters.
autotune_mode	str	Autotune search granularity ("fast" or "max").
enable_tma	bool	Enable experimental SM90+ TMA hardware acceleration.
enable_ws	bool	Force warp-specialized configs (requires enable_tma).
persist_dkdv	bool	Keep dK/dV accumulator in fp32 across backward invocations (requires enable_tma and backward=True).
split_launch	bool	Issue separate backward launches for dKdV and dQ for finer-grained scheduling.
preprocess_d_chunk	bool	Split the d_chunk preprocess across tiles.
grad_kv_storage_dtype	dtype | str | None	Optional torch.float32 / torch.float16 storage dtype for dK/dV, workaround for causal bf16 precision.
grad_q_storage_dtype	dtype | str | None	Optional torch.float32 / torch.float16 storage dtype for dQ, workaround for cross-tile atomic-add precision when USE_DKDVDQ_FUSION is enabled (also effective in non-fused path).

ffpa_attn.CUDABackend dataclass

Bases: Backend

Hand-written CUDA forward-only backend.

Attributes:

Name	Type	Description
acc	str	MMA accumulator precision ("f16" or "f32").
stages	int	Pipeline stages for the CUDA kernel (3 on Ampere/Ada, 4 on Hopper+).

ffpa_attn.CuTeDSLBackend dataclass

Bases: Backend

CuTeDSL SM90-specialized backend (Hopper only, dense 320<D<=512, fp16/bf16 training).

Attributes:

Name	Type	Description
grad_kv_storage_dtype	dtype | str | None	Optional torch.float32 / torch.float16 storage dtype for the internal SM80 dK/dV HBM buffer; final gradients are always cast back to k.dtype / v.dtype. Workaround for causal bf16 cross-tile accumulation precision (mirrors the Triton option of the same name). SM90 path currently ignores this knob and will raise if it is set.

ffpa_attn.SDPABackend dataclass

Bases: Backend

PyTorch native scaled_dot_product_attention backend.

Forward always short-circuits via :meth:FFPAAttnMeta.fallback. When used as backward_backend it delegates to :func:_efficient_attn_backward_aten.

Attributes:

Name	Type	Description
high_precision_grad	bool	When True request higher numerical precision for the backward pass (passed through to aten).