45  Efficient Attention: FlashAttention, Sparse, Linear

TipTL;DR

Self-attention costs \(O(n^2)\) in sequence length. Two attack surfaces: systems (FlashAttention, Dao 2022, an exact but IO-aware implementation) and algorithmic (sparse attention like Longformer, or linear attention like Performer, which approximate the softmax to remove the quadratic term entirely).

Depends on: Self-Attention & Multi-Head Attention

45.1 Why this matters

Self-Attention computes a score between every pair of positions, an \(n \times n\) matrix for a sequence of length \(n\). That quadratic cost is affordable at the sequence lengths the original Transformer targeted, but it becomes the binding constraint as context windows grow into the tens or hundreds of thousands of tokens. KV-Cache addressed the memory cost of attention at inference; this note addresses the compute cost of attention itself, at both training and inference.

45.2 The mechanism

The three approaches below attack the same \(n \times n\) score matrix in three different ways: compute it faster without changing it, compute only part of it, or never form it explicitly at all.

Exact (dense) Sparse Linear (factorized) every pair attends local window + global tokens φ(Q) φ(K)ᵀV no n×n matrix ever formed

45.2.1 FlashAttention: exact attention, rewritten for hardware

FlashAttention (Dao et al., 2022) computes the exact same attention output as the standard formula,

\[ \text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V, \]

but never materializes the full \(n \times n\) score matrix in slow GPU memory (HBM). It tiles the computation into blocks small enough to fit in fast on-chip memory (SRAM), computing the softmax incrementally across tiles. The result is a large wall-clock speedup and reduced memory use, with zero change to the model’s output: it is a systems optimization, not an approximation. This is why FlashAttention was adopted almost universally with no accuracy tradeoff to weigh.

45.2.2 Sparse attention: restrict the pattern

Sparse attention methods reduce the quadratic cost algorithmically, by having each position attend to only a subset of others rather than every other position:

  • Local window: each token attends only to a fixed number of nearby tokens (the diagonal band in the sparse diagram above).
  • Global tokens: a small number of designated tokens (e.g. a classification token) attend to, and are attended by, everything.

Longformer (Beltagy et al., 2020) combines both, giving near-linear cost at the price of missing a relevant long-range dependency that falls outside the fixed pattern.

45.2.3 Linear attention: approximate the softmax

Linear attention methods (e.g. Performer (Choromanski et al., 2020)) replace the softmax kernel with a decomposable approximation, so that

\[ \text{softmax}(QK^\top) \approx \phi(Q)\,\phi(K)^\top \]

for some feature map \(\phi\). Because \(\phi(K)^\top V\) can be computed once and reused for every query, the whole attention operation drops from \(O(n^2 d)\) to \(O(n d^2)\), linear in sequence length, and the \(n \times n\) matrix is never formed at all. The tradeoff is quality: an approximated softmax is not the exact attention distribution, and linear attention has historically underperformed exact attention on tasks that need precise, sharp attention to a few specific tokens.

import numpy as np

def exact_attention_cost(n, d):
    return n * n * d          # O(n^2 d): score matrix + weighted sum

def linear_attention_cost(n, d):
    return n * d * d          # O(n d^2): no n x n matrix ever formed

for n in [512, 4096, 32768]:
    d = 64
    print(f"n={n}: exact={exact_attention_cost(n,d):,}  linear={linear_attention_cost(n,d):,}")
n=512: exact=16,777,216  linear=2,097,152
n=4096: exact=1,073,741,824  linear=16,777,216
n=32768: exact=68,719,476,736  linear=134,217,728

45.2.4 What to observe

  • FlashAttention changes how attention is computed, not what it computes; sparse and linear attention change the mathematical operation itself, trading exactness for asymptotic cost.
  • The \(O(n^2)\) vs \(O(n)\) crossover only matters at long sequence lengths; the cost comparison above shows the gap widening sharply as \(n\) grows.
  • In practice, FlashAttention (a free, exact speedup) is adopted almost everywhere by default, while sparse and linear attention are used selectively, when context length demands it and the quality tradeoff is acceptable.

45.3 Application & impact

Concept here What it became Where you see it today
FlashAttention’s IO-aware tiling The default attention kernel Nearly every production Transformer implementation
Sparse attention patterns Long-document and long-context models Longformer-style local+global windows
Linear attention An alternative to softmax attention entirely Overlaps conceptually with SSM/Mamba’s linear-cost sequence mixing
NoteKey takeaway

Attention’s quadratic cost has two independent fixes: make the exact computation faster (FlashAttention, adopted almost universally) or change the computation to be sub-quadratic (sparse or linear attention, adopted selectively where long context makes the tradeoff worthwhile). The two are not mutually exclusive; a sparse or linear attention variant can still be implemented with FlashAttention-style IO awareness.

← Back to: KV-Cache