6.12 — Compute and memory: O(T²d) attention, O(Td²) FFN

1. Why this session exists

Once you can write the forward pass (6.11) you owe yourself an answer to "how much does it cost to run it?" The cost has two halves — attention and FFN — that scale differently with sequence length $T$ and width $d$:

The two costs cross at $T \sim d$. For pretraining-length contexts $T = 2{,}048, 4{,}096$ and modern widths $d = 4096, 8192$, FFN dominates compute. For long-context inference at $T = 32{,}768, 128{,}000, 1{,}000{,}000$, attention dominates. This single crossover dictates almost every architectural choice in modern LLMs: from sliding-window attention (6.13), to KV cache (6.14), to FlashAttention (Cap 10.5), to GQA/MQA (Cap 8.3) — they all chase one or the other of these terms.

This session counts the FLOPs precisely and counts the memory both for inference (KV cache) and for training (activation memory). After you have the numbers in hand, the architectural variants of Chapter 8 and Chapter 10 become straightforward optimisations of one of the two terms.

2. The two regimes, derived

Recall one Transformer block has two sub-layers, attention and FFN, each preceded by a LayerNorm. We will count matrix multiplications (matmuls) — they dominate every other cost (LN, softmax, biases, activations).

Conventions, fixed throughout:

• $T$ = sequence length (number of tokens in context).
• $d$ = residual-stream width ($d_\text{model}$).
• $h$ = number of attention heads; $d_h = d / h$ is the head dimension.
• $d_\text{ff} = 4 d$ = FFN inner width.
• $L$ = number of blocks.
• One "FLOP" = one multiply-add (i.e. we count $2$ scalar ops per FMA, the standard convention).

A matmul of two matrices of shape $(A, B)$ and $(B, C)$ costs $2 \cdot A \cdot B \cdot C$ FLOPs.

2.1 Attention sub-layer FLOPs

The attention sub-layer of one block does (see 6.4 §3.1):

| Op | Shape | FLOPs | |---|---|---| | Project $X$ to $Q, K, V$ | $(T, d) \cdot (d, d)$, three times | $6 T d^2$ | | Score $S = Q K^\top$ (per head) | $h \cdot (T, d_h) \cdot (d_h, T)$ | $2 T^2 d$ | | Apply $S V$ (per head) | $h \cdot (T, T) \cdot (T, d_h)$ | $2 T^2 d$ | | Output projection $W_O$ | $(T, d) \cdot (d, d)$ | $2 T d^2$ | | Total attention | | $\boxed{\ 8 T d^2 + 4 T^2 d \ }$ |

The $8 T d^2$ term is the four projections ($Q, K, V, O$). The $4 T^2 d$ term is the score-and-mix ($Q K^\top$ + $S V$). RoPE, masking, and softmax are $O(T^2)$ each and we drop them.

2.2 FFN sub-layer FLOPs

The FFN (6.8 §5) is two matmuls with $d_\text{ff} = 4 d$:

| Op | Shape | FLOPs | |---|---|---| | $W_1$: $(T, d) \cdot (d, d_\text{ff})$ | $(T, d) \cdot (d, 4d)$ | $8 T d^2$ | | $W_2$: $(T, d_\text{ff}) \cdot (d_\text{ff}, d)$ | $(T, 4d) \cdot (4d, d)$ | $8 T d^2$ | | Total FFN | | $\boxed{\ 16 T d^2 \ }$ |

2.3 Per-block total

Summing:

For $L$ blocks, multiply by $L$. Drop biases, norms, and the embedding/output (those are $\Theta(L d)$ or $\Theta(V T d)$; the latter is comparable in magnitude only at very small models, e.g. GPT-2-small).

3. The crossover at $T = 4 d$

Setting attention compute equal to FFN compute:

(If we drop the $8 T d^2$ projection cost of attention and only count the score-mix $4 T^2 d$, the crossover is $T = 4 d$; depending on what one chooses to "attribute" to attention, the crossover is in the band $[2d, 4d]$. Most authors use $T \approx 4d$.)

For LLaMA-7B ($d = 4096$) this is $T \in [8192, 16384]$. So:

• Pretraining at $T = 2048$ or $T = 4096$ → well below the crossover → FFN dominates.
• Long-context inference at $T = 32768$ or $T = 128000$ → well above the crossover → attention dominates.

This single inequality dictates where every "make Transformers efficient" effort focuses:

• Short context. Optimise FFN: SwiGLU (Cap 8.4) makes it expressive; MoE (Cap 8.6) makes it conditional.
• Long context. Optimise attention: sliding window (6.13), FlashAttention I/O (Cap 10.5), GQA/MQA (Cap 8.3), KV cache (6.14), linear attention (Mamba, Cap 8.7).

4. Memory: activations and KV cache

FLOPs are time; memory is space. Both bottleneck Transformer scaling, often before FLOPs do.

4.1 KV cache (inference)

At autoregressive inference (generating one token at a time), the keys and values of all previous tokens do not change. We cache them to avoid recomputation; the cache stores

The factor $2$ is for K and V together. For LLaMA-7B (one batch, fp16):

At $T = 4096$: $\approx 2$ GB. At $T = 32{,}768$: $\approx 17$ GB. At $T = 128{,}000$: $\approx 67$ GB — larger than a single H100's 80 GB. This is why GQA (Cap 8.3) — sharing K, V heads across multiple Q heads — and MLA (Cap 8.5) — low-rank K, V compression — exist: shrinking the $h \cdot d_h$ factor in the KV-cache formula.

For a full discussion of the KV cache see 6.14.

4.2 Activation memory (training)

During training we must store every intermediate activation for the backward pass (autodiff, 2.11). Per token, per block, the largest activations are:

• The output of each sub-layer: $T \cdot d$ floats.
• The attention scores: $h \cdot T \cdot T$ floats — this is the $T^2$ memory cost.
• The FFN hidden: $T \cdot d_\text{ff} = T \cdot 4d$ floats.

For long context, $h T^2 \gg T d_\text{ff}$. A 32-head, $T = 8192$ model stores $32 \cdot 8192^2 = 2 \times 10^9$ scores per block, in fp16: $4$ GB per block, $128$ GB across 32 blocks. This is the primary reason long-context training is hard, and the primary reason for FlashAttention (Cap 10.5), which recomputes the scores on-the-fly instead of materialising them.

4.3 The 4-quadrant cost grid

Putting it together:

| | Short $T$ ($T \ll 4d$) | Long $T$ ($T \gg 4d$) | |---|---|---| | Compute | FFN dominates ($16 T d^2$) | Attention dominates ($4 T^2 d$) | | Memory | FFN hidden ($4 T d$) | Attention scores ($T^2$) |

Most architectural choices in modern LLMs map onto one cell of this grid.

5. Worked example: LLaMA-7B

Numbers for one forward pass:

• $d = 4096, L = 32, h = 32, d_h = 128$.
• $d_\text{ff} = 4d = 16384$ (LLaMA actually uses $\frac{2}{3} \cdot 4d = 10922$ rounded to $11008$ for SwiGLU; we'll use $4d$ here for clarity).

5.1 $T = 2048$ (pretraining context)

FFN/attention ratio $\approx 1.6$; FFN dominates (matches §3 crossover prediction at $T < 4d$).

5.2 $T = 32768$ (long-context inference)

Attention/FFN ratio $\approx 2.5$; attention dominates — by a wide margin, and growing as $T^2$.

This is the structural reason "long context is expensive". The cost is not abstract; it follows directly from the $T^2 d$ vs $T d^2$ counting we just did.

6. Generation cost vs prefill cost

A subtlety specific to LLMs: during generation, the model processes one token at a time, with all previous keys/values cached. Per generated token, effective $T_\text{generated} = 1$ but $T_\text{context}$ is the cache length. The FLOPs become:

• Prefill (one forward over the prompt of length $T_p$): the cost we just derived, scaling as $T_p d^2$ + $T_p^2 d$.
• Per generated token: one row of $Q$ attends to all $T$ cached K, V; the FFN runs on one position. Cost is $\sim 8 d^2 + 4 T d + 16 d^2$ FLOPs per token — linear in $T$ at the attention term.

Generating $N$ tokens, total cost $\sim N (24 d^2 + 4 T d)$. The cumulative cost scales as $N \cdot T$ (attention term) — still quadratic in combined sequence length, but distributed across $N$ separate calls rather than one big one. This is what makes the KV cache an enormous practical win for inference; see 6.14.

7. Putting it together

Three numbers to remember:

1. Per-block FLOPs: $24 T d^2 + 4 T^2 d$. Roughly: $24 T d^2$ for FFN-and-projections, $4 T^2 d$ for the score-and-mix.
2. Crossover: at $T = 2d$ to $4d$, the two terms are equal. Below: FFN dominates. Above: attention dominates.
3. KV-cache memory: $2 L h d_h T$ elements per batch. Linear in $T$, quadratic in model depth × width product. The dominant memory of long-context inference.

Next session turns to the long-context regime — what architectural variants exist to bend the $T^2 d$ term back down to $T d$ or $T \log T$, while preserving model quality.