The Performance Model

Same Chip, Two Orders of Magnitude Apart

One square matrix multiplication, synchronized and timed at two sizes on the same GPU delivers a percent or so of peak at one size and nearly the full specification number at the other.

Explaining that gap yields the chapter’s whole toolkit:

  • a cost model you can compute with pencil and paper,
  • a measurement discipline the frameworks’ async dispatch makes non-optional,
  • a diagnosis → fix table the next six sections fill in.

Arithmetic Intensity and the Roofline

Every op asks for FLOPs and bytes. Intensity = FLOPs/byte.

\textrm{performance} \leq \min(P,\; I \cdot \beta)

Matmul \mathbf{X}_{B\times D}\mathbf{W}_{D\times F}: intensity \approx 2B/b — batch size is intensity. Elementwise ops: intensity below 1, forever bandwidth-bound.

The Ridge Point, on Our Card

device = jax.devices()[0]
peak_tflops = 165.0   # RTX 4090 dense bf16 tensor-core spec
bandwidth_tbs = 1.008  # GDDR6X spec
print(f'{device.device_kind}')
print(f'ridge point = {peak_tflops / bandwidth_tbs:.0f} FLOP/byte')
NVIDIA GeForce RTX 4090
ridge point = 164 FLOP/byte

~165 FLOP/byte: the machine wants about 165 ops for every byte fetched just to break even. Almost nothing you write naturally gets there. Performance work is mostly about bytes.

Correct GPU Timing

Dispatch is asynchronous — Python enqueues, the GPU runs behind.

key = jax.random.PRNGKey(0)
a = jax.random.normal(key, (4096, 4096), dtype=jnp.bfloat16)
jnp.dot(a, a).block_until_ready()  # Warmup: triggers compilation

t0 = time.perf_counter()
for _ in range(10):
    b = jnp.dot(a, a)
naive = time.perf_counter() - t0

b.block_until_ready()  # Drain the naive loop's in-flight work first
t0 = time.perf_counter()
for _ in range(10):
    b = jnp.dot(a, a)
b.block_until_ready()
honest = time.perf_counter() - t0
print(f'naive timer: {1000 * naive:.2f} ms   '
      f'with block_until_ready: {1000 * honest:.2f} ms')
naive timer: 0.52 ms   with block_until_ready: 8.09 ms

Implicit Barriers

Anything that needs a value waits for the device: .item(), .numpy(), print(x), if loss < 0.1:, nonzero().

x = jnp.ones((256, 256))

t0 = time.perf_counter()
s = 0.0
for _ in range(1000):
    y = (x * 1.01).sum()
    s += float(y)          # Reads the value: a barrier on every step
sync_every = time.perf_counter() - t0

t0 = time.perf_counter()
s = jnp.zeros(())
for _ in range(1000):
    y = (x * 1.01).sum()
    s = s + y              # Accumulate on the device: no barrier
s = float(s)               # One read at the very end
sync_once = time.perf_counter() - t0
print(f'read every step: {sync_every:.3f} s   '
      f'read once: {sync_once:.3f} s')
read every step: 0.776 s   read once: 0.298 s

Rule: synchronize once per minibatch at most — and only when the host actually needs the value.

Synchronized Benchmarking

Warmup (kernel selection, compilation), sync, time, sync:

class Benchmark:
    """Time a callable: warmup, then device-synchronized average seconds."""
    def __init__(self, f, warmup=3, repeats=10, desc='time'):
        self.desc = desc
        for _ in range(warmup):
            out = f()
        jax.block_until_ready(out)
        t0 = time.perf_counter()
        for _ in range(repeats):
            out = f()
        jax.block_until_ready(out)
        self.time = (time.perf_counter() - t0) / repeats

    def __repr__(self):
        return f'{self.desc}: {1000 * self.time:.2f} ms/call'

One Kernel across Three Regimes

sizes = [256, 512, 1024, 2048, 4096, 8192]
achieved = []
for n in sizes:
    a = jax.random.normal(key, (n, n), dtype=jnp.bfloat16)
    compiled = jax.jit(jnp.dot).lower(a, a).compile()
    flops = compiled.cost_analysis()['flops']  # Analytic count, no timing
    t = Benchmark(lambda: compiled(a, a)).time
    achieved.append(flops / t / 1e12)
d2l.plot(sizes, [achieved], 'matrix size $n$', 'achieved TFLOP/s',
         xscale='log', yscale='log')
print([f'{n}: {tf:.1f}' for n, tf in zip(sizes, achieved)])

['256: 0.4', '512: 3.4', '1024: 26.3', '2048: 134.0', '4096: 157.9', '8192: 178.9']

Small matrices are overhead-bound because launch time rivals arithmetic. Intermediate matrices use an increasing fraction of the 128 SMs. Large matrices approach the compute roof. The roofline is the ceiling; the measured knee (~2048–4096) sits well past the nominal crossover (~500) — that gap is utilization and overhead.

Diagnosis Determines the Fix

x = jax.random.normal(key, (4000, 4000))
for desc, f in [('add', lambda: x + 1), ('mul', lambda: x * 1.5),
                ('sin', lambda: jnp.sin(x)),
                ('sigmoid', lambda: jax.nn.sigmoid(x))]:
    print(Benchmark(f, desc=desc))
add: 0.20 ms/call
mul: 0.27 ms/call
sin: 0.16 ms/call
sigmoid: 0.16 ms/call

Below the ridge, sin and addition take similar time because memory traffic, rather than arithmetic, determines elapsed time.

The Method

An unfused elementwise chain performs one memory round trip per operation. The compiler section measures how fusion removes the redundant traffic.

measure → classify → fix → re-measure

The rest of the chapter is this loop, applied: hardware explains the constants; compilation attacks bandwidth and overhead; memory and precision buy headroom; more GPUs buy more roof.