Attention From Scratch: How Transformers Replaced a Bottleneck with a Lookup

Table of Contents

Take a vanilla encoder-decoder RNN doing machine translation. It reads a source sentence token by token, updating one hidden-state vector as it goes, and by the time it hits the end-of-sentence token, that single vector — a few hundred floats — is all the decoder gets. Every word, every dependency, every bit of structure in a sentence that might be sixty tokens long has to be squeezed into that one fixed-size summary before generation even starts.

For short sentences this is fine. For long ones it isn’t, and by 2014 there were empirical numbers to prove it: translation quality degrades sharply as sentence length grows, because the fixed-length vector “does not have enough capacity to encode a long sentence with complicated structure and meaning” (Cho et al., 2014). That’s not a training problem — it’s a representational one. The architecture itself has nowhere to put the information.

Attention was invented to fix exactly that. Bahdanau, Cho & Bengio’s 2014 paper is refreshingly direct about it in the abstract: they “conjecture that the use of a fixed-length vector is a bottleneck” and propose letting the decoder “automatically (soft-)search for parts of a source sentence that are relevant to predicting a target word,” instead of compressing everything upfront.

That’s the whole idea, and it’s worth holding onto as a single sentence because everything else in this post — scaled dot products, multi-head splitting, causal masks, the FlashAttention systems machinery — is elaboration on it: replace a fixed-size, sequentially-updated summary with a direct, content-based lookup across the whole sequence. Once that clicks, the rest is implementation detail. Implementation detail is where this post is going to spend most of its time, because that’s where the actual understanding lives.

The bottleneck, precisely

It’s worth being precise here, because this is where casual explanations tend to blur two genuinely different problems together.

Problem 1: optimization. Plain RNNs are hard to train over long sequences because gradients have to flow backward through many sequential, multiplicative steps. Bengio, Simard & Frasconi showed this formally in 1994: long-term gradient components shrink exponentially with the number of steps between two positions, so learning a dependency between token 1 and token 100 via backpropagation becomes numerically close to impossible — not because the information isn’t there, but because the gradient signal that would teach the network to use it has vanished by the time it arrives. LSTMs (Hochreiter & Schmidhuber, 1997) substantially mitigate this with gating: the cell state updates additively rather than through repeated matrix multiplications, giving gradients a more direct path backward.

Problem 2: capacity. Even a perfectly-optimizable RNN — imagine gradients flow beautifully, no vanishing, no exploding — still has to funnel an entire sequence-so-far through one hidden-state vector of fixed width. That’s an architectural ceiling, not an optimization artifact. A 512-dimensional hidden state has exactly 512 floats of room, whether the sentence so far is five tokens or five hundred. LSTMs don’t touch this. Better gradient flow doesn’t create more room in the vector.

Attention’s original motivation, in Bahdanau et al., was squarely problem 2, not problem 1. Instead of forcing the decoder to work from one compressed summary, let it look back at every encoder hidden state at every decoding step, and learn — via a small feedforward scoring network, in the 2014 version — which ones matter for the token it’s generating right now. To be precise about what’s actually different here: RNNs already process the whole sequence, token by token. The problem was never visibility. It’s that everything gets compressed into one fixed-size vector along the way. Attention isn’t a “bigger visibility window” bolted onto the same recurrent mechanism — it doesn’t maintain a running summary at all. It recomputes relevance from scratch at every position. That’s a different mechanism, not a bigger version of the same one.

This also predates, and is logically separate from, the parallelism argument that motivated dropping recurrence entirely three years later in “Attention Is All You Need.” Keep the two straight: 2014 attention was bolted onto RNN/LSTM encoder-decoders to solve a representational-capacity problem. 2017 self-attention removed the RNN altogether, for a related but distinct reason — sequential computation itself.

A sharper version of “why”: path length

Vaswani et al. (2017) give a cleaner, more falsifiable version of the same argument, framed as maximum path length: the number of computation steps information has to travel through to connect any two positions in a sequence.

Layer type	Complexity per layer	Sequential operations	Maximum path length
Self-attention	O(n²·d)	O(1)	O(1)
Recurrent	O(n·d²)	O(n)	O(n)
Convolutional	O(k·n·d²)	O(1)	O(log_k(n))

n = sequence length, d = representation dimension, k = kernel width. Table 1, Vaswani et al. 2017.

In a recurrent layer, connecting token 1 to token n takes n sequential steps — the update at position 5 has to happen after position 4, which happened after position 3, and so on down the chain. That’s exactly the setup Bengio et al.’s 1994 result describes: more steps between two positions means more opportunities for the gradient connecting them to vanish.

Self-attention collapses that path to O(1). Token 1 and token n are connected by exactly one dot product, regardless of how far apart they sit in the sequence. There’s no chain to traverse — the “path” is a single matrix multiplication that touches every position at once.

The paper is also upfront about the cost of that trade, and this detail is the one that sets up everything after it: self-attention is only cheaper than a recurrent layer, per that same table, when n < d — sequence length smaller than representation dimension. That holds for a lot of NLP (word-piece and byte-pair token sequences are often shorter than the model’s hidden dimension), which made it a reasonable bet in 2017. But flip the inequality — long documents, long contexts, anything where n grows past d — and the O(n²) term stops being a rounding error and becomes the dominant cost. That’s not a footnote; it’s the entire reason FlashAttention needed to exist five years later, and we’ll get there by the end of this post.

Scaled dot-product attention, mechanically

Vaswani et al.’s structural move was to drop the encoder-decoder framing entirely and run the same kind of lookup within a single sequence — every position attending to every other position in that same sequence, including itself. That’s what “self” means in self-attention.

The formula (Eq. 1 in the paper):

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

Three matrices, three roles:

Q (query) — what is this position looking for?
K (key) — what does each position have to offer, as a label to match against?
V (value) — what does each position actually contain, as content?

For self-attention, Q, K and V are all linear projections of the same input sequence X, through three separate learned weight matrices:

Q = X W^Q
K = X W^K
V = X W^V

Nothing here is magic — W^Q, W^K, W^V are ordinary weight matrices trained by backprop like any other layer. The mechanism lives entirely in what happens once you have Q, K and V.

    flowchart TD
	    X["Input embeddings X — shape (n, d_model)"] --> WQ["Linear W^Q"]
	    X --> WK["Linear W^K"]
	    X --> WV["Linear W^V"]
	    WQ --> Q["Q — shape (n, d_k)"]
	    WK --> K["K — shape (n, d_k)"]
	    WV --> V["V — shape (n, d_v)"]
	    Q --> MM1["MatMul: Q · K^T"]
	    K --> MM1
	    MM1 --> Scale["Scale by 1/sqrt(d_k)"]
	    Scale --> Mask["Optional mask: illegal positions -> -1e9"]
	    Mask --> SM["Softmax, row-wise"]
	    SM --> W["Attention weights — (n, n), rows sum to 1"]
	    W --> MM2["MatMul: weights · V"]
	    V --> MM2
	    MM2 --> Out["Output — shape (n, d_v)"]

Here’s the same thing worked out on actual numbers, in plain numpy, for a toy 4-token sequence with 8-dimensional embeddings:

import numpy as np

np.random.seed(0)

n, d_model, d_k = 4, 8, 8            # 4 toy tokens, 8-dim embeddings
X = np.random.randn(n, d_model)       # stand-in for token embeddings

W_Q = np.random.randn(d_model, d_k) * 0.1
W_K = np.random.randn(d_model, d_k) * 0.1
W_V = np.random.randn(d_model, d_k) * 0.1

Q, K, V = X @ W_Q, X @ W_K, X @ W_V

scores = Q @ K.T / np.sqrt(d_k)       # (n, n) similarity matrix

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)   # numerical stability
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

weights = softmax(scores)             # each row sums to 1
output = weights @ V                  # (n, d_k): weighted sum of V's rows

print(weights.sum(axis=1))            # -> [1. 1. 1. 1.]
print(output.shape)                   # -> (4, 8)

weights is an n×n matrix where row i tells you how much position i attends to every position in the sequence, including itself. Every row is a genuine probability distribution — nonnegative, sums to 1 — which means output[i] is a convex combination of every value vector, weighted by that row: a weighted average where the weights are nonnegative and sum to 1, so the result is guaranteed to land inside the convex hull of V’s rows. That’s a precise way to say “weighted sum” rather than a hand-wave, and it’s why attention output magnitudes stay bounded regardless of sequence length.

That’s the whole mechanism in about ten lines. Here’s the form you’ll actually find in most transformer codebases, lifted from The Annotated Transformer (Harvard NLP), a faithful PyTorch reimplementation of the paper:

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

Same four operations as the numpy version — matmul, scale, optional mask, softmax, matmul — just batched, with a mask argument we haven’t used yet. That argument does an important job, and it’s the subject of a later section.

Why a dot product measures “relevance”

Geometrically, a dot product q·k is large when q and k point in a similar direction and have large magnitude — it rewards both alignment and size. The model has no built-in notion of relevance; it learns W^Q and W^K such that this score comes out high exactly when a key’s content is useful for a given query, because that’s what minimizes the training loss. There’s no guarantee this lands on anything a human would call “semantic,” and it’s worth being precise that a dot product is not cosine similarity: there’s no normalization step dividing out the vector norms, so magnitude matters, not just direction. Two vectors pointing the same way, one twice as long as the other, get a higher score than two identical unit vectors would. Whatever the projections learn, they learn it to serve the loss function, not to match human intuitions about meaning.

There’s also a practical reason the paper chose a plain dot product over the additive attention Bahdanau et al. used — a small feedforward network scoring each query-key pair. The two are similar in theoretical complexity, but dot-product attention reduces to a matrix multiplication, and matrix multiplication is precisely the operation GPUs (and TPUs) are built to run at full throughput — GEMM (general matrix multiply) kernels are among the most heavily optimized primitives in all of deep learning. Additive attention needs a nonlinearity evaluated per query-key pair; dot-product attention needs one GEMM call over the whole sequence at once. That’s a systems argument as much as a modeling one, and it’s the same theme that resurfaces later with FlashAttention: the shape of the math determines how fast the hardware can actually run it.

Why divide by sqrt(d_k)

This one has an exact derivation, not just a “seems to help empirically” justification. Assume the components of q and k are independent random variables with mean 0 and variance 1 — roughly true early in training, after reasonable initialization. Their dot product is

q · k = Σ_{i=1}^{d_k} q_i k_i

a sum of d_k independent, mean-zero terms. Each term q_i k_i has mean 0 (independence) and variance 1 (the variance of a product of two independent, mean-0, unit-variance variables is 1). Variance of a sum of independent variables is the sum of the variances, so:

Var(q · k) = d_k

The variance of the raw dot product grows linearly with the dimensionality of the query/key space. Left unscaled, a larger d_k means larger-magnitude scores, which pushes the softmax input into a regime where one entry dominates and the rest get squashed toward zero — softmax saturates toward a near-one-hot output, and a saturated softmax has tiny gradients almost everywhere, the same failure mode as a saturated sigmoid, just multivariate. Dividing by sqrt(d_k) renormalizes the variance back to 1, regardless of d_k, keeping the softmax input in a range where it actually has gradient to propagate.

Multi-head attention: splitting the lookup, not multiplying it

One attention head produces one weighted average per position. That’s a real limitation: a single weighted average can only express one notion of relevance at a time. If a token needs to simultaneously track, say, “what’s the subject of this verb” and “what noun does this pronoun refer to,” one softmax distribution over one set of keys has to do both jobs at once — and averaging across competing demands tends to wash both signals out, because a query can’t cleanly point in two directions inside a single similarity space.

The fix is almost embarrassingly simple: instead of one attention computation over the full d_model-dimensional space, run h independent, smaller attention computations in parallel, each with its own learned Q/K/V projections, and combine the results.

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O
head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)

The base model in the paper uses h = 8 heads, d_model = 512, with each head projecting down to d_k = d_v = d_model / h = 64. Each head gets its own 512→64 projection for Q, K and V, runs the exact same scaled dot-product mechanism from the previous section independently, in its own 64-dimensional subspace, and the eight 64-dim outputs get concatenated back to 512 dimensions and passed through one more learned projection, W^O (512×512), that mixes information across heads before the result moves to the next layer.

    flowchart LR
	    X["X — shape (n, d_model)"] --> LQ["Linear W^Q (d_model x d_model)"]
	    X --> LK["Linear W^K (d_model x d_model)"]
	    X --> LV["Linear W^V (d_model x d_model)"]
	    LQ --> SPLIT["reshape into h heads: (n, h, d_k)"]
	    LK --> SPLIT
	    LV --> SPLIT
	    SPLIT --> H1["head 1: Attention(Q1,K1,V1)"]
	    SPLIT --> H2["head 2: Attention(Q2,K2,V2)"]
	    SPLIT --> HN["... heads 3 .. h-1 ..."]
	    SPLIT --> H8["head h: Attention(Qh,Kh,Vh)"]
	    H1 --> CAT["Concat back to (n, d_model)"]
	    H2 --> CAT
	    HN --> CAT
	    H8 --> CAT
	    CAT --> WO["Linear W^O (d_model x d_model)"]
	    WO --> OUT["MultiHead output — (n, d_model)"]

One detail worth being precise about, because it trips people up: multi-head attention is not 8x the compute of single-head attention. Each head works in a 64-dimensional space instead of a 512-dimensional one, so the total compute across all 8 heads is roughly the same budget a single full-width attention computation would cost — it’s the same total compute, split across 8 independent similarity subspaces, not 8 independent full-size computations stacked on top of each other.

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super().__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)  # W^Q, W^K, W^V, W^O
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        query, key, value = [
            lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        ]
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

(This excerpt assumes import torch, torch.nn as nn, math, and a clones helper that deep-copies a module h times — see the Annotated Transformer source for the full file.) Worth staring at this for a second, because it clarifies something the formula alone doesn’t: in a real implementation there usually aren’t 8 separate small weight matrices sitting in memory. There’s one big (d_model × d_model) linear layer for Q, one for K, one for V, and “splitting into heads” is just a .view() and .transpose() — a reshape from (batch, seq_len, d_model) into (batch, heads, seq_len, d_k). Mathematically it’s equivalent to h separate smaller projections; computationally it’s one matmul sliced afterward, which is friendlier to how GPUs like to work.

Why it actually helps

Jay Alammar’s Illustrated Transformer frames the practical benefit in two parts: it lets a given position attend to several different other positions at once (useful for something like resolving what a pronoun refers to, where “relevance” might mean several different things simultaneously), and it hands the model multiple independent representation subspaces to work in — since each head has its own learned projections, head 3’s notion of “similar” can be geometrically unrelated to head 5’s.

Here’s a genuinely common overclaim worth flagging, though: it’s tempting to describe each head as learning one clean, interpretable job — “this head does syntax, that head does coreference.” The original paper makes no such claim; it only frames the benefit in terms of multiple representation subspaces. The “heads specialize in interpretable linguistic roles” framing comes from later interpretability work — Clark et al. (2019), “What Does BERT Look At?,” and Vig (2019) — and even there the findings are mixed, not tidy one-head-one-job stories. Cutting the other way, Michel, Levy & Neubig (2019), “Are Sixteen Heads Really Better than One?,” found that many heads in a trained model can simply be pruned afterward with little to no performance loss — evidence of real redundancy across heads, not clean specialization. Treat “each head learns an interpretable role” as a simplification people reach for when explaining attention, not as something the mechanism guarantees.

Causal masking: why GPT can’t look ahead

Everything above works for an encoder — a layer allowed to look at the entire sequence at once, in both directions, because the whole input is already known. That’s fine for classifying a complete sentence. It’s not fine for a language model generating text one token at a time, because during training you don’t want position 5 cheating by looking at position 6’s answer — the very thing it’s supposed to be predicting.

This is where GPT-style decoder-only transformers diverge from the original encoder-decoder architecture. GPT-1 (Radford et al., 2018) describes itself plainly: a 12-layer decoder-only transformer with masked self-attention (768-dimensional states, 12 heads), where “every token can only attend to its left context.” No encoder, no cross-attention into a separate source sequence — just the masked self-attention decoder stack from the original Transformer, used on its own. GPT-2 follows essentially the same decoder-only, masked-self-attention backbone at larger scale; the architectural description above is sourced specifically to the GPT-1 paper, which states it explicitly, rather than to GPT-2’s own paper.

The mask itself is simple once you see it: force every “illegal” score — any key position that comes after the query position — to a large negative number before the softmax, so it contributes effectively nothing to the output. Combined with shifting the target sequence by one position during training, this guarantees a prediction at position i can only depend on known outputs at positions before i.

Concretely, again from The Annotated Transformer:

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)
    return subsequent_mask == 0

torch.triu(..., diagonal=1) builds a strictly upper-triangular matrix of ones — every cell where the column index is strictly greater than the row index, i.e. every “future” position relative to each row. Inverting it with == 0 flips that into a boolean matrix that’s True exactly where a position is allowed to look: itself and everything before it. Plug that into the attention() function from earlier, and scores.masked_fill(mask == 0, -1e9) pushes every illegal position to a huge negative number before the softmax — not literal -inf, because exp(-inf) mixed with real gradient computations starts producing NaNs, but -1e9 is close enough that after exponentiating, those positions contribute effectively zero to the weighted sum.

The parallelism this buys you — and the parallelism it doesn’t

This is the part that gets mixed up constantly, so it’s worth stating precisely: causal masking is what makes training fast, not generation.

During training you already have the full ground-truth sequence — teacher forcing, meaning the model is fed the correct previous tokens rather than its own predictions at each step. With a causal mask, you can compute the loss for every position in that sequence in a single masked matrix multiplication: row i of the output only “sees” columns ≤ i, but all rows get computed simultaneously, in one forward pass, over the whole sequence at once. That’s a huge win over an RNN, which has no choice but to finish computing the hidden state at position t before it can even start on position t+1.

    sequenceDiagram
	    participant GT as Ground-truth sequence
	    participant Fwd as Training forward pass
	    participant Gen as Autoregressive generation
	
	    Note over GT,Fwd: Training: full sequence already known (teacher forcing)
	    GT->>Fwd: one masked matmul, all positions at once
	    Fwd-->>Fwd: loss computed for positions 1..n in parallel
	
	    Note over Gen: Inference: future tokens don't exist yet
	    Gen->>Gen: sample token 1
	    Gen->>Gen: sample token 2 (needs token 1)
	    Gen->>Gen: sample token 3 (needs tokens 1-2)
	    Gen->>Gen: ... strictly sequential, one token per step

Generation is a different story. At inference time you don’t have the future tokens — they don’t exist yet, because generating them is the whole point. Token t+1 depends on having actually sampled token t. No amount of masking changes that; it’s a property of autoregressive generation itself, not of how attention is computed. Causal masking parallelizes the training forward pass over a known sequence; it says nothing about making sampling parallel, because sampling has an inherent sequential dependency the mask can’t remove. KV-caching helps generation speed by avoiding recomputation of past keys and values at every step, but it doesn’t remove the token-by-token dependency — that’s really a topic for the serving/inference post later in this series, not this one.

Where this breaks down: the O(n²) wall

Go back to that complexity table. Self-attention’s per-layer cost is O(n²·d) — quadratic in sequence length. That’s the price of O(1) path length: connecting every position to every other position directly means computing n² pairwise scores. For short sequences that’s a fantastic trade. For a 32K-token context, n² means well over a billion score entries per attention layer, per head, and that stops being a rounding error.

The bigger practical problem in a naive implementation, though, isn’t even the raw FLOP count — it’s memory traffic. A naive attention kernel computes the full n×n score matrix, writes it out to the GPU’s main memory (HBM — high-bandwidth memory, physically separate from the compute cores), reads it back for the softmax, writes the softmax output back out, then reads it again for the final matmul with V. On an A100, HBM bandwidth is roughly 1.5–2.0 TB/s, while on-chip SRAM — much smaller, much faster memory sitting right next to the compute units on each streaming multiprocessor (SM) — runs at roughly 19 TB/s, call it an order of magnitude faster. The catch is capacity: about 192KB of SRAM per SM, against 40–80GB of total HBM. (Multiply 192KB by an A100’s roughly 108 SMs and you land, as a rough back-of-envelope estimate rather than a number stated directly in the paper, somewhere in the tens of megabytes of total on-chip SRAM — versus tens of gigabytes of HBM. Right order of magnitude, not a precise figure.) A naive attention kernel spends most of its wall-clock time shuffling that n×n matrix in and out of the slow, large memory tier, not doing the matrix multiplications themselves. It’s memory-bandwidth-bound, not compute-bound — a different bottleneck than the O(n²·d) FLOP count alone would suggest.

That’s precisely the problem FlashAttention (Dao, Fu, Ermon, Rudra, Ré, 2022) targets — not a smarter approximation of attention, but the exact same math, computed so the n×n matrix never gets fully materialized in slow memory.

How, mechanically

FlashAttention tiles the computation. Q gets split into blocks; K and V get split into blocks. The outer loop iterates over K/V blocks, loading each one into fast on-chip SRAM once; the inner loop iterates over Q blocks, computing partial attention output for each Q block against the current K/V block, entirely within SRAM. Block sizes are derived from how much SRAM is actually available on the device.

The genuinely clever algorithmic piece is computing a numerically stable softmax when you only ever see one block of a row at a time, instead of having the whole row in memory to normalize against. The trick — “online softmax” — is really just the standard “subtract the max before exponentiating” numerical-stability trick, applied incrementally: keep a running max m and a running sum ℓ for each query block, and every time a new K/V block arrives, rescale the previously-accumulated numerator and denominator by exp(old_max − new_max) before folding in the new block’s contribution. Done correctly, the final result is identical (up to ordinary floating-point rounding) to computing softmax over the full row at once — it’s a reformulation of the arithmetic, not an approximation of it.

Because the running statistics live in SRAM and only the small per-block output ever gets written to HBM, the full n×n matrix is never written out in full. It’s still computed — just transiently, in tiles, inside fast memory, and discarded as soon as the next tile is ready. That distinction matters: FlashAttention doesn’t avoid computing pairwise scores, it avoids ever storing the whole score matrix in slow memory at once.

In the backward pass, FlashAttention goes a step further and deliberately recomputes chunks of the forward attention matrix on the fly — from the stored Q, K, V and the saved softmax statistics — instead of keeping the full n×n matrix from the forward pass around to reuse. That’s extra arithmetic, traded directly for less memory traffic, in the same spirit as gradient checkpointing, and it’s worth it because on modern GPUs moving data is far more expensive than computing on it.

Which is the single easiest thing to get backwards about this whole paper, so it’s worth stating as plainly as possible: FlashAttention does not reduce the number of FLOPs, and it does not change attention’s O(n²) complexity class. If anything, it does somewhat more total arithmetic than a naive implementation, because of that backward-pass recomputation. The entire speedup comes from moving less data through slow memory, not from doing less math. And it computes exact attention — mathematically identical to standard softmax attention, not an approximation. The same paper separately proposes a block-sparse variant that is approximate; that’s a different algorithm bolted on afterward, not the headline result.

What it actually bought, in verified numbers

Straight from the paper’s abstract:

3× speedup training GPT-2 at sequence length 1K
15% end-to-end wall-clock speedup training BERT-large (sequence length 512) against the MLPerf 1.1 training speed record
2.4× speedup on Long-Range Arena (sequence lengths 1K–4K)
Longer context, made affordable, translated into better models, not just faster ones: 0.7 better perplexity on GPT-2 and a 6.4-point lift on a long-document classification task, plus the first Transformers to beat chance on the Long Range Arena’s Path-X (61.4% accuracy at sequence length 16K) and Path-256 (63.1% at 64K) tasks — synthetic long-range-dependency benchmarks nothing had cracked before
Memory scales linearly with sequence length instead of quadratically. The authors’ own repository (not the paper itself) reports roughly 10× memory savings at sequence length 2K and 20× at 4K
An IO-complexity theorem: standard attention needs Θ(nd + n²) HBM accesses; FlashAttention needs Θ(n²d²/M) accesses, where M is SRAM size — and the paper proves this is asymptotically optimal across the range of SRAM sizes real GPUs have. (You may see other multipliers — a 7.6× GPT-2 figure circulates in some secondary write-ups — but the paper’s own abstract states 3× for GPT-2 at sequence length 1K, and that’s the number worth citing.)

FlashAttention-2 (Dao, 2023) pushed further on GPU work partitioning and parallelism — roughly 2× faster than the original FlashAttention, up to 230 TFLOPs/s on an A100 for the attention operation itself. Hopper-generation GPUs have their own successor (FlashAttention-3, built around H100-specific asynchrony and FP8) — worth knowing it exists, though I’m not citing specific multipliers for it here since I haven’t verified them directly against the paper.

What this bought, and what’s next

Zoom out and the arc looks like this: RNNs bottleneck long-range dependencies because everything has to pass through one fixed-size, sequentially-updated vector. Attention fixes that by replacing the bottleneck with a direct, content-based lookup — Q, K, V, a scaled dot product, a softmax, a weighted sum — collapsing the path between any two positions to O(1). Multi-head splits that lookup into several independent subspaces so the model isn’t limited to one notion of relevance at a time. Causal masking adapts the same mechanism to autoregressive generation and buys a huge parallelism win at training time, though not at inference time. And the O(n²) cost that direct lookup introduces turns out to be a memory-bandwidth problem in disguise — exactly what FlashAttention was built to solve, without changing the underlying math at all.

It worked empirically from the start: the original paper’s “big” Transformer trained in 3.5 days on 8 P100 GPUs and posted a new single-model state of the art on WMT 2014 English-French translation, 41.8 BLEU (a standard machine-translation quality metric) — “a small fraction of the training costs of the best models from the literature,” in the paper’s own words.

This post is the mechanism. It’s deliberately not the whole story: the next post in this series gets into the theory side — what determines how well a model of a given size, trained on a given amount of data, actually performs — and the one after that covers what it takes to serve a transformer in production, where KV-caching, batching, and quantization end up mattering just as much as the attention formula itself.