Skip to content

Backends Overview

FFPA ships four interchangeable attention backends behind a single entry point, ffpa_attn_func:

Backend Config class Forward Backward Where it runs
SDPA SDPABackend ✔ (PyTorch native) ✔ all devices
CUDA CUDABackend ✔ ✖ fp16/bf16 + FP8/FP4 C++ kernels
Triton TritonBackend ✔ ✔ NVIDIA sm_80+ and AMD ROCm
CuTe-DSL CuTeDSLBackend ✔ ✔ sm_80 (generic), sm_90a, sm_100a

Choosing a backend

from ffpa_attn import ffpa_attn_func, CUDABackend, TritonBackend, CuTeDSLBackend

# FlashAttention-style training shape (D=512, needs backward): Triton default.
o = ffpa_attn_func(q, k, v)

# Large-D SM90/SM100 dense training: CuTe-DSL (D=512), see cutedsl.md.
o = ffpa_attn_func(q, k, v, backend=CuTeDSLBackend())

# Forward-only inference with the C++ CUDA kernels (fp16/bf16 or FP8/FP4).
o = ffpa_attn_func(q, k, v, forward_backend=CUDABackend(enable_fp8=True))

The backend= shorthand fills both forward_backend and backward_backend; explicit forward_backend / backward_backend kwargs take priority. Forward and backward may use different backends, e.g. the CUDA forward kernels with a Triton backward.

Capability matrix

Backend Arch Fwd Bwd Dtype Headdim Causal GQA / MQA attn_mask Dropout Varlen (THD) NHD zero-copy Autotune
SDPA sm>=75 ✔ ✔ fp16/bf16/fp32 all ✔ ✔ ✔ ✔ ✖ ✖ ✖
CUDA (fp16/bf16) sm>=80 ✔ ✖ fp16/bf16 320~1024 ✔ ✔ ✔ ✔ ✖ ✔ ✖
CUDA FP8 sm_89 ✔ ✖ fp16/bf16 in 64~224 ✔ ✔ ✔ ✖ ✖ ✔ ✖
CUDA FP8 sm_120{a,f} ✔ ✖ fp16/bf16 in 64~1024 ✔ ✔ ✔ ✖ ✖ ✔ ✖
CUDA FP4 sm_120{a,f} ✔ ✖ fp16/bf16 in 64~512 ✔ ✔ ✔ ✖ ✖ ✔ ✖
Triton sm>=80, ROCm ✔ ✔ fp16/bf16 320~1024 ✔ ✔ ✔ (bias grad ✔) ✔ ✖ ✖ ✔
CuTe-DSL sm>=80 ✔ ✔ fp16/bf16 320~1024 ✔ ✔ ✖ ✖ ✔ ✖ ✖
CuTe-DSL sm_90a ✔ ✔ fp16/bf16 320~512 ✔ ✔ ✖ ✖ ✔ ✖ ✖
CuTe-DSL sm_100a ✔ ✔ fp16/bf16 512 ✔ ✔ ✖ ✖ ✔ ✖ ✖

Notes:

  • Headdim ranges above are the FFPA-tuned ranges. Small-D (D<=256) also runs every backend by default; set the matching FFPA_{CUDA,TRITON,CUTE}_ALLOW_SMALL_D=0 to restore the SDPA delegation (see Automatic SDPA fallback).
  • The CUDA FP8 / FP4 rows take fp16/bf16 activations and quantize Q/K/V inside the kernel; see the FP8 / FP4 guide.
  • NHD zero-copy means the kernel consumes diffusers-style [B, N, H, D] (or fused-QKV chunk) views without materializing a packed copy — a CUDA-backend inference-path feature only.

Automatic SDPA fallback

ffpa_attn_func keeps the SDPA-aligned contract: instead of failing, it falls back to torch._C._nn.scaled_dot_product_attention when the requested shape/device is outside a backend's contract:

Condition Behavior
D <= 256 (small-D) the selected backend's kernels by default; SDPA only when FFPA_CUDA_ALLOW_SMALL_D=0 / FFPA_TRITON_ALLOW_SMALL_D=0 / FFPA_CUTE_ALLOW_SMALL_D=0 is set explicitly
D > 1024 SDPA
8 <= Nq < 512 or Nkv < 512 (short sequences) SDPA (Nq in [1, 8) still uses the native decode path)
CuTe-DSL hardware/shape mismatch SDPA
forward_backend='sdpa' always SDPA

Everything else (unsupported dtype, attn_mask under CuTe-DSL, dropout_p > 0 under CuTe-DSL, invalid backend combinations) raises NotImplementedError / TypeError / ValueError — there is no silent fallback for those. Set FFPA_FORCE_NO_SDPA_FALLBACK=1 in tests to assert the backend kernels instead of the fallback.

Backends in this section

  • CUDA Backend — CUDABackend: native / TMA / CuTe fp16 C++ kernels plus FP8 and NVFP4 attention, forward only.
  • Triton Backend — TritonBackend: the default forward + backward backend, autotune and persistent tuned configs.
  • CuTe-DSL Backend — CuTeDSLBackend: Python CuTe kernels for sm_80/sm_90a/sm_100a, including the tcgen05 2-CTA D=512 path and the varlen (packed THD) API.
  • FP8 / FP4 Attention — quantized attention usage, knobs and precision notes for the CUDA backend.