flagos-ai/FlagAttention

A collection of memory efficient attention operators implemented in the Triton language.

What it solves

FlagAttention provides memory-efficient attention operators that are easier to modify and customize than standard CUDA-based implementations like FlashAttention. It addresses the need for non-standard attention mechanisms in language modeling where project-specific customizations to attention score computation, KV-cache layout, or inference paths are required.

How it works

Implemented in the Triton language, the project uses tiling and re-computation to avoid materializing the full attention score matrix, which significantly reduces memory footprint and traffic. It offers several specialized operators:

  • flash_attention: A Triton-based implementation of FlashAttention v2 supporting MQA/GQA and dropout.
  • piecewise_attention: An extension for Non-Linearized position embedding (NLPE) that uses two sets of queries and keys to compute attention scores based on a distance threshold.
  • flash_attention_split_kv: A split-KV flash decoding operator designed for long KV sequences.
  • paged_attention: An operator for paged KV-cache management during inference.

Who it’s for

It is designed for AI researchers and developers who need high-performance attention operators but require more flexibility than what is provided by out-of-the-box libraries, or those who want to implement custom transformations to attention scores.

Highlights

  • Triton-based: Easier to understand and modify than raw CUDA code.
  • Memory Efficient: Reduces memory traffic and footprint using tiling and re-computation.
  • Customizable: Specifically supports non-standard operators like piecewise attention for NLPE.
  • Hardware Support: Tested on Nvidia Ampere GPUs (A100, RTX-3090) and Iluvatar CoreX GPUs.
  • Comprehensive Features: Supports float16, bfloat16, causal/non-causal modes, and MQA/GQA.

Related

  • Project
    Dao-AILab/flash-attention

    A fast and memory-efficient implementation of the attention mechanism for Transformers, reducing memory IO to accelerate LLM training and inference.

  • Dispatch
    vLLM Triton Attention Backend Deep Dive

    vLLM has implemented a portable Triton-based attention backend that achieves state-of-the-art performance across NVIDIA, AMD, and Intel GPUs using a single source code implementation.

  • Project
    HKUSTDial/flash-sparse-attention

    A high-performance trainable sparse attention implementation that combines Flash Attention's memory efficiency with sparse computation for handling extremely long sequences in Transformers.

  • Project
    xlite-dev/ffpa-attn

    A fast and memory-efficient exact attention implementation for large head dimensions (up to 1024), providing significant speedups over PyTorch SDPA through SRAM and register optimizations.

  • Project
    microsoft/vattention

    A memory manager for LLM KV-caches that uses CUDA virtual memory APIs to enable dynamic allocation while maintaining virtual contiguity for unmodified attention kernels.