Efficient Transformers: Linear Attention and Flash Attention

The 16,000-token wall

A legal-tech team building a High Court judgment-search tool wants to feed an entire judgment — often 15,000 to 40,000 tokens once you count facts, arguments, citations, and the order — straight into an open-weight transformer as context, rather than chunking it and risking a summary that misses the operative paragraph. The prototype works beautifully on short judgments. Past roughly 16,000 tokens the GPU throws an out-of-memory error, and the team cannot explain why: they added more tokens, not more parameters. The model did not get bigger. The input got longer, and somewhere inside self-attention that alone was enough to exhaust 40 GB of GPU memory.

The culprit is a single matrix. Self-attention compares every token to every other token, so the intermediate "who attends to whom" table has one entry per pair of tokens — n² entries for a sequence of length n. Double the judgment length and that table quadruples. This chapter derives exactly why that happens, then covers the two dominant fixes used in production LLM systems today: linear attention, which changes the algorithm so the n×n table is never computed at all, and FlashAttention, which keeps the exact same math but never lets that table sit in slow GPU memory. They solve related but distinct problems, and confusing them is the single most common mistake students make about this topic — addressed explicitly later in this chapter.

Why self-attention costs O(n²)

Recall the scaled dot-product attention from the G11 transformer-internals unit. Given query, key, and value matrices Q, K, V ∈ ℝ^n×d (n tokens, d-dimensional head), the output is

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

Trace the shapes. QK^T is (n×d)·(d×n) = an n×n matrix S of raw similarity scores — computing it costs O(n²d) multiply-adds. softmax(S) normalizes each row independently, another O(n²) elementwise pass. Multiplying that n×n weight matrix by V (n×d) costs another O(n²d). Total compute: O(n²d). Total memory for the intermediate score/weight matrix: O(n²), independent of d.

For an 8,000-token document this table has 64 million entries. For a 128,000-token context — the scale modern long-context LLMs advertise — it has 128,000² = 16,384,000,000 entries. Stored in fp16 (2 bytes each), that is 32.77 GB, for a single attention head, in a single layer, for a single sequence in the batch. A model with 32 heads per layer would need over a terabyte just to hold that one matrix for one layer if it materialized every head's scores simultaneously — impossible on any current GPU. This is the arithmetic behind the OOM error: nobody wrote a bug: the algorithm, as textbooks present it, is intrinsically quadratic in memory.

Linear attention: reordering the matmuls, not shrinking them

The key algebraic fact linear attention exploits is that matrix multiplication is associative: (QK^T)V and Q(K^TV) are mathematically the same n×d result, but they are computed in a different order. The first order — the standard one — builds the expensive n×n matrix first. The second order builds K^TV first, which is only d×d, and only then multiplies by Q. If you can avoid the softmax's nonlinearity (which prevents this reordering, since softmax cannot be pulled apart across the matrix product), the whole computation collapses from O(n²d) to O(nd²) time and from O(n²) to O(nd + d²) memory — linear in sequence length.

Katharopoulos, Vyas, Pappas & Fleuret ("Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention," ICML 2020) formalized this by replacing softmax(QK^T) with a kernel feature map φ applied to each query and key row, so that similarity becomes φ(Q_i)·φ(K_j) instead of exp(Q_i·K_j/√d). The attention output for query i becomes

Attention_i = [ φ(Q_i) · Σ_j φ(K_j) V_j^T ] / [ φ(Q_i) · Σ_j φ(K_j) ]

The sums Σ_jφ(K_j)V_j^T (a d×d matrix) and Σ_jφ(K_j) (a d-vector) are computed once, shared across every query, then each query does one cheap d-dimensional lookup. φ must produce non-negative outputs for this to behave like a valid attention distribution — Katharopoulos et al. use φ(x) = elu(x) + 1; Choromanski et al. ("Rethinking Attention with Performers," ICLR 2021) instead use random Fourier features (FAVOR+) that provably approximate the softmax kernel exp(Q·K) itself, trading a small amount of approximation error for the same linear-time guarantee with a tighter fidelity bound.

Worked example: the reassociation trick, by hand

Take a tiny 3-token, 2-dimensional case with a simple non-negative feature map φ(x) = x (valid here because every Q, K entry below is non-negative, so all dot products stay non-negative):

Q = [[1,0], [0,1], [1,1]]
K = [[2,1], [1,3], [1,1]]
V = [[1,0], [0,1], [1,1]]

Standard path. S = QK^T gives rows [2,1,1], [1,3,1], [3,4,2] (e.g. row 3 = Q₃·K₁=1·2+1·1=3, Q₃·K₂=1·1+1·3=4, Q₃·K₃=1·1+1·1=2). Row-normalizing (row sums 4, 5, 9) gives weights [0.5,0.25,0.25], [0.2,0.6,0.2], [1/3,4/9,2/9]. Multiplying by V: row 1 output = 0.5(1,0)+0.25(0,1)+0.25(1,1) = (0.75, 0.5); row 2 = (0.4, 0.8); row 3 = (5/9, 2/3) ≈ (0.556, 0.667).

Linear path. Compute K^TV once: K^T = [[2,1,1],[1,3,1]], and K^TV = [[3,2],[2,4]] (a 2×2 matrix — check row 1: 2·1+1·0+1·1=3, 2·0+1·1+1·1=2). Also compute the column sums of K once: k_sum = (4, 5). Now for each query: numerator = Q_i·(K^TV), denominator = Q_i·k_sum.

Q₁=(1,0): numerator=(3,2), denominator=1·4+0·5=4 → output=(0.75, 0.5). Q₂=(0,1): numerator=(2,4), denominator=5 → (0.4, 0.8). Q₃=(1,1): numerator=(5,6), denominator=9 → (5/9, 6/9) ≈ (0.556, 0.667). Every output matches the standard path exactly, because associativity guarantees Q(K^TV) = (QK^T)V and Q(K^Tone) = (QK^T)one — this is not a coincidence of the numbers chosen, it is a law of matrix algebra.

The saving is invisible at n=3 but decisive at n=40,000: the standard path builds a 40,000×40,000 matrix; the linear path never builds anything bigger than d×d (say 64×64) regardless of n.

import numpy as np

Q = np.array([[1,0],[0,1],[1,1]], dtype=float)
K = np.array([[2,1],[1,3],[1,1]], dtype=float)
V = np.array([[1,0],[0,1],[1,1]], dtype=float)

# Standard path: O(n^2 d)
S = Q @ K.T
weights = S / S.sum(axis=1, keepdims=True)
out_standard = weights @ V

# Linear path: O(n d^2), same result by associativity
KV = K.T @ V              # (d, d) -- computed once
k_sum = K.sum(axis=0)     # (d,)   -- computed once
numerator = Q @ KV
denominator = Q @ k_sum
out_linear = numerator / denominator[:, None]

print(out_standard)
# [[0.75       0.5       ]
#  [0.4        0.8       ]
#  [0.55555556 0.66666667]]
print(out_linear)
# identical, up to floating-point rounding

The recurrent view: attention as an RNN

The "Transformers are RNNs" title is not a metaphor. Define a running state S_t = Σ_j≤tφ(K_j)V_j^T and normalizer z_t = Σ_j≤tφ(K_j). During causal (autoregressive) generation, each new token updates the state incrementally: S_t = S_t−1 + φ(K_t)V_t^T, and the output at step t is φ(Q_t)·S_t / φ(Q_t)·z_t. Because S_t is always a fixed-size d×d matrix, generation needs O(1) memory per step, exactly like an RNN's hidden state — no growing key-value cache. This is the structural reason linear attention (and its modern descendants, state-space models like Mamba) are attractive for very long, streaming generation: the "memory" of the whole past is compressed into one fixed-size matrix instead of kept as a list that grows with every token.

FlashAttention: same math, IO-aware execution

FlashAttention (Dao, Fu, Ermon, Rudra & Ré, "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," NeurIPS 2022) attacks a completely different bottleneck and makes no approximation at all — it computes exact softmax attention, bit-for-bit the same output as the standard formula. Its target is the GPU memory hierarchy. An A100 GPU has two very different kinds of memory: HBM (the large, off-chip memory holding the whole model and activations — tens of gigabytes, but only ~1.5–2 TB/s bandwidth) and SRAM (small on-chip memory local to each streaming multiprocessor — tens of megabytes total, but roughly an order of magnitude faster to access). The naive attention implementation reads Q, K, V from HBM, writes the full n×n score matrix S back to HBM, reads it again for softmax, writes the normalized weights back to HBM, reads them again to multiply by V. Every one of those n² elements crosses the slow HBM boundary multiple times. GPUs have far more raw compute (FLOPs) than memory bandwidth, so for long sequences this implementation is memory-bound, not compute-bound — the GPU spends most of its time waiting on data movement, not arithmetic.

FlashAttention's fix is tiling plus online softmax: split Q, K, V into small blocks that fit in SRAM, stream key/value blocks in one at a time, and update a running maximum m, running normalizer l, and running (unnormalized) output o incrementally, so the full n×n matrix is never written to HBM — only small tiles ever exist, briefly, on-chip. The trick that makes this exact (not approximate) is rescaling: because softmax needs a global max for numerical stability, and you do not know the global max until you have seen every block, each time a new block raises the running max you must retroactively rescale everything accumulated so far by exp(old_max − new_max) before adding the new block's contribution.

Worked example: online softmax, traced by hand

Take four keys with raw (pre-softmax) scores against one query: s = [2, 0, 1, 3], with matching values v = [30, 40, 10, 20]. Standard (batch) softmax: global max = 3; exp(s−3) = [e⁻¹, e⁻³, e⁻², e⁰] = [0.367879, 0.049787, 0.135335, 1]; sum l = 1.553002; weighted sum o = 0.367879·30 + 0.049787·40 + 0.135335·10 + 1·20 = 11.03638 + 1.99148 + 1.35335 + 20 = 34.38122; output = 34.38122 / 1.553002 ≈ 22.1385.

Now process the same four keys in two streamed blocks of size 2, deliberately ordered so the global maximum arrives in the second block — this forces the rescale step to actually fire, which is the part of the algorithm worth checking by hand. Block 1 = keys with scores [2, 0], values [30, 40]. Block 2 = keys with scores [1, 3], values [10, 20].

Block 1. Local max = 2, so running max m₁ = 2. exp(2−2)=1, exp(0−2)=0.135335. l₁ = 1 + 0.135335 = 1.135335. o₁ = 1·30 + 0.135335·40 = 30 + 5.41341 = 35.41341.

Block 2. Local max = 3, so new running max m₂ = max(2, 3) = 3. Correction factor = exp(m₁ − m₂) = exp(−1) = 0.367879. Rescale the block-1 accumulators before adding anything new: l₁_corrected = 1.135335 × 0.367879 = 0.417667; o₁_corrected = 35.41341 × 0.367879 = 13.02787. New block's contribution: exp(1−3)=0.135335, exp(3−3)=1, so l_block=1.135335, o_block = 0.135335·10 + 1·20 = 1.35335+20 = 21.35335. Final: l₂ = 0.417667 + 1.135335 = 1.553002; o₂ = 13.02787 + 21.35335 = 34.38122.

Output = 34.38122 / 1.553002 ≈ 22.1385 — identical to the batch computation, to five decimal places, and the intermediate l and o values matched the full-batch sum exactly at every step. The full n×n (here, 1×4) matrix of exponentials was never held all at once; only two elements existed in memory at any moment. That is the entire mechanism of FlashAttention's forward pass, generalized from one query to n queries and tiled in both dimensions.

import numpy as np

def flash_attention_1d(scores, values, block_size):
    n = len(scores)
    m = -np.inf   # running max
    l = 0.0       # running softmax denominator
    o = 0.0       # running unnormalized weighted output
    for start in range(0, n, block_size):
        block_scores = scores[start:start + block_size]
        block_values = values[start:start + block_size]
        m_new = max(m, block_scores.max())
        correction = np.exp(m - m_new) if m != -np.inf else 0.0
        p = np.exp(block_scores - m_new)
        l = l * correction + p.sum()
        o = o * correction + (p * block_values).sum()
        m = m_new
    return o / l

scores = np.array([2.0, 0.0, 1.0, 3.0])   # block 1 = [2,0], block 2 = [1,3]
values = np.array([30.0, 40.0, 10.0, 20.0])
print(flash_attention_1d(scores, values, block_size=2))
# 22.1385...  (matches the full-batch softmax exactly)

FlashAttention-2 (Dao, 2023) keeps this same IO-aware algorithm and restructures the parallelism — more of the work is split across thread blocks and warps, and fewer non-matmul operations sit on the critical path — to close the remaining gap to a GPU's theoretical peak throughput; it does not change the asymptotic complexity established here.

Seeing both mechanisms at once

The misconception: "FlashAttention makes attention subquadratic"

Students who learn linear attention and FlashAttention back to back almost always merge them into one idea: "both make attention faster than O(n²)." That is true only for linear attention. FlashAttention performs exactly O(n²d) floating-point operations — the same count as the standard formula, because it computes the exact same softmax, over the exact same n² query-key pairs. Nothing about the arithmetic shrinks. What shrinks is data movement: the number of times bytes cross between HBM and SRAM. Doubling the sequence length still quadruples FlashAttention's wall-clock cost, in the same asymptotic sense as standard attention (with a much smaller constant factor and none of the O(n²) memory blow-up); it does not make an 8× longer context "only" 8× more expensive the way true linear attention does. If you need genuinely sub-quadratic scaling for extreme context lengths — a million-token document, not 40,000 — FlashAttention alone does not get you there; you need an algorithmic change like linear attention, or a hybrid architecture. If you need exact, unmodified softmax attention that simply runs efficiently on the GPU you already have, FlashAttention is the right tool and linear attention is unnecessary risk (its kernel approximation can measurably hurt quality on tasks requiring sharp, needle-in-a-haystack retrieval over long context, precisely because it replaces softmax's sharp exponential weighting with a smoother kernel similarity). These are answers to two different questions — "can I avoid computing n² comparisons at all?" versus "can I compute them without drowning the GPU in memory traffic?" — and production systems increasingly use FlashAttention as the default kernel for standard attention layers, reserving linear/hybrid attention for architectures deliberately designed around it from the start.

Active recall

Q1. Why is materializing the n×n score matrix the actual bottleneck for long sequences, rather than "attention is slow" in a vague sense?

Q2. True or false: "FlashAttention makes transformers sub-quadratic, the same way linear attention does." Justify your answer.

Q3. For n = 65,536 tokens (2¹⁶), d_head = 128, fp16 storage, compute the memory needed to hold the full attention score matrix for a single head in a single layer. Compare it to a typical 80 GB GPU.

Q4. In linear attention's kernel trick, why must the feature map φ produce non-negative outputs? What breaks if you use φ(x) = x on inputs with mixed-sign entries?

Q5. Why can linear attention generate autoregressively with O(1) memory per step (like an RNN), while standard softmax attention needs a key-value cache that grows with sequence length?

Q6. In the online-softmax worked example, block 1 was scores [2,0] and block 2 was [1,3], giving output ≈22.1385. Suppose you instead use one single block covering all four scores — does the output change? Now suppose a fifth key arrives with score 5 and value 50, still processed as a final one-element block after the first two blocks — trace the correction step and find the new output. Does linear attention's kernel-weighted average respond to this new high-scoring token in the same way?

A1. The n² term is a memory-traffic problem, not primarily a raw-arithmetic problem. GPUs have far more compute throughput (hundreds of TFLOPs) than memory bandwidth (a few TB/s), so any kernel that reads and writes an n×n array multiple times spends most of its wall-clock time waiting on HBM rather than computing — it is memory-bound. This is exactly why FlashAttention's fix targets memory traffic rather than FLOPs.

A2. False. FlashAttention performs the identical O(n²d) floating-point operations as standard attention and produces bit-identical (up to floating point rounding) output — it is an exact algorithm. Its speedup comes from minimizing HBM reads/writes via tiling and online softmax, not from doing less arithmetic. Linear attention is the one that actually changes the FLOP count, from O(n²d) to O(nd²), by giving up the exact softmax nonlinearity.

A3. n² = 65,536² = 4,294,967,296 = 2³². At 2 bytes per entry (fp16): 2³³ bytes = 8,589,934,592 bytes = exactly 8 GiB — for one head, one layer, one sequence. A model with, say, 32 heads and 32 layers would need 8 GiB × 32 × 32 = 8,192 GiB if every layer's full score matrix were kept simultaneously — over 100× the memory of an 80 GB GPU. This is precisely why no production system materializes these matrices; FlashAttention-style tiling (or an algorithmic change like linear attention) is not an optimization, it is a requirement at this scale.

A4. The denominator in the attention formula, Σ_jφ(Q_i)·φ(K_j), must stay positive to serve as a normalizer that turns the numerator into a weighted average of the value rows. If φ can output negative values, individual similarity terms can cancel or make the denominator zero or negative — the "weights" are no longer a valid convex combination (they can be negative or sum to something nonsensical), the output is no longer bounded within the range of V's values, and division by a near-zero or negative denominator causes numerical blow-up. This is exactly why practical linear-attention papers pick φ(x) = elu(x) + 1 (always ≥ 0) or an explicit non-negative kernel approximation (Performer's FAVOR+), rather than the identity map used in the worked example above, which only worked because every entry of Q and K was chosen non-negative.

A5. Linear attention's summary of the past is the fixed-size matrix S_t = Σ_j≤tφ(K_j)V_j^T (size d×d) plus the vector z_t = Σ_j≤tφ(K_j) (size d): both can be updated incrementally as each new token arrives, so only the current state — never the individual past tokens — needs to be kept. Softmax attention cannot do this because the normalizer (and every weight) depends nonlinearly on the full set of raw scores computed fresh against every past key; there is no fixed-size summary that lets you reconstruct softmax(Q·K₁, ..., Q·K_t) without holding onto every individual K_j, V_j — hence the KV cache, which grows with t.

A6. Single block (block_size = 4): this is simply the batch computation done directly, so m = 3, l ≈ 1.553002, o ≈ 34.38122, output ≈ 22.1385 — identical to the two-block version. Online softmax with any valid block partition always reproduces the exact batch result; only the intermediate correction steps differ, never the final answer.

Adding a fifth key (score 5, value 50) as a third block: after blocks 1–2 (the original four scores), the running state was m=3, l≈1.553002, o≈34.38122. New local max = 5, so m_new = max(3,5) = 5. Correction = exp(3−5) = exp(−2) ≈ 0.135335. Rescale: l_corrected ≈ 1.553002 × 0.135335 ≈ 0.210177; o_corrected ≈ 34.38122 × 0.135335 ≈ 4.65299. New block: exp(5−5)=1, so l = 0.210177 + 1 = 1.210177; o = 4.65299 + 1×50 = 54.65299. Output = 54.65299 / 1.210177 ≈ 45.16.

The single new token — one with a much larger raw score than any before it — pulls the output from ≈22.14 almost all the way to its own value of 50 (≈45.16, more than double the previous output), because softmax's exponential weighting concentrates almost all the probability mass on whichever score is largest: exp(5) dominates exp(3), exp(2), exp(1), exp(0) combined. This is the ripple effect worth noticing: a single highly-relevant new token can swing an exact-softmax output dramatically, which is exactly the "needle in a haystack" retrieval behavior long-context systems rely on. Linear attention's kernel similarity (say, elu(x)+1 dot products) grows much more gently than an exponential, so an equally large jump in raw similarity produces a proportionally smaller shift in the weighted average — it approximates softmax's overall shape but blurs its sharpest peaks, which is the concrete, measurable cost of trading exactness for linear-time scaling.

Think About It

Think about this: How would you explain efficient transformers: linear attention and flash attention to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where efficient transformers: linear attention and flash attention is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting efficient transformers: linear attention and flash attention to at least 3 other topics you have studied.