flagos-ai/FlagAttention

FlagAttention – Fast, memory‑efficient attention kernels

What it is – FlagAttention is a Python library that provides a collection of custom attention operators written in Triton (a language for writing high‑performance GPU kernels). The kernels implement the same functionality as the popular FlashAttention algorithm but also add many specialised variants such as sparse, paged‑cache, piece‑wise, quantised, and recurrent attention. They are intended for training and inference of large language models or other transformer‑style networks on a variety of accelerators (NVIDIA GPUs, AMD, Ascend, etc.).

---

Key operator families

Family	Typical use‑case	API entry point	Gradient support	Maturity
FlashAttention	Standard dense attention (scaled‑dot‑product)	flag_attn.flash_attention	Forward + backward	Stable
Piecewise Attention	Switch between two Q/K pairs based on token distance (e.g., NLPE)	flag_attn.piecewise_attention	Forward + backward	Stable
Split‑KV FlashAttention	Long‑KV decoding, low query parallelism	flag_attn.flash_attention_split_kv	Forward only	Alpha
Paged Attention	Single‑token queries over a paged KV cache (common in decoder‑only models)	flag_attn.paged_attention	Forward only	Alpha
MiniMax M3 Sparse Attention	Block‑sparse prefill/decode with 128‑token pages	flag_attn.minimax_m3_*	Inference only	Alpha
Chunk GLA / GDN2 / KDA	Recurrent gated‑linear or delta‑rule attention for long sequences	flag_attn.chunk_*	Forward (+ backward for GLA)	Alpha
SageAttention	Per‑block INT8 quantised Q/K with optional mask	flag_attn.sage_attention.forward	Forward only	Alpha
Parallel NSA	Native sparse attention + compression	flag_attn.parallel_nsa.parallel_nsa	Forward + backward	Alpha

---

Installation (quick guide)

# Clone the repo
git clone https://github.com/flagos-ai/FlagAttention.git
cd FlagAttention

# Install the editable package with Triton and test extras
pip install -e "[triton,test]"

You must have a compatible PyTorch build for your accelerator already installed (e.g., torch==2.x+cu118). The package deliberately does not depend on a specific PyTorch wheel because the correct build varies by GPU vendor.

---

Minimal usage example – FlashAttention

import torch
from flag_attn import flash_attention

B, Hq, Hkv, M, N, D = 2, 16, 4, 2048, 4096, 128
q = torch.randn(B, Hq, M, D, device='cuda', dtype=torch.float16, requires_grad=True)
k = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)
v = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)

out = flash_attention(q, k, v, causal=True)
out.sum().backward()   # gradient flow works

The function signature mirrors the original FlashAttention API and adds optional flags for dropout, returning the log‑normalizer, total attention mass, and Philox RNG state.

---

When to use which operator

Scenario	Recommended operator
Standard dense training or inference on a single GPU	FlashAttention (stable)
Need to switch attention matrices after a distance threshold (e.g., positional‑embedding tricks)	Piecewise Attention (stable)
Very long KV caches where queries are few (decoding)	Split‑KV FlashAttention (alpha)
Decoder‑only models that keep a paged KV cache (e.g., LLM serving)	Paged Attention (alpha)
Block‑sparse models such as vLLM‑style token‑sparse decoding	MiniMax M3 (alpha)
Recurrent linear attention for streaming or very long sequences	Chunk GLA / GDN2 / KDA (alpha)
Want INT8 quantisation per block to save memory bandwidth	SageAttention (alpha)
Need native sparse attention with optional compression	Parallel NSA (alpha)

---

Maturity & testing

• The repository tracks operator maturity in conf/operators.yaml. Only FlashAttention and Piecewise Attention are marked stable; all others are alpha and may change API signatures.
• A full pytest suite (pytest tests) validates numerical accuracy against FP32 PyTorch references. Accuracy tolerances are roughly twice the FP32 error, with a small absolute margin.
• Benchmarks are provided under benchmark/ and can be run via the helper script tools/run_tests.py or directly (e.g., python benchmark/flash_benchmark.py). Results are emitted as JSON compatible with the FlagGems platform.

---

Supported hardware & backends

• Primary: NVIDIA CUDA GPUs (via Triton). The library also contains vendor‑specific backends for AMD, Ascend, Enflame, MetaX, MThreads, Intel, etc. Backend selection is automatic based on the detected PyTorch device, but not every operator is implemented for every vendor – consult conf/operators.yaml and the src/flag_attn/runtime/backend directory for details.
• CPU: You can import the package and inspect metadata, but the kernels cannot execute without a GPU‑compatible runtime.

---

Where to find more

• Documentation – The README itself contains the full API reference, installation notes, and example scripts under examples/.
• Tests & benchmarks – Look at tests/ for concrete usage patterns and benchmark/ for performance numbers.
• Vendor‑specific extensions – See src/flag_attn/runtime/backend/_<vendor> for low‑level implementations.

---

TL;DR

FlagAttention is a real, production‑grade library that gives researchers and engineers fast, memory‑saving attention kernels for a wide range of transformer workloads. Its stable FlashAttention and PiecewiseAttention APIs can replace the default PyTorch attention in training loops, while the alpha operators provide experimental features such as paged caches, block‑sparse decoding, quantised kernels, and recurrent linear attention for advanced LLM serving and research.