What the workload needs

You know a neural network runs mostly on matrix multiplication. To reason about hardware for it, you need two counts for every operation: the FLOPs it performs and the bytes it must move to and from off-chip memory. Their ratio, the arithmetic intensity, placed on a roofline next to the chip’s peak compute and bandwidth, bounds the performance any implementation can reach.¹

This chapter derives those counts for the operations that matter and applies them:

• GEMM intensity as a function of $M$, $N$, $K$ and bytes per value, and why it is set by the smallest dimension.
• Attention’s quadratic FLOPs and traffic, fused and unfused.
• Why prefill is compute-bound and decode is memory-bound, and how batch size and the KV cache move decode on the chart.
• Training’s $6N$ FLOPs per token and 16 bytes per parameter of state, against inference’s $2N$ and latency limits.
• Where the roofline bound is loose: ceilings, latency, quantization effects and multiple memory levels.

The AI chapters that follow assume this vocabulary. Number formats change the bytes per value (Number formats); GPUs, systolic arrays and spatial meshes are different ways to raise on-chip reuse; and The memory wall follows decode’s bandwidth problem down to HBM.

For $C = AB$ with $A \in \mathbb{R}^{M \times K}$, $B \in \mathbb{R}^{K \times N}$, the work is $MNK$ MACs, or $2MNK$ FLOPs.¹⁹ The layers map onto it as follows:

• Fully connected / linear: $M = \text{batch} \times \text{positions}$, $K$ = input features, $N$ = output features. With one input vector ($M = 1$) it degenerates to a matrix-vector product (GEMV).⁹
• Convolution: lowered (im2col) to a $K \times CRS$ filter matrix times a $CRS \times NPQ$ data matrix, where $C$ is input channels, $R \times S$ the filter, $K$ output channels and $NPQ$ the batch times output pixels. The price is duplicating each input up to $RS$ times, which lowers intensity unless the lowering is done on the fly.⁸
• Transformer block: the $Q$, $K$, $V$ and output projections ($4 d_{\mathrm{model}}^2$ weights per layer) and the two MLP matrices ($2 d_{\mathrm{model}} d_{\mathrm{ff}}$), giving $N \approx 12\, n_{\mathrm{layer}} d_{\mathrm{model}}^2$ parameters with $d_{\mathrm{ff}} = 4 d_{\mathrm{model}}$; plus the activation-by-activation products $QK^T$ and $PV$ inside attention.³

Kaplan et al.’s count gives a forward pass of about $2N + 2\, n_{\mathrm{layer}} n_{\mathrm{ctx}} d_{\mathrm{model}}$ FLOPs per token; the second term is attention over the context, and it is small while $d_{\mathrm{model}} > n_{\mathrm{ctx}}/12$.³ For a 6.7B model with width 4,096 and 32 layers¹⁶, that crossover is at a context of about 49,000 tokens, so at ordinary lengths the weight GEMMs dominate the FLOPs.

That FLOP share is not the time share. In BERT training on GPUs, tensor contractions were 99.80% of the FLOPs but 61.0% of the runtime; normalization (0.17% of FLOPs) took 25.5% and elementwise operations (0.03%) took 13.5%.⁷ The reason is bytes, not FLOPs, and that is what the rest of the chapter is about.

Per head, with sequence length $S$ and head dimension $d_h$: $Q, K, V \in \mathbb{R}^{S \times d_h}$; $S = QK^T \in \mathbb{R}^{S \times S}$; $P = \mathrm{softmax}(S/\sqrt{d_h})$ row-wise; $O = PV$.² Across $h$ heads with $h d_h = D$:

• Projections: $4D^2$ weights per layer ($Q$, $K$, $V$, output), 2 FLOPs per weight per token. GEMMs with $M$ = tokens.
• $QK^T$ and $PV$: $2S^2 D$ FLOPs each, so $4S^2 D$ per layer without masking; a causal mask lets a kernel skip about half. Both operands are activations, so there is no weight reuse to amortize, and the inner dimension of $PV$ is $S$.
• Softmax, scaling, masking: a few operations per score, but over $S^2 h$ scores; memory-bound reductions and elementwise ops.⁵

The textbook kernel sequence writes $S$ and $P$ to HBM, giving $\Theta(S d_h + S^2)$ HBM accesses per head. FlashAttention tiles $Q$, $K$ and $V$ through on-chip SRAM and recomputes the softmax normalization incrementally, so the $S \times S$ matrices never leave the chip; it reports up to 7.6× faster attention on GPT-2.⁵ In roofline terms, fusion moves attention from the slope to the flat roof without changing its FLOPs (the Try it sim has a toggle for this).

During decode, a step processes one new query per sequence against $S$ cached keys and values per layer. The FLOPs are $4SD$ per layer per token, but the KV cache must be streamed in for every step. Shazeer’s analysis puts the ratio of memory traffic to arithmetic for incremental multi-head attention at $\Theta(n/d + 1/b)$, with $n$ the context and $b$ the batch, and says both terms must be made small; a larger batch shrinks only the second.¹² Multi-query attention (one shared K/V head) and grouped-query attention (a few) shrink the first.¹²¹³

Define intensity $I = F/Q$, with $F$ the FLOPs and $Q$ the bytes moved between the chip and the memory level you are analyzing, normally off-chip DRAM or HBM. Williams et al. call this operational intensity and count traffic after it has been filtered by the caches, so that cache and tiling optimizations show up as a higher $I$.¹

For a GEMM with $b$ bytes per value and each operand read once and the result written once:¹⁹

• Square, $n = M = N = K$: $I = 2n/(3b)$. At $b = 2$, an $8192^3$ GEMM gives 2,731 FLOP/B; NVIDIA’s guide reports 2,730.¹⁹
• Skinny: if $N \ll M, K$, the $MK$ term dominates the bytes and $I \approx 2N/b$. NVIDIA’s $8192 \times 128 \times 8192$ example gives 124.1 FLOP/B.¹⁹ Intensity is set by the smallest dimension.
• GEMV ($N = 1$): $I \approx 2/b$, about 1 FLOP/B at 16 bits. Always memory-bound on any accelerator.¹⁹

For a weight GEMM, the dimension that sets intensity is usually the number of tokens processed together: batch × sequence length. That one fact explains most of the chapter: large-batch training and prefill have thousands of tokens per GEMM; decode has one per sequence.

The formula assumes ideal reuse. Real kernels reach it only if a tile of the output and the matching panels of the inputs fit on chip (derivation in Under the hood). Lowering a convolution through im2col, or writing intermediates between unfused kernels, adds traffic and lowers the intensity actually achieved.⁸⁵ Operations with no reuse, such as elementwise operations and reductions, sit at or below 1 FLOP/B, which is why they took 39% of BERT’s training time for 0.2% of its FLOPs.⁷

With peak compute $\pi$ (FLOP/s) and bandwidth $\beta$ (byte/s), the roofline bound for a kernel of intensity $I$ is¹

On log-log axes the bandwidth term is a line of slope 1 and the compute term a horizontal line, meeting at the ridge point $I^* = \pi/\beta$.¹ Williams et al. treat $I^*$ as a measure of how hard a machine is to program to peak: when the AMD Opteron X4 quadrupled the X2’s peak compute on the same memory system, the ridge moved from 1.0 to 4.4 FLOP/B, and kernels below 1 FLOP/B gained nothing from the upgrade.¹ AI accelerators sit two orders of magnitude further right (see By the numbers).

Three refinements make it a working tool rather than a slogan:

• Ceilings. Lines below the roof for missing optimizations: unbalanced multiply/add mix, no SIMD, no prefetching, poor memory affinity. A kernel cannot break through a ceiling without the matching optimization.¹ On an accelerator, a vector unit is a lower compute ceiling than the matrix unit; the AMD MI300X, for example, lists 163.4 TFLOPS of vector FP32 next to 1,307.4 TFLOPS of dense FP16 matrix math.²¹
• One roof per memory level. Nothing requires DRAM: with L2 or on-chip SRAM bandwidth on the slope and intensity counted against that level’s traffic, the ridge moves left.¹
• Measured, not datasheet, bandwidth. The paper measures sustainable bandwidth with microbenchmarks rather than using pin bandwidth.¹

The first TPU paper adapted the model to inference by counting integer operations per byte of weights read, since weights didn’t fit on chip. Its ridge sat at 1,350 operations per weight byte; the MLPs and LSTMs (at intensities of 64 to 200, equal to their batch sizes) were memory-bound and the CNNs compute-bound.⁶ The paper’s Table 1 shows intensity tracking batch size directly, the same lever this chapter returns to for decode.

Training. $\text{FLOPs} \approx 6N$ per token: $2N$ forward, about $4N$ backward (gradients with respect to both activations and weights).³ GPT-3 175B at 300 billion tokens is $3.14 \times 10^{23}$ FLOPs by exactly this count.¹⁷ Weight GEMMs have $M$ = tokens per device per step, typically thousands, so they sit far right on the roofline; the attention core and the normalization and elementwise ops are where time leaks.⁷ The binding constraint is often capacity: mixed-precision Adam needs $2\Psi + 2\Psi + 12\Psi = 16\Psi$ bytes for parameters, gradients and optimizer states,¹¹ plus saved activations, which is what pushes training across many devices (see the parallelism chapter in Systems).

Inference. Forward only, $2N$ FLOPs per token,⁴ under a latency target. Pope et al. split it into⁴

• Prefill: one forward pass over $B \times L_{\mathrm{in}}$ tokens in parallel. Weight GEMMs have $M = B L_{\mathrm{in}}$, so prefill is compute-bound at any realistic prompt length; they report 76% model FLOPs utilization (MFU) for large-batch prefill on PaLM 540B.⁴ It sets time to first token.
• Decode: $L_{\mathrm{gen}}$ sequential passes with $M = B$. Per step, the chip streams all weights plus every sequence’s KV cache. At small batch and short context the weights dominate; at large batch or long context the KV cache does. For a 500B+ model with multi-head attention at batch 512 and context 2,048, the KV cache is 3 TB, three times the weights, reloaded for every token while the compute sits mostly idle.⁴

For a decoder with $P$ parameters, $L$ layers, width $D$, context $S$, batch $B$ and $b$ bytes per value (multi-head attention, activations ignored):

At $B = 1$, $I \approx 2/b$: about 1 FLOP/B at 16 bits. As $B$ grows, the weight term is amortized and $I$ approaches $(2P + 4LSD)/(2bLSD)$; with $P \approx 12LD^2$ that is $(12D/S + 2)/b$. For $D = 4{,}096$ and $S = 1{,}024$ at 16 bits the ceiling is 25 FLOP/B no matter the batch, an order of magnitude below a typical ridge point. Shazeer’s $\Theta(n/d + 1/b)$ is the same statement.¹² This is why decode is attacked on the bytes side: fewer KV heads (MQA, GQA),¹²¹³ fewer bits per weight and cache entry, and more bandwidth (The memory wall).

The Expert view adds the FLOP and byte counts, compute and memory time, a bytes-per-value switch, and controls for each workload’s shape: $M$, $N$, $K$ for the GEMMs; prompt length and fused versus unfused for attention prefill; batch and context for decode. The decode model is a 6.4B-parameter decoder (32 layers, width 4,096, multi-head attention) and includes KV-cache reads. Try:

• Shrink one GEMM dimension and confirm that intensity tracks the smallest of $M$, $N$, $K$.
• Turn off fusion for prefill and see how far attention falls once the $S \times S$ matrices go to memory.
• Push decode batch to 512 at 1,024 and then 8,192 tokens of context. Where does intensity level off, and why?
• Switch to 8-bit. What moves, and what would you expect to happen to peak compute on a real chip?

The model is the bound, not a prediction: every point sits on the roof because it assumes ideal reuse and perfect overlap of memory and compute. The same peak is used at every precision, which real chips don’t do.

Two details: the attention products in prefill land at 114 because this analysis materializes the $S \times S$ matrix; with a fused kernel they would sit near $S/(2b)$. And the decode figure of 1 is the $2/b$ limit for GEMV at 16 bits; nothing about the layer shapes can lift it without batching.¹⁵¹⁹

Design pressures from the workload

• Training and inference want different chips. Training is dominated by large GEMMs (compute) and by capacity for 16 bytes of state per parameter plus activations;¹¹ decode is dominated by bandwidth and KV-cache capacity at a latency target.⁴ A design tuned for one is mis-provisioned for the other, which is why products and system configurations split along that line.
• The ridge keeps moving right. Compute has outpaced memory bandwidth, so more operations end up memory-bound.⁵ The first TPU’s authors estimated that giving it the GPU’s GDDR5 memory would have moved its ridge from 1,350 to 250 and raised the weighted mean speedup on their workload to 3.9×, more than a faster clock would.⁶
• Batch versus latency. Decode intensity rises with batch only until the KV cache dominates, and every step of batch adds per-token latency. MQA and GQA trade a little model quality for smaller caches;¹³ speculative decoding verifies several draft tokens per large-model pass, reporting 2–3× faster decoding with identical outputs on T5-XXL.¹⁴
• Fusion versus training state. Fusion helps less in training, because intermediates must be saved for the backward pass; FlashAttention instead recomputes attention on chip in the backward pass, spending FLOPs to save bytes.⁵

Failure modes of the analysis

• Counting FLOPs and ignoring bytes. BERT’s 0.2% of FLOPs in normalization and elementwise ops took 39% of its time.⁷
• Using the wrong roof. If a working set fits in an on-chip level, the DRAM roofline is the wrong bound; use that level’s bandwidth.¹
• Assuming ideal reuse. The GEMM formula assumes each operand crosses the memory boundary once; tile sizes limited by on-chip capacity, im2col duplication or unfused intermediates all add traffic.⁸
• Ignoring latency and quantization. Small GEMMs underfill the machine; tile and wave quantization waste work at awkward sizes.¹⁹
• Mixing up peak precisions. Ridge points computed from a sparse or 8-bit peak and a 16-bit workload are off by 2× or more.

The counting rules behind every number in this chapter, with their assumptions, and the model the simulation uses.

1. FLOPs of a GEMM

$C_{ij} = \sum_k A_{ik} B_{kj}$ is $K$ multiplies and $K - 1$ adds per output, so about $2K$ FLOPs for each of $MN$ outputs: $F = 2MNK$.¹⁹ Counting a MAC as 2 FLOPs is the convention in Kaplan et al. and Pope et al.; vendor peaks follow it.³⁴ Some papers count multiply-adds as one “FLOP” (ResNet-50’s “3.8 billion FLOPs” are multiply-adds), so check before comparing.¹⁰

2. GEMM traffic and the tiling argument

The compulsory traffic, each operand read once and the result written once, is $b\,(MK + KN + MN)$, giving the intensity formula above.¹⁹ Whether a kernel reaches it depends on on-chip capacity. A standard blocked schedule, step by step:

1. Keep a $T \times T$ tile of $C$ in on-chip storage (registers or SRAM) for its whole computation.
2. For each step along $K$, load a $T \times 1$ column slice of $A$ and a $1 \times T$ row slice of $B$: $2T$ values, $2Tb$ bytes.
3. Those feed $T^2$ MACs, $2T^2$ FLOPs, into the resident tile.
4. So the intensity of the inner loop is $2T^2/(2Tb) = T/b$, ignoring the one-time write of $C$.

Intensity grows linearly with tile size, and tile area grows with on-chip capacity. To clear a ridge of 200 FLOP/B at $b = 2$, $T$ must be at least 400, a $400 \times 400$ tile of accumulators (640 KB at 32-bit accumulation) plus staging buffers. This is why accelerators devote so much area to SRAM next to the multipliers, and why GPUs, systolic arrays and spatial meshes can all be read as different ways to hold a big tile close to the math (GPUs, Systolic arrays, Dataflow). It is also why the roofline paper’s operational intensity is measured after the caches: a better tiling is a higher $I$, not a different roof.¹

3. Attention, fused and unfused

For one layer with $S$ tokens, width $D = h d_h$ and $b$ bytes per value:

• FLOPs of $QK^T$ and $PV$: $4S^2 D$, or about $2S^2 D$ with causal masking skipped.²
• Fused traffic (read $Q$, $K$, $V$; write $O$): $4SDb$. Intensity with masking: $2S^2 D/(4SDb) = S/(2b)$, 1,024 FLOP/B for $S = 4{,}096$ at 16 bits.
• Unfused, add a write and a read of $S$ and of $P$ for each head: about $4hS^2 b$ more. For $S = 4{,}096$, $D = 4{,}096$, $h = 32$, $b = 2$ that is 4.3 GB against 134 MB, and intensity drops to about 31 FLOP/B.

Dao et al. prove FlashAttention needs $\Theta(S^2 d_h^2 / M)$ HBM accesses with on-chip memory $M$, against $\Theta(S d_h + S^2)$ for the standard algorithm, and that no exact algorithm does asymptotically better over all $M$.⁵

4. Decode intensity with a KV cache

1. Weights: $P = 12LD^2$ with the standard $d_{\mathrm{ff}} = 4D$.³ For $L = 32$, $D = 4{,}096$: $6.44 \times 10^9$.
2. FLOPs per decode step: $B\,(2P + 4LSD)$, weights plus attention over $S$ cached positions.³
3. KV cache per sequence: $2\ (\text{K and V}) \times L \times S \times D$ values. At 16 bits that is 512 KiB per token for this model, 537 MB for 1,024 tokens.
4. Bytes per step: $bP + 2bBLSD$. With $B = 1$, $S = 1{,}024$: $I = (13.42 \times 10^9) / (13.42 \times 10^9) \approx 1.0$. With $B = 64$: $I \approx 18$. As $B \to \infty$: $I \to (12D/S + 2)/b = 25$.

Grouped-query attention with $g$ key-value heads instead of $h$ divides the KV term by $h/g$, raising the ceiling accordingly; multi-query attention is $g = 1$.¹²¹³ Pope et al. reach the same conclusion from the other side: weight loading dominates memory time at small batch and short context, KV loading at large batch and long context.⁴

5. Training compute

Backward through a linear layer computes the gradient with respect to its input (one GEMM the size of the forward) and with respect to its weights (another), so $\text{backward} \approx 2 \times \text{forward}$ and the total is $6N$ per token, ignoring attention’s context term.³ Check: $6 \times 175 \times 10^9 \times 300 \times 10^9 = 3.15 \times 10^{23}$, against the $3.14 \times 10^{23}$ GPT-3’s authors report.¹⁷ Memory per parameter for mixed-precision Adam: 2 (16-bit weights) + 2 (16-bit gradients) + 12 (32-bit master weights, momentum and variance) = 16 bytes.¹¹

6. What the bound assumes

• Perfect overlap. $t = \max(F/\pi,\ Q/\beta)$ assumes memory and compute run concurrently. With no overlap, $t = F/\pi + Q/\beta$, at most 2× the bound, with the worst case at the ridge point.
• Enough concurrency. Sustaining bandwidth takes many outstanding requests, by Little’s law; the roofline paper’s “no prefetching” ceilings are what happens without them.¹
• One bottleneck at a time. In a multi-level hierarchy, each level gets its own roof, and the binding one is the lowest at that level’s intensity.¹

7. The simulation’s model

Each workload computes $F$ and $Q$ from its shape with the formulas above ($b = 2$ unless switched), then plots $I = F/Q$ against $\min(\pi,\ \beta I)$. Defaults: big GEMM $8192^3$ ($I = 2{,}731$), small GEMM $256^3$ (85), causal fused attention prefill at $S = 4{,}096$ for one layer of the 6.4B model (1,024), decode at batch 1 and 64 with 1,024 tokens of context (1.0 and 18), and a ReLU over 16 M values (0.25). It ignores activation traffic in decode, softmax and normalization FLOPs, tile and wave quantization, latency, and the fact that real chips change peak compute with precision.