XProf, and the three accounts

XProf is the profiler for the whole XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → stack: a capture library wired into jax, plus a viewer that renders what it caught. You reach it two ways. Programmatically, jax.profiler.trace writes a capture you can parse in code, which is how this course uses it. Interactively, the TensorBoard profile plugin (now the standalone xprof project) opens the same capture as a set of tools, each built to answer one kind of question.

Know the tools by the question each answers. The overview page answers where the step went: device compute, input pipeline, or host. The trace viewer is the timeline itself, every op on every plane against time: overlap and bubbles live here, and GYM·08 is a reading of exactly this pane's data. The op profile is the league table: time by op and category, which five ops actually matter. The memory tools answer OOM forensics: allocation over time, and which buffer owns the peak. The graph viewer is the compiled HLO, clickable. Newer builds add a rooflineThe floor model: latency is at least the larger of FLOPs over peak compute and bytes over bandwidth. Predict first, measure second.taught in /l/tpu → view that places each op against the machine's ceiling automatically, which is stage 0's discipline as a built-in pane.

The routing table, by symptom: a slow step starts at the overview page. Low MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → utilization goes to the op profile, then the trace viewer to find the bubbles between compute. An OOM goes to the memory viewer. An op you cannot explain goes to the graph viewer, and if it is a custom call, to this course's museum of reasons the models cannot see inside it. Multi-chip timing questions go to the pod tools, which is gate 04 country.

Everything before this section predicts; this section checks. The whole rooflineThe floor model: latency is at least the larger of FLOPs over peak compute and bytes over bandwidth. Predict first, measure second.taught in /l/tpu → discipline earns its keep only when you can hold the prediction next to what the machine actually did, and the machine keeps an account: the profiler's device timeline, one timed event per op, straight from the hardware. Capturing it costs four lines. Compile outside the trace window, run enough steady-state iterations to swamp noise, and read the xplane files that land.

# capture: compile OUTSIDE the window, then trace steady-state iterations
fn = jax.jit(naive_attention)
fn(*xs).block_until_ready()

with jax.profiler.trace("/tmp/xprof-run"):
    for _ in range(20):
        fn(*xs).block_until_ready()

# what lands: plugins/profile/<run>/<host>.xplane.pb, one per host

The working loop has four beats, and the order is the discipline. Predict the number from constants first, the way stage 0 drills. Profile to get the machine's account. Attribute the gap op by op: an unattributed gap is a guess wearing a number. Then change exactly one thing and measure again. Every optimization story this site tells, from the dimension_semantics retune to the byte-confirm, is this loop run honestly.

A program's cost has three tellings, and this site caught them disagreeing in public. The first is the cost model: compiled.cost_analysis() returns XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s own estimate of flops and bytes, instantly, with no hardware. It said naive attention at seq 8192 moves 155.2 MB. The second is the round-trip arithmetic from chapter 04: the score matrix written and read is 276.8 MB with I/O. Nearly a factor of two apart, and both cannot be right.

The third account settled it. The device timeline put 217.6 µs of hardware time on the two ops that carry the score matrix, and time converts to bytes through bandwidth. Not nameplate bandwidth: no real stream hits the datasheet ceiling. A pure 1 GB copy measured this chip at 1287.5 GB/s, 80.5% of nameplate, and at that achieved rate 217.6 µs is 280.2 MB. The estimate said 276.8. Agreement within 1.2%, and the cost model's 155.2 turned out to be an artifact: it cannot see inside custom calls.

Why was there a custom call inside a program made of jnp ops? Because the timeline held a surprise worth the whole exercise: XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →:TPU had pattern-matched the softmax and second matmul and dispatched its own kernel, named %online-softmax in the timeline, tiled 1024,1024. The compiler recognized a famous shape and swapped in a hand-written crossing. Recognition, not derivation: chapter 05 draws that line precisely, and this trace is its best evidence.