After the backend pipeline runs, the same block arrives fused and scheduled.
ENTRY main.40 {
Arg_0.1 = f32[16,32]{1,0} parameter(0), metadata={op_name="x"}
Arg_2.3 = f32[32,32]{1,0} parameter(2), metadata={op_name="wk"}
dot = f32[32,16]{1,0} dot(Arg_2.3, Arg_0.1), lhs_contracting_dims={0}, rhs_contracting_dims={1}, metadata={op_name="jit(block)/jit(main)/transpose"}
Arg_1.2 = f32[32,32]{1,0} parameter(1), metadata={op_name="wq"}
dot.11 = f32[16,32]{1,0} dot(Arg_0.1, Arg_1.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
dot.15 = f32[16,16]{1,0} dot(dot.11, dot), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
multiply_reduce_fusion = f32[16]{0} fusion(dot.15), kind=kLoop, calls=fused_computation.2, metadata={op_name="jit(block)/jit(main)/reduce_max"}
subtract_exponential_fusion = f32[16,16]{1,0} fusion(multiply_reduce_fusion, dot.15), kind=kLoop, calls=fused_computation.1, metadata={op_name="jit(block)/jit(main)/exp"}What XLA does with your program
XLA takes StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → and makes the performance decisions: which ops fuse into one kernel, what layouts arrays get, how work schedules across the chip. Its central optimization is fusion: merging adjacent elementwise chains, reductions, and data movement so intermediates stay in fast memory instead of round-tripping through HBM. On most programs it is very good, which is why the honest bar for any hand-written kernel is "beats XLA," not "beats naive."
Print what it actually decided with jax.jit(fn).lower(*args).compile().as_text(). The dump is backend-specific: run it on the TPU runtime to see TPU decisions. Inside you will find fusion ops with their operands: those are the compiler's actual choices, not documentation about them.
The ceiling, precisely
Fusion has an exact limit: XLA merges along dataflow edges, but it cannot apply algebraic identities. It cannot notice that a two-pass softmax could become a one-pass streaming computation, because that is a theorem about exponentials, not a graph transformation. So in naive attention, the seq x seq score matrix is computed, written to HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu →, and read back, no matter how well everything around it fuses. At seq 8192 in bf16 that is 134 MB written and read: 268 MB of traffic that exists because the algorithm is multi-pass.
Say the conclusion precisely, because the whole stack turns on it: the spill is not a fusion failure, it is an algorithm failure. No compiler pass fixes it. A theorem fixes it, and a human currently has to apply that theorem by hand, one layer down, on the other side of the gap.
One update from this site's own bench: XLA:TPU now pattern-matches this exact program and dispatches a hand-written online-softmax custom kernel; the profiler shows it by name, and GYM·08 holds the trace. That is recognition of a famous shape, not derivation of the rewrite: perturb the pattern and the ceiling returns. Chapter 05 states what follows.
Lessons
- 01Two dumps, two truthsYour program passes through XLA twice: once translated, once decided. Only one of those dumps can answer a performance question. ·
- 02The fusion taxonomyEvery fusion op in a compiled dump carries a kind, and the kinds are the difference between a wall of names and a diagnosis. ·
- 03LayoutsThe compiled dump is the only place you can see how a tensor sits in memory, and the annotation reads in one breath once you know the order. ·
- 04Reading the memory reportWhen a kernel asks for more VMEM than exists, the compiler prints the most literal error message in the stack. This lesson reads it line by line. ·
- 05Dumps on demand, and where XLA stopsYou can make the compiler write out every pass it runs, and you should know which decisions were never its to make. ·
- 06The one-level worldBelow every tensor compiler sits an IR with no tensors in it. Reading one small dump tells you exactly which facts get dropped on the way down, and why the wall this unit teaches has to exist where it does. ·
- 07Levels as a first-class ideaOne elementwise add, written three ways in one framework: whole tensors, then loops over buffers, then register-width vectors. What is computed never changes. What has been decided changes at every step. ·