XLA

XLA takes 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.

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 HBM, 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.