What the clock caught

Chapter 1 established that JAX dispatches asynchronously, so a line returns before the device is done. The consequence is easier to feel as a measurement than as a sentence. Time the same jitted matmul at three sizes, once without blocking and once with, and the unblocked column barely moves while the blocked column grows with the work.

The middle column below spans 156.4 to 5929.7 microseconds, a factor of 38 for a factor of 512 more arithmetic. The left column sits between 14.9 and 18.4 microseconds the whole way. That left number is the cost of a Python call plus an enqueue, and it is nearly independent of the program you enqueued.

So a benchmark that forgets to block is not merely off by a constant. It reports a number that tracks your dispatch path and not your computation, which is why an unblocked benchmark can show a rewrite getting faster when the rewrite did nothing at all.

An unblocked timing loop measures your Python, not your program.

import time

import jax
import jax.numpy as jnp

f = jax.jit(lambda x: x @ x)

def best(g, x, r=25):
    return min(g(x) for _ in range(r))

def dispatch(x):
    t0 = time.perf_counter()
    f(x)
    return time.perf_counter() - t0

def compute(x):
    t0 = time.perf_counter()
    f(x).block_until_ready()
    return time.perf_counter() - t0

for n in (128, 512, 1024):
    x = jnp.ones((n, n))
    f(x).block_until_ready()
    d, c = best(dispatch, x), best(compute, x)
    print(f"{n:5d} {d * 1e6:7.1f} {c * 1e6:9.1f} {c / d:7.0f}")

#   128    14.9     156.4      11
#   512    18.4     736.4      40
#  1024    16.5    5929.7     359

The first call to a jitted function traces it, lowers it, compiles it, and then runs it. Timing that call gives you an upper bound on the compile rather than the compile itself, and for a six-layer MLP on this machine the bound is 651.6 milliseconds against a steady step of 5.9 milliseconds.

Divide and the number stops being an abstract complaint about warmup. One compile of this function costs 111 steady steps. Across four runs of the same script the ratio landed between 105 and 190, because the steady step varies more than the compile does when the machine is busy.

That ratio is the exchange rate for a mid-run recompile. LAB·J2 is where you find out which part of the cache key moved and drive the log to silence; this is what one entry in that log costs you before you go looking.

import time

import jax
import jax.numpy as jnp

def mlp(p, x):
    for w, b in p:
        x = jnp.tanh(x @ w + b)
    return x.sum()

key = jax.random.key(0)
p = [(jax.random.normal(k, (512, 512)) * 0.05, jnp.zeros(512))
     for k in jax.random.split(key, 6)]
x = jnp.ones((256, 512))
f = jax.jit(mlp)

t0 = time.perf_counter()
f(p, x).block_until_ready()
first = time.perf_counter() - t0

t1 = time.perf_counter()
for _ in range(20):
    out = f(p, x)
out.block_until_ready()
steady = (time.perf_counter() - t1) / 20

print(f"first call {first * 1e3:7.1f} ms")
print(f"steady     {steady * 1e3:7.1f} ms")
print(f"ratio      {first / steady:7.0f}")

# first call   651.6 ms
# steady         5.9 ms
# ratio          111

Take fifty calls of one function and move a single line. Block once after the loop and the mean is 727 microseconds per call. Block inside the loop, on every call, and the mean is 1524. Same function, same fifty calls, same array.

The gap is async dispatch doing its job. When nothing blocks mid-loop, the host runs ahead and enqueues the next call while the device is still on the current one. When every call blocks, the host has nothing queued and each step pays the full round trip before the next one is submitted.

Neither number is a lie, and picking between them is a question about your program rather than about JAX. A training loop that fires steps back to back is described by the first. A service that waits on each result before deciding what to do next is described by the second. Say which one you measured, because a reader cannot tell from the number.

import time

import jax
import jax.numpy as jnp

f = jax.jit(lambda x: x @ x)
x = jnp.ones((512, 512))
f(x).block_until_ready()

def block_at_the_end(k=50):
    t0 = time.perf_counter()
    for _ in range(k):
        out = f(x)
    out.block_until_ready()
    return (time.perf_counter() - t0) / k

def block_every_call(k=50):
    t0 = time.perf_counter()
    for _ in range(k):
        f(x).block_until_ready()
    return (time.perf_counter() - t0) / k

a = min(block_at_the_end() for _ in range(7)) * 1e6
b = min(block_every_call() for _ in range(7)) * 1e6
print(f"block at the end {a:6.0f} us/call")
print(f"block every call {b:6.0f} us/call")

# block at the end    727 us/call
# block every call   1524 us/call

Two sections back, a 512 matmul measured 736.4 microseconds. One section back, the same matmul on the same machine measured 1524. Both are honest, and they disagree because one is the minimum of 25 isolated calls and the other is the mean of 50 back-to-back ones.

A minimum reports the best the machine managed with the scheduler out of the way, which is the right statistic when you are comparing two implementations and want the noise gone. A mean over a loop reports what a run of that length actually delivered, which is the right statistic when you are quoting throughput to someone who will plan around it. The method is not a footnote to the number; it is half of the number.

The spread is worth knowing before you celebrate a change. Five consecutive runs of the identical script, on this machine, put the 1024 blocked figure between 5929.7 and 7167.1 microseconds, a 21 percent range with nothing changed. A difference smaller than that is not a result yet.