Fusion

Picture two operations chained together, one feeding the next. Run them unfused and the first op writes its whole result to memory before the second op is allowed to start reading it back. Merge them into a single kernel instead and the picture changes completely: the first op's output for one element flows straight into the second op's computation on that element, inside one loop, and the intermediate value never leaves fast memory at all. That merge is what XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → calls 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 →, and the instruction it produces is a kFusion, one op standing in for what used to be several.

It is easy to picture 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 → as the compiler finding a smarter formula, some algebraic shortcut dressed up as an optimization pass. That picture is wrong in a way worth fixing early, because it will mislead you at every later chapter. A fused kernel computes the identical sequence of arithmetic the unfused version computed; nothing about the math is rewritten, simplified, or reordered at the level of results. The only thing fusion changes is whether an intermediate value round-trips through memory or stays resident where the next op can reach it immediately.

Fusion is about the intermediates, never the math.

Two questions gate 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 → decision XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → makes, and the compiler runs the cheaper one first. Fusibility asks a structural question: can these two instructions even be merged, given their opcodes and how their operands connect? Profitability asks the harder one: would merging them actually help? XLA answers profitability with a cost model, backed by an is_fusible callback each backend gets to customize for its own hardware, since what counts as a win on a CPU and what counts as a win on a GPU are not the same arithmetic.

On GPU that cost model runs as a priority queue: every candidate producer gets ranked by the time it would save if fused, and the pass fuses in that order, highest expected benefit first. Nothing here is exhaustive search. It is a greedy ordering over real cost estimates, and you can watch it work directly: the 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 → x-ray at /gym/xla#fusion holds three real before-and-after pairs pulled from a TPU, each one a case where this exact ranking made a call you can check against your own intuition. The kernel path already measured what a fusion like this is worth in practice, a fused softmax beating the unfused chain outright, recorded with full provenance at /bench.

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 → pass, however aggressive, is bounded by one hard rule: it can only merge instructions that already sit on a dataflow edge in the graph it was handed. It never invents a new algorithm and it never restructures the computation into a shape that does less work. If two values are related in your head but not connected by a direct producer-consumer edge in the HLO, fusion has nothing to offer them.

This is the wall the kernel path spends an entire stage teaching you to see, at /l/xla: naive softmax needs the row maximum before it can compute a single exponential, so the score matrix gets written to memory, read back, and only then reduced. No pass in this pipeline can fold that spill away, because folding it away would mean changing the algorithm from multi-pass to streaming, and 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 → does not do algorithms. It does edges.

The dump below is real, pulled from the optimized module for the attention program this path has been tracking. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → looked at a broadcast feeding directly into a divide and merged the two into one %fused_computation: the broadcast materializes its input into the divide's shape, and the divide consumes it, all inside a single computation with its own parameters and one ROOT instruction. Nothing outside this fused block ever sees the broadcast's output on its own.

%fused_computation (param_0: f32[64,64], param_1.1: f32[64]) -> f32[64,64] {
  %param_0 = f32[64,64]{1,0} parameter(0)
  %param_1.1 = f32[64]{0} parameter(1)
  %broadcast.3 = f32[64,64]{1,0} broadcast(f32[64]{0} %param_1.1), dimensions={0}, metadata={op_name="jit(attend)/jit(main)/div"}
  ROOT %divide.0 = f32[64,64]{1,0} divide(f32[64,64]{1,0} %param_0, f32[64,64]{1,0} %broadcast.3), metadata={op_name="jit(attend)/jit(main)/div" source_file="attend.py" source_line=12}
}

This section used to publish a disagreement. LAB·X2 timed this chapter's attention on a Colab TPU v6e at 2048 by 128 in bf16, medians of twenty runs in isolated processes, and reported 187.5 microseconds with the pipeline untouched against 139.3 microseconds with --xla_disable_hlo_passes=fusion, a ratio of 0.74x. Disabling 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 → appeared to make the program faster, which is the opposite of everything above, so it went up here as an open question.

The question is now closed, and not in the direction anyone expected. Compiling both versions and reading what each one emitted returned the same module twice: 66 instructions either way, seven fusionsSeveral 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 → either way, no custom calls in either, and the same fusion at the root. The flag changed nothing. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → accepted a pass name, found nothing registered under it in this backend's pipeline, and carried on without a word. Both timings measured the same program, so the ratio was never about fusion at all.

A flag that names a pass the backend does not have is not an experiment. It is a typo the compiler agreed to.

Pass names are backend-specific, which chapter 4's dump shows plainly: the TPU pipeline files its work under names like Phase_1_pre_layout_assignment_passes and X64_elimination that a CPU compile never mentions. A name lifted from one backend's vocabulary can be silently absent from another's, and nothing in the output distinguishes a disabled pass from a misspelled one. The habit that costs ten seconds and would have caught this before it was published: compile both ways, diff the modules, and only time them once you have proof the flag moved something.

What remains open is smaller and better posed. The 26% gap between two identical compilations is still unexplained, though it is now a question about measurement rather than about 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 →, and the first thing to try is running the two processes in the opposite order to see whether the gap follows the order rather than the flag. Beyond that sits the original question, which was never actually asked on this hardware: find what the fusion stage is called in the TPU pipeline, disable that, and measure again.