Computational Complexity for ML

Counting operations tells you how a cost grows; counting bytes moved tells you how long it actually takes — and in machine learning the second one usually wins.

Below a certain amount of arithmetic per byte fetched, a kernel spends its time waiting for memory and never reaches the chip's advertised speed, however many operations you remove

Why Does This Exist?

You halve the number of operations in a kernel and it runs no faster. You increase the batch size eightfold and throughput goes up more than eightfold. Neither makes sense under a pure operation count, and both are routine.

The reason is that modern accelerators can do arithmetic far faster than they can fetch data. A GPU might sustain hundreds of teraflops while its memory delivers a couple of terabytes per second — a ratio of a hundred or more operations per byte. Feed it a kernel that does less arithmetic per byte than that and the arithmetic units sit idle waiting, and no amount of removing operations helps.

So the useful question is not "how many operations" but "how many operations per byte moved". That single reframing explains why batching pays, why FlashAttention was a breakthrough without changing the maths, why quantisation speeds up inference beyond the memory saving, and why a decoder generating one token at a time is so inefficient.

Counting operations in the abstract is algorithm complexity and big-O. The matrix cost model underneath is matrix multiplication. Where the retained values those bytes carry live is computational graphs.

Think of It Like This

A fast chef and a slow corridor

A chef can chop vegetables extraordinarily quickly. The ingredients are in a storeroom down a long corridor, and each trip takes a while.

Give the chef one carrot per trip and the kitchen's output is set entirely by walking speed. The chopping is instant; the corridor is everything. Buy a faster knife and nothing changes.

Give the chef a whole crate per trip and the picture flips. Now they spend most of their time chopping, the corridor barely matters, and a faster knife genuinely helps.

The deciding quantity is how much chopping happens per trip. Below some threshold you are corridor-limited; above it, knife-limited. And the threshold is a property of the kitchen, not of the recipe.

Accelerators are exactly this. The corridor is memory bandwidth, the knife is the arithmetic units, and the threshold is a fixed ratio for each chip. Most of performance engineering in ML is arranging to carry crates instead of carrots — which is what batching, fusion, and tiling all do.

How It Actually Works

Two costs, always count both.

Arithmetic — the operation count, what big-O usually describes.

Memory traffic — the bytes moved between slow memory and the compute units.

Their ratio is the arithmetic intensity:

I = \frac{\text{operations}}{\text{bytes moved}}

Each chip has a threshold ratio, its peak flops divided by its peak bandwidth. Below it you are memory-bound; above it, compute-bound. That is the roofline in the diagram: a rising line where bandwidth limits you, and a flat ceiling where arithmetic does.

Matrix multiplication, both ways

An $n \times n$ by $n \times n$ product does about $2n^3$ operations and moves about $3n^2$ numbers — read both inputs, write the output. So intensity grows with $n$ . Large matmuls are compute-bound and hit near-peak throughput. Small or thin ones are memory-bound and run at a fraction of peak no matter how fast the chip is.

That is the whole reason batching works. Multiplying one $1 \times 768$ input by a $768 \times 3072$ weight matrix reads the entire weight matrix — 2.4 million numbers — to do very little arithmetic. Stack 32 inputs and you read the same weights once for 32 times the work. The weights were going to be fetched either way, so the extra rows are nearly free.

Attention, and why FlashAttention mattered

For sequence length $n$ and head dimension $d$ , attention computes $QK^{T}$ ( $n^2 d$ operations), a softmax over $n^2$ entries, then multiplies by $V$ (another $n^2 d$ ). So arithmetic is $O(n^2 d)$ — quadratic in sequence length, which is the well-known scaling problem.

The less-quoted problem is memory. The score matrix has $n^2$ entries, and at $n = 4096$ that is 16.7 million numbers per head — 33.5 MB in float16. With 32 heads, over a gigabyte, written to and read back from slow memory.

FlashAttention changed no mathematics. It computes attention in tiles that fit in fast on-chip memory and never materialises the full score matrix. Same $O(n^2 d)$ operations, dramatically less memory traffic, and a large real speedup. It is the clearest possible demonstration that operation counts are the wrong thing to optimise.

Worked example

A dense layer with $d_{\text{in}} = 768$ and $d_{\text{out}} = 3072$ , weights in float16 (2 bytes).

The weight matrix holds $768 \times 3072 = 2{,}359{,}296$ numbers, so 4.7 MB must be read regardless of batch size.

Batch of 1. Operations $= 2 \times 1 \times 768 \times 3072 = 4.72$ million. Bytes moved $\approx 4.7$ MB. So intensity is about $1$ operation per byte. Against a chip threshold near 100, that is firmly memory-bound — roughly 1 percent of peak arithmetic throughput.

Batch of 128. Operations $= 2 \times 128 \times 768 \times 3072 = 604$ million. The weights are still read once; the activations add $128 \times 768 \times 2 = 0.2$ MB in and $128 \times 3072 \times 2 = 0.8$ MB out, so about 5.7 MB total. Intensity is about $106$ operations per byte. Now compute-bound. Operations rose 128-fold while bytes rose only 1.2-fold.

The attention numbers, for contrast. At $n = 4096$ and $d = 64$ per head: about $2 \times 4096^2 \times 64 \approx 2.1$ billion operations per matmul, against a score matrix of $4096^2 = 16.8$ million entries, or 33.5 MB in float16 — per head, written then read again. Thirty-two heads puts that over a gigabyte of avoidable traffic, and avoiding it is what FlashAttention does.

Code

1def dense_cost(batch: int, d_in: int, d_out: int, bytes_per: int = 2) -> tuple[float, float]:2    ops = 2 * batch * d_in * d_out                      # multiply + add3    weights = d_in * d_out * bytes_per                   # read once, any batch size4    activations = (batch * d_in + batch * d_out) * bytes_per5    return ops, weights + activations6
7
8for b in (1, 8, 128, 1024):9    ops, byts = dense_cost(b, 768, 3072)10    print(b, f"{ops/1e6:9.1f}M ops", f"{byts/1e6:6.2f}MB", f"intensity {ops/byts:7.1f}")11# -> 1      4.7M ops  4.73MB  intensity     1.0   memory-bound12# -> 8     37.7M ops  4.78MB  intensity     7.913# -> 128  604.0M ops  5.70MB  intensity   105.9   compute-bound14# -> 1024 4831.8M ops 12.58MB intensity   384.015
16# Attention's hidden cost is the score matrix, not the operation count.17n, d, heads = 4096, 64, 3218scores_mb = n * n * 2 / 1e619print(round(scores_mb, 1), round(scores_mb * heads / 1000, 2))   # -> 33.6 1.07 (MB, GB)

The intensity of 1.0 at batch 1 and about 106 at batch 128 match the hand calculation, and the jump is the batching argument in numbers. The last line is the gigabyte of score-matrix traffic that FlashAttention removes without changing a single operation.

Watch Out For

Optimising the operation count of a memory-bound kernel

Removing arithmetic from a kernel that is waiting on memory changes nothing, and this is where a great deal of optimisation effort goes to waste.

The classic case is a chain of small elementwise operations — a bias add, an activation, a dropout mask, a scale. Each has almost no arithmetic per element and each reads and writes the whole tensor, so the sequence is bandwidth-limited end to end. Making the activation function cheaper is invisible. Fusing the operations into one kernel so the tensor is read once and written once is a large win, and it removes no arithmetic at all.

The habit is to measure before optimising. Compute the arithmetic intensity of the kernel you care about and compare it against your hardware's ratio — that one number tells you which side you are on. Then profile: if the arithmetic units are idle, the answer is fusion, tiling, better layout, or a narrower dtype, not a cleverer formula.

Reading big-O as a prediction of runtime

Asymptotic complexity describes how cost grows. It says nothing about constants, memory behaviour, or cache locality, and in ML those often dominate at the sizes you actually run.

Two $O(n^2)$ implementations can differ by an order of magnitude purely from access pattern — a row-major traversal versus column-major on the same data. And a hand-written triple-loop matmul and a BLAS call are both $O(n^3)$ with a performance gap of a hundredfold or more, entirely from blocking and vectorisation.

So use big-O for what it is good at — deciding whether an approach can scale to your target size, and spotting the accidental quadratic in a data pipeline. Then benchmark at realistic sizes with realistic data, because crossover points are empirical.

The Quick Version

Count operations and bytes moved. Their ratio is arithmetic intensity.
Each chip has a threshold (peak flops over peak bandwidth). Below it you are memory-bound, above it compute-bound.
An $n\times n$ matmul is $2n^3$ operations against $3n^2$ numbers, so intensity grows with $n$ .
Batching works because the weights are read once regardless: batch 1 gives intensity near 1, batch 128 gives about 106.
Attention is $O(n^2 d)$ operations and $O(n^2)$ memory. At $n=4096$ the score matrix is 33.5 MB per head.
FlashAttention changed no maths — it tiles the computation so the score matrix is never materialised.
Autoregressive decoding runs at batch 1 per step, so it is memory-bound and reads every weight per token.
Quantisation gives close to a linear speedup in memory-bound regimes because time is set by bytes.
Training costs roughly 3× a forward pass.
Fusing elementwise chains is a large win and removes no arithmetic.
Big-O flatters algorithms with poor locality. Benchmark at real sizes.