The Performance Model

Same Chip, Two Orders of Magnitude Apart

One kernel — a square matmul — timed honestly 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.

Timing GPUs: the Trap

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.74 ms   with block_until_ready: 8.84 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.983 s   read once: 0.361 s

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

A Benchmark That Does Not Lie

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'

The Sweep: One Kernel, 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: 4.7', '1024: 21.8', '2048: 119.9', '4096: 158.7', '8192: 173.1']

Tiny: overhead-bound (launches rival arithmetic). Middle: a steep climb while the kernel learns to fill 128 SMs. Large: flat at the roof — compute-bound. 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.29 ms/call
mul: 0.30 ms/call
sin: 0.17 ms/call
sigmoid: 0.16 ms/call

sin costs no more than add: below the ridge, memory traffic sets the time and the arithmetic hides under it.

The Method

An unfused elementwise chain pays one memory round trip per op — we measured it, we diagnosed it, and we leave it bleeding until the compiler section cures it with one line.

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.