xlite-dev/ffpa-attn

What it solves

FFPA addresses the performance and memory bottlenecks associated with the attention mechanism in transformers when using large head dimensions (specifically those greater than 256, up to 1024). Standard implementations like PyTorch SDPA often struggle with SRAM limits and register pressure as the head dimension increases, leading to slower execution and potential register spilling.

How it works

FFPA extends FlashAttention by implementing two primary technical optimizations:

• Split-D: This technique chunks the head dimension axis to keep SRAM usage constant ($O(1)$ complexity), preventing the SRAM bottleneck that occurs when $Q$ is too large to fit in memory.
• TiledMMA: This uses a specific layout (M4N2) to reduce register pressure. By splitting the $N$ axis, it halves the register cost for the $PV$ matrix multiplication, allowing the system to support head dimensions up to 1024 without spilling.

The project provides multiple backends including CUDA, Triton, and CuTe-DSL, and supports various precision formats including BF16, FP8, and FP4.

Who it’s for

Developers and researchers working with large-scale transformer models (such as Gemma4-31B or Diffusion models like FLUX.1-dev) that utilize large head dimensions and require high-throughput training or inference on NVIDIA GPUs (including Blackwell and Hopper architectures).

Highlights

• Large Headdim Support: Supports head dimensions from 320 up to 1024.
• Cuda-Optimized: Achieves 1.5x to 15x speedup over PyTorch SDPA depending on the hardware and backend.
• Multi-Precision: Experimental support for FP8 and FP4 attention for high-speed inference.
• Broad Compatibility: Supports self-attention, GQA, MQA, cross-attention, causal masking, and dropout.
• Hardware Support: Compatible with NVIDIA GPUs (sm_80 through sm_120) and AMD ROCm/HIP GPUs via the Triton backend.