Abstract
FlashAttention [1] is an IO-aware algorithm for computing exact self-attention [2] with substantially reduced memory traffic. It reduces the memory and memory-bandwidth cost of attention, not its arithmetic (FLOP) complexity: the result is exact, and the wall-clock speedups come from never materializing the full attention matrix in high-bandwidth memory. This note sets out what that buys for learning distributed representations, and draws the line — precisely — between kernel fusion (exact) and the block-sparse variant (an approximation).
Note
Standard self-attention forms an n × n score matrix for a length-n sequence, which is
quadratic in both time and memory [2]. For long sequences the dominant practical
cost is often not raw arithmetic but moving that large intermediate matrix between GPU memory
tiers. FlashAttention [1] addresses exactly this: it tiles the computation and fuses
the score, softmax, and value-aggregation steps so that the full n × n matrix is never
written to high-bandwidth memory, reducing memory use from quadratic to linear in n while
computing the same result as ordinary attention. The paper reports measured wall-clock
speedups and training on longer contexts than a naive implementation permits.
For distributed representation learning, the consequence is direct: longer contexts become affordable at a fixed memory budget, so representations can be learned from larger spans — whole documents, long audio — without approximating the attention pattern. That is the benefit to claim, anchored to [1].
Two distinctions carry the argument and are worth stating exactly:
- Exact, not sub-quadratic in FLOPs. Base FlashAttention preserves the asymptotic arithmetic cost of exact attention; it lowers memory traffic and thus wall-clock time [1]. Sub-quadratic behavior comes only from the separate block-sparse variant, which changes the attention pattern — an approximation, and labeled as one.
- Kernel fusion and block-sparsity are distinct. Fusion (avoiding the materialized matrix) preserves exactness; block-sparsity trades exactness for further savings. They are different design points and should not be run together.
Evidence & Scope
Every quantitative statement here is the measured result of the FlashAttention paper [1] over the transformer backbone [2]: exact FlashAttention is a memory/IO optimization that keeps the quadratic FLOP count of attention, and its speedups are wall-clock, from reduced memory traffic. Where sub-quadratic compute is the requirement, the design point shifts to approximate-attention methods — sparse or low-rank families — which trade exactness for scale. This note is the map of that terrain, drawn to the boundary between the two.
References
- Tri Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2205.14135.
- Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.