The fusion taxonomy

Fusion is 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 main lever for cutting 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 → traffic, and its taxonomy is what turns a compiled dump from a wall of unfamiliar names into an actual diagnosis. Every fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → op in the dump carries a kind, and that kind tells you which shape of computation the compiler decided could share a single pass over memory, without you having to infer it from the op names alone.

XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → groups computations under three fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → kinds. A loop fusion covers a chain of elementwise work, the add then multiply then cast pattern, merged into a single pass so nothing round-trips through 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 → between steps. An input fusion is rooted at a reduction, with the operations that feed a sum or a max pulled inside the reduction itself, so the reduction reads its inputs once instead of reading a separately materialized intermediate. An output fusion feeds directly into a store, folding the last computation into the write instead of writing a temporary first.

Fusion has limits, and they follow from what a single pass over memory can express. Elementwise chains fuse cleanly because every output element depends only on the matching input elements. Reductions fuse the ops that feed them, but two independent reductions over the same tensor will not merge into one fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla →, since each still needs its own pass. Anything that needs the whole tensor materialized before the next op can proceed, a transpose feeding a matmul with an incompatible layout, for instance, forces a boundary, and boundaries are exactly where extra fusion ops show up in the dump.

The clearest evidence of a fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → boundary is the naive attention kernel at sequence length 8192 on a v6e-1 chip. The compiled dump carries the full bf16[8192,8192] score matrix through three separate fusion ops: one to compute it, one to run softmax over it, one to consume it in the second matmul. That score matrix never had to exist in 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 → if the kernel kept it in VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → the whole time, and end to end the naive version costs 414.4 microseconds. The three fusion ops are the spill made visible, printed right there in the dump.

Reading fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → kinds is worth building into a habit deliberately: scan the compiled dump for fusion ops, read the kind attached to each, and ask whether the boundary between two fusions was necessary or just a missed opportunity. Sometimes it's necessary, a genuine dependency the compiler cannot see through. Sometimes it's a shape mismatch or a layout disagreement you could remove by restructuring the kernel, and the fusion boundary disappears along with the extra 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 → pass it implied.