The Cost of Attention

Dive into Deep Learning · §10.5

The cost of attention
quadratic complexity · online softmax and FlashAttention · sliding windows · recurrent linear attention

Sequence-layer complexity

Map n tokens of dimension d to n tokens of dimension d:

layer complexity sequential ops max path length
convolution (k) \mathcal{O}(knd^2) \mathcal{O}(1) \mathcal{O}(n/k)
recurrence \mathcal{O}(nd^2) \mathcal{O}(n) \mathcal{O}(n)
self-attention \mathcal{O}(nd^2+n^2d) \mathcal{O}(1) \mathcal{O}(1)

Sequence-layer connectivity for a convolution, a recurrence, and self-attention.

Self-attention combines full parallelism and constant path length with quadratic complexity in sequence length.

Measured quadratic scaling

One layer: 4n^2d FLOPs, two n \times n score buffers = 8n^2 bytes (fp32), per head, per sequence. Dominates the layer as soon as n > 2d.

n =  2048: measured    40.6 MiB, predicted 8n^2 B =    32.0 MiB
n =  4096: measured   128.0 MiB, predicted 8n^2 B =   128.0 MiB
n =  8192: measured   512.0 MiB, predicted 8n^2 B =   512.0 MiB
n = 16384: measured  2048.0 MiB, predicted 8n^2 B =  2048.0 MiB
  • Exact from n = 4096 up (the smallest run adds a few MiB of allocator overhead); doubling n quadruples this cost.
  • At n = 131{,}072: one fp32 attention map \approx 69 GB. Nobody stores it.

Measured runtime is quadratic

def wall_clock(f, *args, reps=10):
    f(*args)  # Warm up
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(reps):
        f(*args)
    torch.cuda.synchronize()
    return (time.time() - start) / reps

for n in [2048, 4096, 8192, 16384]:
    Q = torch.randn(1, n, d_h, device=d2l.try_gpu())
    with torch.no_grad():
        t = wall_clock(attention, Q, Q, Q)
    print(f'n = {n:5d}: {t*1e3:6.2f} ms')
n =  2048:   0.09 ms
n =  4096:   0.47 ms
n =  8192:   1.85 ms
n = 16384:   7.37 ms
  • Launch-bound while the GPU is idle; from a few thousand tokens on, each doubling of n roughly quadruples the time, consistent with the 4n^2d operation count.

Online softmax: the matrix never needs to exist

Per query, carry a running max, normalizer, and output; process keys in blocks:

m' = \max(m, \max_j a_j), \quad s' = s\,e^{m-m'} + \sum_j e^{a_j - m'}, \quad \mathbf{o}' = \mathbf{o}\,e^{m-m'} + \sum_j e^{a_j-m'}\mathbf{v}_j

\mathbf{o}/s equals \mathrm{softmax}(\mathbf{a})\mathbf{V}. When the running maximum increases, the factors e^{m-m'} rescale earlier partial sums.

Online softmax carries a running maximum, normalizer, and weighted-value sum across key blocks.

Chunked attention: exact, in linear memory

Twenty lines carry the update over all queries, one key block at a time:

torch.manual_seed(0)
n, d_h = 2048, 64
Q, K, V = (torch.randn(n, d_h, device=d2l.try_gpu()) for _ in range(3))
err = (chunked_attention(Q, K, V) - causal_attention(Q, K, V)).abs().max()
print(f'maximum deviation: {float(err):.2e}')
maximum deviation: 4.77e-07

  • Same answer, same FLOPs; the n^2 footprint is gone.

The bottleneck is memory traffic, not FLOPs

GPUs multiply hundreds of times faster than they fetch. Naive attention writes n^2 scores to slow memory and reads them back twice. FlashAttention (Dao et al., 2022): tiles in on-chip SRAM, softmax stats in registers, backward pass recomputes instead of storing.

B, H, n = 2, 8, 8192
X = torch.randn(B, H, n, d_h, device=d2l.try_gpu(), dtype=torch.float16)

def naive_heads(X):
    i = torch.arange(X.shape[-2], device=X.device)
    scores = X @ X.transpose(-1, -2) / math.sqrt(X.shape[-1])
    scores.masked_fill_(i[None, :] > i[:, None], torch.finfo(X.dtype).min)
    return torch.softmax(scores, dim=-1) @ X

def fused_heads(X):
    with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        return F.scaled_dot_product_attention(X, X, X, is_causal=True)

for name, f in [('naive', naive_heads), ('fused', fused_heads)]:
    t, mem = wall_clock(f, X), peak_memory(f, X)
    print(f'{name}: {t*1e3:6.2f} ms, peak memory {mem/2**20:7.1f} MiB')
naive:  19.82 ms, peak memory  4112.1 MiB
fused:   1.11 ms, peak memory    16.5 MiB
  • Exact attention with no quadratic score buffer; use scaled_dot_product_attention or jax.nn.dot_product_attention when the corresponding optimized backend is available.

Attention through a window

If each query attends to the w most recent positions, the causal mask becomes a band:

  • \mathcal{O}(nw) work instead of \mathcal{O}(n^2).
  • Deployed: Mistral 7B (w = 4096, 32 layers); Longformer adds global tokens; NSA learns the sparsity pattern itself.

Depth restores the reach

Information hops w-1 positions per layer: receptive field 1 + L(w-1) after L layers, as the following check confirms:

depth 1: last query reaches 8 positions (formula: 8)
depth 2: last query reaches 15 positions (formula: 15)
depth 4: last query reaches 29 positions (formula: 29)
  • The CNN bet (ch. 7), re-made: locality compounds through depth.

Linear attention: kernelize the score

Choose a factorizing kernel \phi(\mathbf{q})^\top\phi(\mathbf{k}) instead of \exp(\mathbf{q}^\top\mathbf{k}/\sqrt{d}). The query factors out of the attention sum:

\mathbf{o}_t = \frac{\phi(\mathbf{q}_t)^\top \mathbf{S}_t}{\phi(\mathbf{q}_t)^\top \mathbf{z}_t}, \qquad \mathbf{S}_t = \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\,\mathbf{v}_t^\top

A recurrent network: fixed d \times d matrix state, outer-product writes, queried like an associative memory (Katharopoulos et al., 2020; fast weight programmers, 1991).

Two forms, one answer

Parallel form (cumulative sums) for training; recurrent form (carry \mathbf{S}, \mathbf{z}) for generation:

torch.manual_seed(0)
n = 512
Q, K, V = (torch.randn(n, d_h, device=d2l.try_gpu()) for _ in range(3))
err = (linear_attention_parallel(Q, K, V)
       - linear_attention_recurrent(Q, K, V)).abs().max()
print(f'maximum deviation: {float(err):.2e}')
print(f'recurrent state: {d_h}x{d_h} + {d_h} floats '
      f'= {(d_h * d_h + d_h) * 4 / 1024:.0f} KiB at any sequence length')
maximum deviation: 2.38e-07
recurrent state: 64x64 + 64 floats = 16 KiB at any sequence length
  • 16 KiB of state at any context length, vs. a KV cache that grows forever. The fixed-size state \mathbf{S}_t is a lossy summary of the past.

Measured time and memory

  • Dense approaches slope 2; windowed remains near the launch-overhead floor; linear attention uses nd_h^2 memory and bandwidth-bound cumulative sums. Its principal advantage is the constant-memory recurrent mode.
  • Current approaches emphasize memory-efficient exact attention, windowed attention, and linear attention in recurrent form.

Recap: three scaling strategies

  • Reorganize: online softmax → FlashAttention. Exact, linear memory; the bottleneck was traffic, not FLOPs.
  • Restrict: sliding windows; depth restores reach as 1 + L(w-1).
  • Kernelize: linear attention = linear recurrence with matrix state:

\mathbf{S}_t = \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\,\mathbf{v}_t^\top

With identity decay, this recurrence has the state-space form developed in Chapter 12. Parallel scans and learned gates lead to related models such as Mamba; attention and recurrence occupy two ends of this design space (SSD, 2024).