What is FlashAttention?

FlashAttention is an exact attention algorithm that computes transformer attention in small tiles inside fast on-chip memory, so the full attention matrix is never written out to the GPU's main memory. Because it returns the same result as standard attention (it is not an approximation), it works as a drop-in replacement that cuts reads and writes to HBM, runs faster, and lets models handle much longer sequences.

It was introduced in the paper "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (Dao et al., 2022).

The problem it fixes

Attention multiplies queries by keys to get a score for every pair of tokens, applies a softmax, then multiplies by the values. A straightforward implementation runs these as separate GPU kernels, and each kernel writes its result to HBM for the next one to read back. The score matrix has one entry per pair of tokens, so it grows with the square of sequence length.

For one sequence of 32,768 tokens in 16-bit, one head in one layer produces 32,768 x 32,768 x 2 bytes = 2,147,483,648 bytes, which is 2 GiB of scores. The arithmetic on those scores is light, so the GPU spends its time moving bytes rather than computing: a textbook case of being memory-bound on the roofline model.

How FlashAttention works

It splits the queries, keys and values into blocks sized to fit in on-chip SRAM, which on a GPU means shared memory and registers inside each streaming multiprocessor. For each block it computes the scores, folds them into a running softmax (using an online softmax technique that tracks a running maximum and sum, so the exact answer can be assembled tile by tile), and accumulates the output. All of this is fused into one kernel, so the intermediate scores never leave the chip. In the backward pass during training, it recomputes scores instead of storing them, trading a little extra arithmetic for far less memory traffic.

The result is that the extra memory attention needs grows linearly with sequence length instead of quadratically. The total arithmetic is not reduced; the win comes from moving fewer bytes. Later versions of FlashAttention from the same research line improve how work is split across the GPU and take advantage of newer hardware features.

One thing it does not do is shrink the KV cache. During token-by-token generation each new token still reads the whole cache from HBM, so decoding stays limited by memory bandwidth.

What it means when you pick a GPU

Memory bandwidth still matters. FlashAttention removes the quadratic score traffic, but weights and the KV cache still stream from HBM on every generated token. These are the bandwidth figures for the cards we list:

On-chip memory sets tile size. Each streaming multiprocessor has a limited amount of shared memory and registers, and that limit decides how big a tile can be and how often the kernel has to go back to HBM. Check this on a card's spec sheet rather than assuming two GPUs behave alike.

Support depends on GPU generation. The fastest FlashAttention kernels are tuned for particular GPU generations. An older card may be limited to an earlier version or fall back to a slower attention implementation, and AMD cards go through a separate software stack (see ROCm vs CUDA). Before you rent, confirm that your framework's attention backend actually supports the card.

Context length is then limited by the KV cache. With FlashAttention in place, the thing that usually runs out first on a long prompt is VRAM for the cache, not the attention matrix. Estimate how much you need with the VRAM calculator before you pick a card. Aquanode rents GPUs by the hour.